// MNIST MLP: 784 -> 300 -> 100 -> 10. Fixed five epochs; test used once at the end.
// First run tools/benchmarks/fetch_mnist.sh from the repository root.
import aner.dataset
import aner.nn
import aner.ml
import aner.dev
import aner.viz
import aner.metrics

let seed = 42
let epochs = 5
let batch_size = 128
let official = dataset.mnist(directory: "build/benchmarks/mnist/data", partition: "train")
let testing = dataset.mnist(directory: "build/benchmarks/mnist/data", partition: "test")
let partitions = official.partition(validation: 0.1, test: 0.0, seed: 2026, stratify: true, group_aware: false)
let training = partitions.train()
let validation = partitions.validation()
print(partitions.audit())
dataset.require_disjoint(training, testing)
dataset.require_disjoint(validation, testing)
training.require_training()
// Placement follows CLI options; preserve is the default precision policy.
print(nn.execution())
print(nn.backend_details())
dev.automatic(false)
nn.reset_memory_peaks()
let started = nn.clock()
// Exploration: topology, training contract and evaluation are explicit.
let network = nn.mlp(inputs: 784, hidden: List(300, 100), outputs: 10, seed: seed)
let settings = ml.training(epochs: epochs, batch: batch_size, rate: 0.001, seed: seed)
    .validation(validation)
let fitted = ml.classifier(network, standardize: false, loss: "cross_entropy", optimizer: "adam")
    .fit(training, settings, progress: true)

print("BENCHMARK_FIT_SECONDS")
print(nn.clock() - started)
print("BENCHMARK_TRAIN_LOSS")
print(fitted.training_loss())
print("BENCHMARK_TRAIN_ACCURACY")
print(fitted.training_accuracy())
let test_started = nn.clock()
let evaluation = fitted.evaluate(testing)
print("BENCHMARK_TEST_SECONDS")
print(nn.clock() - test_started)
print("BENCHMARK_TEST_LOSS")
print(evaluation.loss())
print("BENCHMARK_TEST_ACCURACY")
print(evaluation.accuracy())
print("BENCHMARK_MEMORY")
print(nn.memory())
print(fitted.summary())
dev.capture("Final model", fitted)