// 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()
// Functional control: each Dense initialization and activation is visible.
fn build_network(seed: Int64) {
    return nn.sequential(seed: seed)
        .dense(inputs: 784, outputs: 300, initialization: "xavier")
        .relu()
        .dense(inputs: 300, outputs: 100, initialization: "xavier")
        .relu()
        .dense(inputs: 100, outputs: 10, initialization: "xavier")
}
fn train_network(network: NeuralNetwork, samples: Dataset, held_out: Dataset,
                 epochs: Int64, batch: Int64, seed: Int64) {
    let contract = ml.training(epochs: epochs, batch: batch, rate: 0.001, seed: seed)
        .validation(held_out).select_best(false)
    let plan = ml.classifier(network, standardize: false, loss: "cross_entropy", optimizer: "adam")
    return plan.fit(samples, contract, progress: true)
}
let fitted = train_network(build_network(seed), training, validation, epochs, batch_size, seed)

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)