// 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()
// Research control: parameters, forward, objective, batches and Adam updates.
class MnistNetwork {
    var w1: Tensor
    var b1: Tensor
    var w2: Tensor
    var b2: Tensor
    var w3: Tensor
    var b3: Tensor

    fn forward(input: Tensor) {
        let hidden1 = nn.relu(nn.add_bias(nn.matmul(input, w1), b1))
        let hidden2 = nn.relu(nn.add_bias(nn.matmul(hidden1, w2), b2))
        return nn.add_bias(nn.matmul(hidden2, w3), b3)
    }
    // Detached parameter leaves avoid constructing a training graph during evaluation.
    fn infer(input: Tensor) {
        let hidden1 = nn.relu(nn.add_bias(nn.matmul(input, nn.detach(w1)), nn.detach(b1)))
        let hidden2 = nn.relu(nn.add_bias(nn.matmul(hidden1, nn.detach(w2)), nn.detach(b2)))
        return nn.add_bias(nn.matmul(hidden2, nn.detach(w3)), nn.detach(b3))
    }
    fn update(state: Optimizer) {
        w1 = state.update(w1, key: "w1")
        b1 = state.update(b1, key: "b1")
        w2 = state.update(w2, key: "w2")
        b2 = state.update(b2, key: "b2")
        w3 = state.update(w3, key: "w3")
        b3 = state.update(b3, key: "b3")
    }
}
fn objective(logits: Tensor, target: Tensor) {
    return nn.cross_entropy(logits, target)
}
class Scores {
    public let loss: Float64
    public let accuracy: Float64
}
fn evaluate_network(network: MnistNetwork, samples: Dataset, size: Int64) {
    var loss_sum = 0.0
    var correct_sum = 0.0
    for batch in samples.batches(size: size, seed: 0, shuffle: false) {
        let input = nn.place(batch.features())
        let target = nn.place(batch.one_hot())
        let output = network.infer(input)
        let count = Float64(batch.rows())
        loss_sum = loss_sum + nn.item(nn.cross_entropy(output, target)) * count
        correct_sum = correct_sum + metrics.accuracy(nn.argmax(output), batch.targets()) * count
    }
    return Scores(loss: loss_sum / Float64(samples.rows()), accuracy: correct_sum / Float64(samples.rows()))
}
let network = MnistNetwork(
    w1: nn.parameter(nn.place(nn.xavier(rows: 784, cols: 300, seed: seed))),
    b1: nn.parameter(nn.place(nn.zeros(1, 300))),
    w2: nn.parameter(nn.place(nn.xavier(rows: 300, cols: 100, seed: seed + 1))),
    b2: nn.parameter(nn.place(nn.zeros(1, 100))),
    w3: nn.parameter(nn.place(nn.xavier(rows: 100, cols: 10, seed: seed + 2))),
    b3: nn.parameter(nn.place(nn.zeros(1, 10))))
let state = nn.optimizer(name: "adam", rate: 0.001)
fn train_network(network: MnistNetwork, state: Optimizer, training: Dataset, validation: Dataset,
                 epochs: Int64, batch_size: Int64, seed: Int64, started: Float64) {
    let initial = evaluate_network(network, training, batch_size)
    print("BENCHMARK_INITIAL_LOSS")
    print(initial.loss)
    var epoch = 0
    while epoch < epochs {
        var weighted_loss = 0.0
        var batch_index = 0
        let schedule = training.batches(size: batch_size, seed: seed + epoch, shuffle: true)
        for batch in schedule {
            let input = nn.place(batch.features())
            let target = nn.place(batch.one_hot())
            let loss = objective(network.forward(input), target)
            weighted_loss = weighted_loss + nn.item(loss) * Float64(batch.rows())
            nn.backward(loss)
            // One gradient/topology sample per epoch; no per-batch inspection overhead.
            if batch_index == schedule.count() - 1 { viz.capture(epoch + 1, loss) }
            network.update(state)
            batch_index = batch_index + 1
        }
        epoch = epoch + 1
        let held_out = evaluate_network(network, validation, batch_size)
        let online_loss = weighted_loss / Float64(training.rows())
        dev.progress(epoch: epoch, epochs: epochs, loss: online_loss, accuracy: held_out.accuracy, seconds: nn.clock() - started, validation_loss: held_out.loss)
        print("BENCHMARK_EPOCH")
        print(epoch)
        print(online_loss)
        print(held_out.loss)
        print(held_out.accuracy)
    }
    return evaluate_network(network, training, batch_size)
}
let final_training = train_network(network, state, training, validation, epochs, batch_size, seed, started)
print("BENCHMARK_FIT_SECONDS")
print(nn.clock() - started)
print("BENCHMARK_TRAIN_LOSS")
print(final_training.loss)
print("BENCHMARK_TRAIN_ACCURACY")
print(final_training.accuracy)
let test_started = nn.clock()
let evaluation = evaluate_network(network, testing, batch_size)
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())
dev.capture("Final model", network)
print(state.summary())