The gallery

use topos::{Numerics, Tape, Tensor};

fn main() {
    // A tape is a notebook of operations: every expression over a
    // `Value` appends one line. `weights` is a parameter, the thing
    // training will change; `input` is fed per run and carries a
    // default, which this file uses throughout.
    let tape: Tape<f32> = Tape::new();
    let weights = tape.parameter(Tensor::new([2, 2], [0.5_f32, -0.25, 0.25, 0.4]));
    let input = tape.input(Tensor::new([1, 2], [1.0_f32, 2.0]));
    let hidden = input.matmul(weights).tanh();
    let loss = (hidden * hidden).sum();

    banner("1. the spec");
    println!("One line per recorded operation: its number, its opcode, the numbers of");
    println!("the operands it reads, and its shape. Shapes were checked as each line");
    println!("was written; a mismatch panics at the expression that records it.\n");
    print!("{}", tape.describe());
    let spec_nodes = tape.len();

    // A `Value` borrows the tape and cannot outlive it. A `Symbol`
    // is the detached name every later phase speaks.
    let (weights, hidden, loss) = (weights.symbol(), hidden.symbol(), loss.symbol());

    banner("2. the derivative, as more spec");
    // The chain rule is applied by the same rules the interpreter's
    // reverse scan uses; here they record instead of compute. The
    // gradient becomes ordinary nodes: printable, schedulable,
    // emittable, and differentiable again.
    let adjoints = tape.differentiate(loss, [weights]);
    let gradient = adjoints.of(weights);
    println!(
        "`differentiate` appended {} nodes to the same tape; nothing above them changed.",
        tape.len() - spec_nodes
    );
    println!("The last one is d loss / d weights.\n");
    for line in tape.describe().lines().skip(spec_nodes) {
        println!("{line}");
    }
    println!("\nHow to read them:");
    println!("  - a Leaf holding one is the seed, spread over the shape the sum reduced");
    println!("  - `hidden * hidden` gets the product rule once per operand");
    println!("  - `1 - hidden * hidden` is the derivative of tanh");
    println!("  - the matmul rule pushes the result back as two transposed products,");
    println!("    one per operand");
    println!("  - each Leaf-then-Add pair starts a gradient at zero and adds one");
    println!("    consumer's share; that is how a value used twice collects both");

    // Sealing turns the tape into an immutable `Network`. From here
    // on nothing writes the graph: training only replaces the
    // caller-owned `Parameters`, which start from the recorded
    // initial values.
    let network = tape.into_network();
    let parameters = network.parameters();

    banner("3. the plan");
    // An entry declares what a run may read: its roots (the loss and
    // the gradient, in that order) plus any interior we also want to
    // look at. Reading an undeclared interior panics. That
    // declaration is the license the scheduler works under: the
    // declared must survive; the undeclared may be skipped, freed,
    // or fused.
    let declare = || network.entry(adjoints.roots()).observe([hidden]);
    let plan = declare().lower();
    println!("The same operations, scheduled. Each line also says what happens to the");
    println!("value once computed: kept for reading, or freed after its last use. The");
    println!("last line is the plan's own memory accounting: the most elements alive at");
    println!("once, against the total if nothing were ever freed.\n");
    print!("{}", plan.describe());
    let scheduled: Vec<usize> = plan.nodes().map(|node| node.symbol().index()).collect();
    let skipped: Vec<String> = network
        .nodes()
        .filter(|node| !scheduled.contains(&node.symbol().index()))
        .map(|node| node.symbol().index().to_string())
        .collect();
    if skipped.is_empty() {
        println!("\nEvery node is scheduled.");
    } else {
        println!(
            "\nNot scheduled: nodes {}. The rules also produced the gradient with",
            skipped.join(", ")
        );
        println!("respect to the input; nobody declared it, so it is never computed.");
    }
    let patterns = plan.patterns();
    if patterns.is_empty() {
        println!("Elected patterns: none. Fusion is an offer the catalog makes for shapes it");
        println!("recognizes (windows, batch-norm); a dense layer this small runs as primitives.");
    } else {
        println!("Elected patterns:");
        for pattern in &patterns {
            println!(
                "  {:?} rooted at node {}",
                pattern.kind(),
                pattern.root().index()
            );
        }
    }

    banner("4. the plan as StableHLO");
    // Emission is `describe` in another syntax: the same schedule,
    // written as text an industrial compiler accepts. Nothing
    // XLA-shaped lives in the crate; the text is the boundary.
    let module = plan
        .emit_stablehlo()
        .expect("every operation in this graph lowers");
    print!("{module}");
    println!("\nThe function takes the parameter and the input, in recording order, and");
    println!("returns exactly what the entry declared, in that order: the loss, the");
    println!("gradient, then `hidden`.");
    assert_eq!(plan.results(), &[loss, gradient, hidden][..]);
    assert!(module.contains("stablehlo.dot_general"));
    assert!(module.contains("stablehlo.tanh"));

    banner("5. two runs that must agree");
    // The interpreter over the whole spec is the oracle. Its bits are
    // the same in every build, on every platform, and they define
    // what the answer is. Its reverse scan applies the derivative
    // rules directly, without recording.
    let oracle = network.forward(&parameters, []);
    let engine_gradient = oracle.backward(loss).parameters(&parameters);

    // A plan under `Exact` must reproduce those bits. Under `Fast`,
    // the default, a backend feature may serve large products with
    // sums in a different order; that difference is a labeled
    // choice on the entry, never a silent effect of a feature flag.
    let exact = declare().numerics(Numerics::Exact).lower();
    let exact_run = exact.forward(&parameters, []);
    let fast_run = plan.forward(&parameters, []);
    let recorded_gradient = exact_run.recorded_gradients(&adjoints);

    println!("loss:");
    println!("  interpreter        {}", oracle.of(loss));
    println!("  exact plan         {}", exact_run.of(loss));
    println!("  fast plan          {}", fast_run.of(loss));
    println!("d loss / d weights:");
    println!("  engine scan        {}", engine_gradient.of(weights));
    println!("  recorded gradient  {}", recorded_gradient.of(weights));
    println!("hidden, declared readable: {}", exact_run.of(hidden));

    assert!(bits_equal(oracle.of(loss), exact_run.of(loss)));
    assert!(bits_equal(
        engine_gradient.of(weights),
        recorded_gradient.of(weights)
    ));
    println!("\nexact plan == interpreter, bit for bit: asserted");
    println!("recorded gradient == engine scan, bit for bit: asserted");
    if bits_equal(fast_run.of(loss), exact_run.of(loss)) {
        println!(
            "fast plan == exact plan: yes; no backend claimed any of this graph in this build"
        );
    } else {
        println!("fast plan == exact plan: no; a backend served part of this graph under Fast,");
        println!("and Exact is how you ask for the reference bits back");
    }

    // A few steps of gradient descent close the loop. The plan was
    // compiled once and holds no state: each step feeds the current
    // parameters, reads the recorded gradient, and mints the next
    // parameters. The network never changes, and a clone of the
    // state is all a what-if costs. Nothing broadcasts by accident,
    // so the scalar learning rate is spread over each gradient
    // explicitly.
    let learning_rate = Tensor::from(0.05_f32);
    let mut current = parameters.clone();
    let mut losses = Vec::new();
    for _ in 0..5 {
        let run = plan.forward(&current, []);
        losses.push(run.of(loss).scalar());
        let gradients = run.recorded_gradients(&adjoints);
        current = current.step(&gradients, |weight, gradient| {
            weight.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }
    losses.push(plan.forward(&current, []).of(loss).scalar());
    let trajectory: Vec<String> = losses.iter().map(|value| format!("{value:.4}")).collect();
    println!("\nloss before training, then after each of five steps of size 0.05:");
    println!("  {}", trajectory.join(", "));
    assert!(losses[losses.len() - 1] < losses[0]);

    banner("6. the spec is executable");
    // The printed lines are an instruction list, and walking them is
    // the interpreter: sources take their stored payloads, computed
    // nodes express their opcode over operands computed earlier. The
    // walk runs under the exact posture, as the whole-spec oracle
    // does by construction.
    let values = Numerics::exactly(|| {
        let mut values: Vec<Tensor<f32>> = Vec::new();
        for node in network.nodes() {
            let value = if node.is_source() {
                network
                    .payload(node.symbol())
                    .expect("sources hold payloads")
                    .clone()
            } else {
                let operands: Vec<&Tensor<f32>> = node
                    .operands()
                    .iter()
                    .map(|symbol| &values[symbol.index()])
                    .collect();
                node.opcode().express(&operands)
            };
            values.push(value);
        }
        values
    });
    assert!(bits_equal(oracle.of(loss), &values[loss.index()]));
    println!(
        "Walking all {} lines by hand, outside the crate, reproduces the interpreter's",
        values.len()
    );
    println!("loss {} bit for bit.", values[loss.index()]);
}

/// Prints a section heading.
fn banner(title: &str) {
    println!("\n== {title} ==\n");
}

/// Compares two tensors bit for bit, the only equality the oracle
/// promises: same shape, same bits in every element.
fn bits_equal(left: &Tensor<f32>, right: &Tensor<f32>) -> bool {
    left.shape() == right.shape()
        && left
            .to_vec()
            .iter()
            .zip(right.to_vec())
            .all(|(left, right)| left.to_bits() == right.to_bits())
}

use topos::Tape;

fn main() {
    let tape = Tape::new();

    // Leaves are the inputs of the graph: learnable parameters or data.
    // The tape owns their state; the returned values are `Copy` proxies
    // borrowing it.
    let a = tape.leaf(2.0_f64);
    let b = tape.leaf(3.0);
    let c = tape.leaf(4.0);
    println!("allocated {} leaves", tape.len());

    // Operators record computed nodes on the same tape. Proxies are
    // never consumed, so the same value can feed any number of expressions.
    let sum = a + b;
    let product = sum * c;
    let expression = -product + a * c;

    println!("chained -((a + b) * c) + a * c as {expression:?}");
    println!("the tape now holds {} values", tape.len());

    // Symbols are the names every phase after recording speaks; the
    // seal consumes the tape and hands back the immutable spec.
    let (a, b, c, expression) = (a.symbol(), b.symbol(), c.symbol(), expression.symbol());
    let network = tape.into_network();
    let parameters = network.parameters();

    // The forward pass materializes every payload into per-run storage,
    // leaving the network untouched.
    let run = network.forward(&parameters, []);
    println!("forward: expression = {}", run.of(expression));

    // The backward pass produces the gradient of the expression with
    // respect to every value. `a` feeds two subexpressions whose
    // contributions cancel exactly, hence its zero gradient.
    let gradients = run.backward(expression);
    println!(
        "gradients: d/da = {}, d/db = {}, d/dc = {}",
        gradients.of(a),
        gradients.of(b),
        gradients.of(c)
    );
}

use std::time::Instant;

use rayon::prelude::*;

use malevich::{Frame, Line, Plot};
use topos::model::{Detach, Tape, Tensor};

fn main() {
    // The whole recording is one closure; its return value is the
    // set of names that leave the tape, detached in one `detach` call.
    let (network, (w, b, loss, sample_losses)) = Tape::record(|tape| {
        // Learnable parameters, starting from zero.
        let w = tape.parameter(0.0_f64);
        let b = tape.parameter(0.0);

        // Training data for the target line `y = 2 * x + 1`, recorded
        // as plain leaves. Each sample's squared error is kept as a
        // separate target; the total loss is their sum.
        let samples = [(1.0, 3.0), (2.0, 5.0), (3.0, 7.0)];
        let mut sample_losses = Vec::new();
        for (x, y) in samples {
            let x = tape.leaf(x);
            let y = tape.leaf(y);
            let error = w * x + b - y;
            sample_losses.push(error * error);
        }
        // Values are `Copy`, so the per-sample losses fold into a total
        // loss with a plain reduce.
        let loss = sample_losses
            .iter()
            .copied()
            .reduce(|total, squared| total + squared)
            .expect("at least one sample");
        (w, b, loss, sample_losses).detach()
    });
    let parameters = network.parameters();

    // One run feeds many backward sweeps: each rayon thread
    // differentiates the same shared run for its own target.
    let run = network.forward(&parameters, []);
    let per_sample: Vec<f64> = sample_losses
        .par_iter()
        .map(|&sample_loss| {
            let gradients = run.backward(sample_loss);
            gradients.of(w).scalar()
        })
        .collect();
    let total_gradient = run.backward(loss).of(w).scalar();
    println!("per-sample d/dw, computed on separate threads: {per_sample:?}");
    println!(
        "their sum {} equals the total-loss d/dw {} by linearity",
        per_sample.iter().sum::<f64>(),
        total_gradient
    );

    // Parallel training: each learning rate clones the parameter state
    // and descends independently over the one shared spec, keeping its
    // whole loss history for the chart.
    let learning_rates = [0.005, 0.02, 0.05];
    let training = Instant::now();
    let runs: Vec<(f64, Vec<f64>, f64, f64)> = learning_rates
        .par_iter()
        .map(|&learning_rate| {
            let mut parameters = parameters.clone();
            let mut losses = Vec::with_capacity(501);
            for _ in 0..500 {
                let run = network.forward(&parameters, []);
                losses.push(run.of(loss).scalar());
                let gradients = run.backward(loss).parameters(&parameters);
                parameters = parameters.step(&gradients, |parameter, gradient| {
                    parameter.clone() - gradient.clone() * Tensor::from(learning_rate)
                });
            }
            let run = network.forward(&parameters, []);
            losses.push(run.of(loss).scalar());
            (
                learning_rate,
                losses,
                parameters.of(w).scalar(),
                parameters.of(b).scalar(),
            )
        })
        .collect();

    println!(
        "trained {} states of 500 steps in {:.3}s",
        learning_rates.len(),
        training.elapsed().as_secs_f64()
    );

    println!("parallel training on cloned states (target: w = 2, b = 1):");
    for (learning_rate, losses, w, b) in &runs {
        let final_loss = losses.last().expect("every run records its final loss");
        println!("  lr = {learning_rate:5.3}: loss = {final_loss:.6}, w = {w:.3}, b = {b:.3}");
    }

    // The same three descents as curves: on a log scale a constant
    // convergence rate is a straight line, so the slope is the rate.
    let mut plot = Plot::new()
        .title("gradient descent per learning rate")
        .x_label("step")
        .y_label("loss")
        .log_y();
    for (learning_rate, losses, ..) in &runs {
        plot = plot.layer(Line::y(&losses[..]).label(format!("lr = {learning_rate}")));
    }
    println!("{}", plot.render_best(&Frame::detect()));
}

use std::time::Instant;

use malevich::{Frame, Line, Plot, Points};
use topos::{Activation, Mlp, Module, Shape, Tape, Tensor, init};

/// How many noisy samples the model trains on; the chart grid reuses
/// the count so the recorded expression serves both.
const SAMPLE_LEN: usize = 96;

/// The half-width of the sampled domain, centered on zero.
const DOMAIN: f64 = 3.0;

/// The standard deviation of the noise added to the sampled curve.
const NOISE_DEVIATION: f64 = 0.1;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 16;

/// How many full-batch gradient descent steps the fit takes.
const STEP_COUNT: usize = 3000;

fn main() {
    // The dataset: x uniform over the domain and y = sin x plus noise,
    // both drawn from the same seeded initializers the parameters use.
    let features: Tensor<f32> = init::uniform(3, DOMAIN)(&Shape::new([SAMPLE_LEN, 1]));
    let noise: Tensor<f32> = init::normal(5, NOISE_DEVIATION)(&Shape::new([SAMPLE_LEN, 1]));
    let feature_values = features.to_vec();
    let target_values: Vec<f32> = feature_values
        .iter()
        .zip(noise.to_vec())
        .map(|(&x, noise)| x.sin() + noise)
        .collect();

    let tape: Tape<f32> = Tape::new();
    let mlp = Mlp::new(
        &tape,
        &[1, HIDDEN_LEN, 1],
        Activation::Tanh,
        init::xavier(7),
    );

    let input = tape.input(features);
    let expected = tape.input(Tensor::new([SAMPLE_LEN, 1], target_values.clone()));
    let predicted = mlp.express(input);
    let error = predicted - expected;
    let loss = (error * error).sum();

    let (input, predicted, loss) = (input.symbol(), predicted.symbol(), loss.symbol());
    let network = tape.into_network();
    let mut parameters = network.parameters();

    let learning_rate = Tensor::new([], [0.001]);
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..STEP_COUNT {
        let run = network.forward(&parameters, []);
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        if step % (STEP_COUNT / 5) == 0 {
            println!("step {step:4}: loss = {batch_loss:.4}");
        }
        let gradients = run.backward(loss).parameters(&parameters);
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    println!(
        "{}",
        Plot::new()
            .layer(Line::y(&losses[..]).label("full batch"))
            .title("regression training")
            .x_label("step")
            .y_label("sum of squared errors")
            .render_best(&Frame::detect())
    );

    // The fit, read over an evenly spaced grid fed through the trained
    // expression in place of the training samples.
    let grid: Vec<f32> = (0..SAMPLE_LEN)
        .map(|index| (((index as f64 / (SAMPLE_LEN - 1) as f64) * 2.0 - 1.0) * DOMAIN) as f32)
        .collect();
    let run = network.forward(
        &parameters,
        [(input, Tensor::new([SAMPLE_LEN, 1], grid.clone()))],
    );
    let fit = run.of(predicted).to_vec();

    println!(
        "{}",
        Plot::new()
            .layer(Points::xy(&feature_values[..], &target_values[..]).label("samples"))
            .layer(Line::function(-DOMAIN..DOMAIN, f64::sin).label("sin x"))
            .layer(Line::xy(&grid[..], &fit[..]).label("fit"))
            .title("a tanh mlp fit to noisy sin x")
            .x_label("x")
            .render_best(&Frame::detect())
    );
}

use std::time::Instant;

use malevich::{Cells, Frame, Line, Plot};
use topos::{Activation, Mlp, Module, Tape, Tensor, init};

/// The resolution of the decision-surface chart: how many grid cells
/// span the unit square along each axis.
const SURFACE_COLUMNS: usize = 24;
const SURFACE_ROWS: usize = 12;

fn main() {
    let tape = Tape::new();
    // A deterministic seeded initializer keeps the run reproducible
    // while hidden-unit symmetry still breaks.
    let mlp = Mlp::new(&tape, &[2, 4, 1], Activation::Tanh, init::uniform(7, 0.5));

    // Declared inputs: the minibatch arrives per run, so the defaults
    // only fix the shapes — two samples of two features, two targets.
    let x = tape.input(Tensor::filled([2, 2], 0.0_f32));
    let y = tape.input(Tensor::filled([2, 1], 0.0));

    let predicted = mlp.express(x);
    let error = predicted - y;
    let loss = (error * error).sum();

    let (x, y, loss, predicted) = (x.symbol(), y.symbol(), loss.symbol(), predicted.symbol());
    let recorded_nodes = tape.len();
    let network = tape.into_network();
    let mut parameters = network.parameters();

    // The XOR truth table with targets in [-1, 1], split into two
    // minibatches that alternate across training steps.
    let minibatches = [
        (
            Tensor::new([2, 2], [0.0, 0.0, 0.0, 1.0]),
            Tensor::new([2, 1], [-1.0, 1.0]),
        ),
        (
            Tensor::new([2, 2], [1.0, 0.0, 1.0, 1.0]),
            Tensor::new([2, 1], [1.0, -1.0]),
        ),
    ];

    let learning_rate = Tensor::new([], [0.05]);
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..4000 {
        let (batch_x, batch_y) = &minibatches[step % minibatches.len()];
        let run = network.forward(&parameters, [(x, batch_x.clone()), (y, batch_y.clone())]);
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        if step % 800 == 0 {
            println!("step {step:4}: minibatch loss = {batch_loss:.6}");
        }
        let gradients = run.backward(loss).parameters(&parameters);
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    assert_eq!(network.len(), recorded_nodes);
    println!("the tape held {recorded_nodes} nodes through every step");

    // On a log scale the descent's exponential tail is a straight
    // line; the two alternating minibatches braid around it.
    println!(
        "{}",
        Plot::new()
            .layer(Line::y(&losses[..]).label("minibatch"))
            .title("xor training")
            .x_label("step")
            .y_label("sum of squared errors")
            .log_y()
            .render_best(&Frame::detect())
    );

    println!("predictions (target in parentheses):");
    for (batch_x, batch_y) in &minibatches {
        let run = network.forward(&parameters, [(x, batch_x.clone()), (y, batch_y.clone())]);
        let outputs = run.of(predicted);
        for (sample, (prediction, target)) in outputs.iter().zip(batch_y.iter()).enumerate() {
            let features = batch_x.as_slice().expect("a fed minibatch is contiguous");
            let features = &features[sample * 2..sample * 2 + 2];
            println!(
                "  {:?} -> {prediction:+.3} ({target:+.0})",
                (features[0], features[1])
            );
        }
    }

    // The learned surface, sampled at grid cell centers over the unit
    // square through the same two-row expression training used, fed
    // two points per run.
    let centers: Vec<(f32, f32)> = (0..SURFACE_ROWS)
        .flat_map(|row| {
            (0..SURFACE_COLUMNS).map(move |column| {
                (
                    (column as f32 + 0.5) / SURFACE_COLUMNS as f32,
                    (row as f32 + 0.5) / SURFACE_ROWS as f32,
                )
            })
        })
        .collect();
    let mut surface = Vec::with_capacity(centers.len());
    for pair in centers.chunks(2) {
        let &[(x0, y0), (x1, y1)] = pair else {
            unreachable!("the even grid splits into exact pairs");
        };
        let run = network.forward(&parameters, [(x, Tensor::new([2, 2], [x0, y0, x1, y1]))]);
        surface.extend(run.of(predicted).to_vec());
    }
    println!(
        "{}",
        Plot::new()
            .layer(Cells::matrix(SURFACE_COLUMNS, surface).extents((0.0, 1.0), (0.0, 1.0)))
            .colorbar()
            .title("the learned xor surface")
            .render_best(&Frame::detect())
    );
}

use std::f32::consts::PI;
use std::time::Instant;

use malevich::{Cells, Frame, Line, Plot, Points};
use topos::{Activation, Mlp, Module, Shape, Tape, Tensor, init};

/// How many points each half-moon holds.
const MOON_LEN: usize = 100;

/// The standard deviation of the noise scattering the moons.
const NOISE_DEVIATION: f64 = 0.1;

/// The resolution of the decision-surface chart grid.
const SURFACE_COLUMNS: usize = 48;
const SURFACE_ROWS: usize = 16;

/// The data window the surface rasterizes: both moons with a margin.
const X_SPAN: (f32, f32) = (-1.5, 2.5);
const Y_SPAN: (f32, f32) = (-1.0, 1.5);

/// How many full-batch gradient descent steps the training takes.
const STEP_COUNT: usize = 3000;

/// Builds one half-moon of `MOON_LEN` noisy points: the upper arch
/// when `flipped` is false, the interleaved lower arch when true.
fn moon(flipped: bool, noise: &[f32]) -> (Vec<f32>, Vec<f32>) {
    let mut moon_x = Vec::with_capacity(MOON_LEN);
    let mut moon_y = Vec::with_capacity(MOON_LEN);
    for index in 0..MOON_LEN {
        let angle = PI * index as f32 / (MOON_LEN - 1) as f32;
        let (x, y) = if flipped {
            (1.0 - libm::cosf(angle), 0.5 - libm::sinf(angle))
        } else {
            (libm::cosf(angle), libm::sinf(angle))
        };
        moon_x.push(x + noise[index * 2]);
        moon_y.push(y + noise[index * 2 + 1]);
    }
    (moon_x, moon_y)
}

fn main() {
    let noise: Tensor<f32> = init::normal(5, NOISE_DEVIATION)(&Shape::new([2 * MOON_LEN, 2]));
    let noise = noise.to_vec();
    let (upper_x, upper_y) = moon(false, &noise[..2 * MOON_LEN]);
    let (lower_x, lower_y) = moon(true, &noise[2 * MOON_LEN..]);

    // The batch interleaves each point's coordinates row by row, the
    // upper moon labeled `+1` and the lower `-1`.
    let mut feature_values = Vec::with_capacity(4 * MOON_LEN);
    for (x, y) in upper_x
        .iter()
        .zip(&upper_y)
        .chain(lower_x.iter().zip(&lower_y))
    {
        feature_values.push(*x);
        feature_values.push(*y);
    }
    let mut target_values = vec![1.0_f32; MOON_LEN];
    target_values.extend(vec![-1.0; MOON_LEN]);

    let tape: Tape<f32> = Tape::new();
    let mlp = Mlp::new(&tape, &[2, 16, 16, 1], Activation::Tanh, init::xavier(7));

    let input = tape.input(Tensor::new([2 * MOON_LEN, 2], feature_values));
    let expected = tape.input(Tensor::new([2 * MOON_LEN, 1], target_values.clone()));
    let predicted = mlp.express(input);
    let error = predicted - expected;
    let loss = (error * error).sum();

    // The rasterizing twin: the same parameters expressed over the
    // chart grid's cell centers, recorded once next to the training
    // expression.
    let mut surface_centers = Vec::with_capacity(2 * SURFACE_COLUMNS * SURFACE_ROWS);
    for row in 0..SURFACE_ROWS {
        for column in 0..SURFACE_COLUMNS {
            let fraction_x = (column as f32 + 0.5) / SURFACE_COLUMNS as f32;
            let fraction_y = (row as f32 + 0.5) / SURFACE_ROWS as f32;
            surface_centers.push(X_SPAN.0 + fraction_x * (X_SPAN.1 - X_SPAN.0));
            surface_centers.push(Y_SPAN.0 + fraction_y * (Y_SPAN.1 - Y_SPAN.0));
        }
    }
    let surface_input = tape.input(Tensor::new(
        [SURFACE_COLUMNS * SURFACE_ROWS, 2],
        surface_centers,
    ));
    let surface_predicted = mlp.express(surface_input);

    let (predicted, surface_predicted, loss) = (
        predicted.symbol(),
        surface_predicted.symbol(),
        loss.symbol(),
    );
    let network = tape.into_network();
    let mut parameters = network.parameters();

    println!(
        "{}",
        Plot::new()
            .layer(Points::xy(&upper_x[..], &upper_y[..]).label("class +1"))
            .layer(Points::xy(&lower_x[..], &lower_y[..]).label("class -1"))
            .x_domain(f64::from(X_SPAN.0), f64::from(X_SPAN.1))
            .y_domain(f64::from(Y_SPAN.0), f64::from(Y_SPAN.1))
            .title("the two moons")
            .render_best(&Frame::detect())
    );

    let learning_rate = Tensor::new([], [0.0003]);
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..STEP_COUNT {
        let run = network.forward(&parameters, []);
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        if step % (STEP_COUNT / 5) == 0 {
            println!("step {step:4}: loss = {batch_loss:.4}");
        }
        let gradients = run.backward(loss).parameters(&parameters);
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    println!(
        "{}",
        Plot::new()
            .layer(Line::y(&losses[..]).label("full batch"))
            .title("two moons training")
            .x_label("step")
            .y_label("sum of squared errors")
            .render_best(&Frame::detect())
    );

    // One forward serves both readouts: the training points' signs for
    // the accuracy line and the grid twin for the surface chart.
    let run = network.forward(&parameters, []);
    let classified = run
        .of(predicted)
        .to_vec()
        .iter()
        .zip(&target_values)
        .filter(|(prediction, target)| prediction.signum() == target.signum())
        .count();
    println!(
        "classified {classified} of {} training points",
        2 * MOON_LEN
    );

    let surface = run.of(surface_predicted).to_vec();
    println!(
        "{}",
        Plot::new()
            .layer(Cells::matrix(SURFACE_COLUMNS, surface).extents(
                (f64::from(X_SPAN.0), f64::from(X_SPAN.1)),
                (f64::from(Y_SPAN.0), f64::from(Y_SPAN.1)),
            ))
            .colorbar()
            .title("the learned decision surface")
            .render_best(&Frame::detect())
    );
}

use std::ops::{Add, Div, Mul, Neg, Sub};

use topos::{Detach, Network, Node, Numerics, Recordable, Shape, Symbol, Tape, Tensor, Trace};

/// A dual number over any recordable payload: the primal value and
/// the tangent riding along. The chain rule is the arithmetic.
#[derive(Clone, Debug)]
struct Dual<R> {
    primal: R,
    tangent: R,
}

impl<R: Recordable> Dual<R> {
    /// Wraps a value carried unchanged through differentiation: a
    /// constant's tangent is zero.
    fn constant(primal: R) -> Self {
        let tangent = primal.zero_like();
        Self { primal, tangent }
    }
}

impl<R: Recordable> Add for Dual<R> {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        Self {
            primal: self.primal + rhs.primal,
            tangent: self.tangent + rhs.tangent,
        }
    }
}

impl<R: Recordable> Sub for Dual<R> {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        Self {
            primal: self.primal - rhs.primal,
            tangent: self.tangent - rhs.tangent,
        }
    }
}

impl<R: Recordable> Mul for Dual<R> {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self {
        Self {
            primal: self.primal.clone() * rhs.primal.clone(),
            tangent: self.primal * rhs.tangent + self.tangent * rhs.primal,
        }
    }
}

impl<R: Recordable> Div for Dual<R> {
    type Output = Self;
    fn div(self, rhs: Self) -> Self {
        let primal = self.primal / rhs.primal.clone();
        let tangent = (self.tangent - primal.clone() * rhs.tangent) / rhs.primal;
        Self { primal, tangent }
    }
}

impl<R: Recordable> Neg for Dual<R> {
    type Output = Self;
    fn neg(self) -> Self {
        Self {
            primal: -self.primal,
            tangent: -self.tangent,
        }
    }
}

impl<R: Recordable> Recordable for Dual<R> {
    fn shape(&self) -> Shape {
        self.primal.shape()
    }

    fn zero_like(&self) -> Self {
        Self::constant(self.primal.zero_like())
    }

    fn one_like(&self) -> Self {
        Self::constant(self.primal.one_like())
    }

    fn exp(&self) -> Self {
        let primal = self.primal.exp();
        let tangent = primal.clone() * self.tangent.clone();
        Self { primal, tangent }
    }

    fn ln(&self) -> Self {
        Self {
            primal: self.primal.ln(),
            tangent: self.tangent.clone() / self.primal.clone(),
        }
    }

    fn sqrt(&self) -> Self {
        let primal = self.primal.sqrt();
        let tangent = self.tangent.clone() / (primal.clone() + primal.clone());
        Self { primal, tangent }
    }

    fn tanh(&self) -> Self {
        let primal = self.primal.tanh();
        let tangent = self.tangent.clone() * (primal.one_like() - primal.clone() * primal.clone());
        Self { primal, tangent }
    }

    fn sin(&self) -> Self {
        Self {
            primal: self.primal.sin(),
            tangent: self.tangent.clone() * self.primal.cos(),
        }
    }

    fn cos(&self) -> Self {
        Self {
            primal: self.primal.cos(),
            tangent: -(self.tangent.clone() * self.primal.sin()),
        }
    }

    fn log1p(&self) -> Self {
        Self {
            primal: self.primal.log1p(),
            tangent: self.tangent.clone() / (self.primal.one_like() + self.primal.clone()),
        }
    }

    fn expm1(&self) -> Self {
        let primal = self.primal.expm1();
        let tangent = self.tangent.clone() * (primal.clone() + primal.one_like());
        Self { primal, tangent }
    }

    fn erf(&self) -> Self {
        Self {
            primal: self.primal.erf(),
            tangent: self.tangent.clone() * self.primal.erf_derivative(),
        }
    }

    fn erf_derivative(&self) -> Self {
        // The scaled Gaussian differentiates to `-2x` times itself.
        let primal = self.primal.erf_derivative();
        let tangent =
            -((self.primal.clone() + self.primal.clone()) * primal.clone() * self.tangent.clone());
        Self { primal, tangent }
    }

    fn powf(&self, exponent: Self) -> Self {
        let primal = self.primal.powf(exponent.primal.clone());
        let tangent = primal.clone()
            * (exponent.tangent * self.primal.ln()
                + exponent.primal * (self.tangent.clone() / self.primal.clone()));
        Self { primal, tangent }
    }

    fn maximum(&self, other: &Self) -> Self {
        let winners = self.primal.step(&other.primal);
        Self {
            primal: self.primal.maximum(&other.primal),
            tangent: winners.clone() * self.tangent.clone()
                + (winners.one_like() - winners) * other.tangent.clone(),
        }
    }

    fn step(&self, threshold: &Self) -> Self {
        Self::constant(self.primal.step(&threshold.primal))
    }

    fn matmul(&self, rhs: &Self) -> Self {
        Self {
            primal: self.primal.matmul(&rhs.primal),
            tangent: self.tangent.matmul(&rhs.primal) + self.primal.matmul(&rhs.tangent),
        }
    }

    fn sum(&self) -> Self {
        Self {
            primal: self.primal.sum(),
            tangent: self.tangent.sum(),
        }
    }

    fn sum_along(&self, axis: usize) -> Self {
        Self {
            primal: self.primal.sum_along(axis),
            tangent: self.tangent.sum_along(axis),
        }
    }

    fn logsumexp(&self, axis: usize) -> Self {
        let primal = self.primal.logsumexp(axis);
        let extent = self.primal.shape().axes()[axis];
        // The softmax weights, from the stable primal itself.
        let weights = (self.primal.clone() - primal.broadcast_along(axis, extent)).exp();
        Self {
            tangent: (weights * self.tangent.clone()).sum_along(axis),
            primal,
        }
    }

    fn log_softmax(&self, axis: usize) -> Self {
        let primal = self.primal.log_softmax(axis);
        let extent = self.primal.shape().axes()[axis];
        let mass = self.tangent.sum_along(axis).broadcast_along(axis, extent);
        Self {
            tangent: self.tangent.clone() - primal.exp() * mass,
            primal,
        }
    }

    fn broadcast(&self, shape: Shape) -> Self {
        Self {
            primal: self.primal.broadcast(shape.clone()),
            tangent: self.tangent.broadcast(shape),
        }
    }

    fn broadcast_along(&self, axis: usize, extent: usize) -> Self {
        Self {
            primal: self.primal.broadcast_along(axis, extent),
            tangent: self.tangent.broadcast_along(axis, extent),
        }
    }

    fn reshape(&self, shape: Shape) -> Self {
        Self {
            primal: self.primal.reshape(shape.clone()),
            tangent: self.tangent.reshape(shape),
        }
    }

    fn permute(&self, order: &[usize]) -> Self {
        Self {
            primal: self.primal.permute(order),
            tangent: self.tangent.permute(order),
        }
    }

    fn narrow(&self, axis: usize, start: usize, len: usize) -> Self {
        Self {
            primal: self.primal.narrow(axis, start, len),
            tangent: self.tangent.narrow(axis, start, len),
        }
    }

    fn pad(&self, axis: usize, start: usize, full_extent: usize) -> Self {
        Self {
            primal: self.primal.pad(axis, start, full_extent),
            tangent: self.tangent.pad(axis, start, full_extent),
        }
    }

    fn unfold(&self, axis: usize, size: usize, step: usize, dilation: usize) -> Self {
        Self {
            primal: self.primal.unfold(axis, size, step, dilation),
            tangent: self.tangent.unfold(axis, size, step, dilation),
        }
    }

    fn fold(&self, axis: usize, size: usize, step: usize, dilation: usize, extent: usize) -> Self {
        Self {
            primal: self.primal.fold(axis, size, step, dilation, extent),
            tangent: self.tangent.fold(axis, size, step, dilation, extent),
        }
    }

    fn gather(&self, selection: &Self) -> Self {
        // The selection is data, not a differentiable dependency.
        Self {
            primal: self.primal.gather(&selection.primal),
            tangent: self.tangent.gather(&selection.primal),
        }
    }

    fn scatter(&self, selection: &Self) -> Self {
        Self {
            primal: self.primal.scatter(&selection.primal),
            tangent: self.tangent.scatter(&selection.primal),
        }
    }
}

/// Replays every node of a spec over dual payloads: `source`
/// supplies the sources' duals, and every computed node is one
/// `Opcode::express` over its operands' duals.
fn jvp_walk<R: Recordable>(
    nodes: &[Node],
    mut source: impl FnMut(&Node) -> Dual<R>,
) -> Vec<Dual<R>> {
    let mut duals: Vec<Dual<R>> = Vec::new();
    for node in nodes {
        let dual = if node.is_source() {
            source(node)
        } else {
            let operands: Vec<&Dual<R>> = node
                .operands()
                .iter()
                .map(|symbol| &duals[symbol.index()])
                .collect();
            node.opcode().express(&operands)
        };
        duals.push(dual);
    }
    duals
}

/// The eager JVP: walks the spec over `Dual<Tensor>` and answers
/// the tangent at `read` for a `seed` planted at `wrt`.
fn jvp_eager(network: &Network<f64>, wrt: Symbol, seed: &Tensor<f64>, read: Symbol) -> Tensor<f64> {
    let nodes: Vec<Node> = network.nodes().collect();
    let duals = Numerics::exactly(|| {
        jvp_walk(&nodes, |node| {
            let primal = network
                .payload(node.symbol())
                .expect("sources hold payloads")
                .clone();
            let tangent = if node.symbol() == wrt {
                seed.clone()
            } else {
                primal.zero_like()
            };
            Dual { primal, tangent }
        })
    });
    duals[read.index()].tangent.clone()
}

fn assert_bitwise(expected: &Tensor<f64>, computed: &Tensor<f64>, subject: &str) {
    let expected = expected.to_vec();
    let computed = computed.to_vec();
    assert_eq!(expected.len(), computed.len(), "{subject}: length differs");
    for (expected, computed) in expected.iter().zip(&computed) {
        assert_eq!(
            expected.to_bits(),
            computed.to_bits(),
            "{subject}: {computed} differs from {expected}"
        );
    }
}

fn main() {
    // Dyadic values keep every product and short sum exact, so the
    // gradings below can demand bits, not tolerances (the
    // transcendental part compares two runs of the same arithmetic,
    // which is bitwise regardless).

    // Part one: a transcendental spec, eager against recorded. The
    // same dual arithmetic runs over `Dual<Tensor>` (computing) and
    // `Dual<Trace>` (recording tangent spec beside resolved
    // sources); the engine then runs the recorded form, and the two
    // answers must agree bit for bit — the closure contract of an
    // out-of-tree mode, exactly as `differentiate` is welded to
    // `Run::backward`.
    let (network, [weights, _inputs, loss]) = Tape::record(|tape| {
        let weights = tape.parameter(Tensor::new(
            [3, 2],
            vec![0.5, -0.25, 0.125, 0.75, -0.5, 0.25],
        ));
        let inputs = tape.input(Tensor::new([2, 3], vec![0.5, -1.0, 0.25, -0.75, 0.5, 1.25]));
        let scores = inputs.matmul(weights).tanh().log_softmax(1);
        let loss = (scores * scores).sum();
        [weights, inputs, loss].detach()
    });
    let seed = Tensor::new([3, 2], vec![0.25, -0.5, 1.0, 0.125, -0.25, 0.5]);

    let eager = jvp_eager(&network, weights, &seed, loss);

    let nodes: Vec<Node> = network.nodes().collect();
    let tape = network.into_tape();
    let duals = jvp_walk(&nodes, |node| {
        // Sources resolve into the same family — parameters are
        // shared, not duplicated — and their tangents are fresh
        // leaves: the seed at `wrt`, zero elsewhere.
        let primal = Trace::of(tape.resolve(node.symbol()));
        let payload = tape.payload(node.symbol()).expect("sources hold payloads");
        let tangent_payload = if node.symbol() == weights {
            seed.clone()
        } else {
            payload.zero_like()
        };
        let tangent = Trace::of(tape.leaf(tangent_payload));
        Dual { primal, tangent }
    });
    let tangent_symbol = duals[loss.index()].tangent.value().symbol();
    drop(duals);
    let jvp_network = tape.into_network();
    let recorded_run = jvp_network.forward(&jvp_network.parameters(), []);
    assert_bitwise(
        &eager,
        recorded_run.of(tangent_symbol),
        "recorded jvp against eager",
    );
    println!("eager and recorded jvp agree bitwise: {}", eager.scalar());

    // Part two: on a rational spec the directional derivative is
    // also the reverse gradient contracted with the seed, and at
    // dyadic values every route computes the same rational number —
    // so forward mode and reverse mode agree bit for bit.
    let (cubic, [x, value]) = Tape::record(|tape| {
        let x = tape.parameter(Tensor::new([3], vec![0.5, -0.25, 1.5]));
        let value = (x * x * x).sum();
        [x, value].detach()
    });
    let direction = Tensor::new([3], vec![0.25, 1.0, -0.5]);
    let forward_slope = jvp_eager(&cubic, x, &direction, value);
    let run = cubic.forward(&cubic.parameters(), []);
    let gradients = run.backward(value);
    let reverse_slope = Numerics::exactly(|| (gradients.of(x).clone() * direction.clone()).sum());
    assert_bitwise(&forward_slope, &reverse_slope, "slope against the gradient");
    println!(
        "forward slope equals <gradient, seed> bitwise: {}",
        forward_slope.scalar()
    );

    // Part three: the recorded gradient is just more spec, so
    // forward mode walks straight over it — forward-over-reverse
    // and reverse-over-reverse compute the same Hessian-vector
    // product, bit for bit at dyadic values.
    let tape: Tape<f64> = Tape::new();
    let x = tape.parameter(Tensor::new([3], vec![0.5, -0.25, 1.5]));
    let value = (x * x * x).sum();
    let adjoints = tape.differentiate(value, [x]);
    let &(_, gradient_symbol) = &adjoints.pairs()[0];
    let direction_leaf = tape.leaf(direction.clone());
    let hessian_adjoints = tape.vjp(gradient_symbol, direction_leaf, [x]);
    let &(_, reverse_hvp_symbol) = &hessian_adjoints.pairs()[0];
    let x = x.symbol();
    let network = tape.into_network();

    let run = network.forward(&network.parameters(), []);
    let reverse_hvp = run.of(reverse_hvp_symbol);
    let forward_hvp = jvp_eager(&network, x, &direction, gradient_symbol);
    assert_bitwise(
        reverse_hvp,
        &forward_hvp,
        "forward-over-reverse against reverse-over-reverse",
    );
    println!(
        "hessian-vector product agrees across modes: {:?}",
        forward_hvp.to_vec()
    );
}

use std::ops::{Add, Div, Mul, Neg, Sub};

use topos::{Detach, Differentiable, Element, Elementary, Tape};

/// A dual number: the primal value and the tangent riding along —
/// how fast the value changes when the variable of differentiation
/// nudges. The chain rule is the arithmetic.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct Dual {
    pub primal: f64,
    pub tangent: f64,
}

impl Dual {
    /// A constant carried unchanged through differentiation: zero
    /// tangent.
    pub fn value(primal: f64) -> Self {
        Self {
            primal,
            tangent: 0.0,
        }
    }

    /// The variable of differentiation: unit tangent.
    pub fn var(primal: f64) -> Self {
        Self {
            primal,
            tangent: 1.0,
        }
    }
}

impl Add for Dual {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        Self {
            primal: self.primal + rhs.primal,
            tangent: self.tangent + rhs.tangent,
        }
    }
}

impl Sub for Dual {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        Self {
            primal: self.primal - rhs.primal,
            tangent: self.tangent - rhs.tangent,
        }
    }
}

impl Mul for Dual {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self {
        Self {
            primal: self.primal * rhs.primal,
            tangent: self.primal * rhs.tangent + self.tangent * rhs.primal,
        }
    }
}

impl Div for Dual {
    type Output = Self;
    fn div(self, rhs: Self) -> Self {
        let primal = self.primal / rhs.primal;
        Self {
            tangent: (self.tangent - primal * rhs.tangent) / rhs.primal,
            primal,
        }
    }
}

impl Neg for Dual {
    type Output = Self;
    fn neg(self) -> Self {
        Self {
            primal: -self.primal,
            tangent: -self.tangent,
        }
    }
}

impl Differentiable for Dual {
    // Duals accumulate in their own precision: promoting a term
    // keeps its tangent, so an inner product's tangent is the sum of
    // the product-rule terms — exactly the derivative of the sum.
    type Accumulator = Self;

    fn promote(&self) -> Self {
        *self
    }

    fn demote(accumulated: Self) -> Self {
        accumulated
    }

    fn zero() -> Self {
        Self::value(0.0)
    }

    fn one() -> Self {
        Self::value(1.0)
    }

    fn from_count(count: usize) -> Self {
        Self::value(count as f64)
    }

    fn is_count(&self, count: usize) -> bool {
        self.primal == count as f64 && self.tangent == 0.0
    }
}

/// Every map is the chain rule on the primal's own kernel, called
/// through `Elementary` explicitly — `f64`'s *inherent* `exp` is the
/// standard library's, whose last bits vary by platform, while the
/// trait's rides `libm` — so the primal half of a dual run is
/// bit-identical to the `f64` run of the same graph.
impl Elementary for Dual {
    fn exp(&self) -> Self {
        let primal = Elementary::exp(&self.primal);
        Self {
            tangent: primal * self.tangent,
            primal,
        }
    }

    fn ln(&self) -> Self {
        Self {
            primal: Elementary::ln(&self.primal),
            tangent: self.tangent / self.primal,
        }
    }

    fn sqrt(&self) -> Self {
        let primal = Elementary::sqrt(&self.primal);
        Self {
            tangent: self.tangent / (primal + primal),
            primal,
        }
    }

    fn tanh(&self) -> Self {
        let primal = Elementary::tanh(&self.primal);
        Self {
            tangent: self.tangent * (1.0 - primal * primal),
            primal,
        }
    }

    fn sin(&self) -> Self {
        Self {
            primal: Elementary::sin(&self.primal),
            tangent: self.tangent * Elementary::cos(&self.primal),
        }
    }

    fn cos(&self) -> Self {
        Self {
            primal: Elementary::cos(&self.primal),
            tangent: -(self.tangent * Elementary::sin(&self.primal)),
        }
    }

    fn log1p(&self) -> Self {
        Self {
            primal: Elementary::log1p(&self.primal),
            tangent: self.tangent / (1.0 + self.primal),
        }
    }

    fn expm1(&self) -> Self {
        let primal = Elementary::expm1(&self.primal);
        Self {
            tangent: self.tangent * (primal + 1.0),
            primal,
        }
    }

    fn erf(&self) -> Self {
        Self {
            primal: Elementary::erf(&self.primal),
            tangent: self.tangent * Elementary::erf_derivative(&self.primal),
        }
    }

    fn erf_derivative(&self) -> Self {
        let primal = Elementary::erf_derivative(&self.primal);
        Self {
            tangent: self.tangent * (-2.0 * self.primal) * primal,
            primal,
        }
    }

    /// The dual of `x^y` on a positive base; elsewhere the `ln`
    /// answers `NaN`, mirroring the operation's own documented
    /// domain.
    fn powf(&self, exponent: Self) -> Self {
        let primal = Elementary::powf(&self.primal, exponent.primal);
        Self {
            tangent: primal
                * (exponent.tangent * Elementary::ln(&self.primal)
                    + exponent.primal * self.tangent / self.primal),
            primal,
        }
    }

    /// The primal wins by the `f64` rule; the tangent follows the
    /// winner, ties to the left like `step`'s ties-answer-one.
    fn maximum(&self, other: &Self) -> Self {
        Self {
            primal: Elementary::maximum(&self.primal, &other.primal),
            tangent: if self.primal >= other.primal {
                self.tangent
            } else {
                other.tangent
            },
        }
    }

    fn step(&self, threshold: &Self) -> Self {
        // The indicator is piecewise constant: zero tangent.
        Self::value(Elementary::step(&self.primal, &threshold.primal))
    }
}

impl Element for Dual {}

fn main() {
    // The README scalar chain, `loss = (w * x - y)^2`, at three
    // dyadic points; the derivative with respect to `w` three ways.
    // The points keep the error nonzero: at a zero error the two
    // modes agree on the value but not the sign of zero (the
    // product rule multiplies into `-0.0` where the reverse
    // accumulation adds into `+0.0`), and the grading demands bits.
    for (w0, x0, y0) in [(0.5, 0.25, 1.5), (-0.75, 1.25, 0.5), (2.0, -0.5, -1.25)] {
        // The f64 twin carries both reverse-mode routes: the engine
        // scan and the recorded gradient.
        let tape: Tape<f64> = Tape::new();
        let w = tape.parameter(w0);
        let x = tape.input(x0);
        let y = tape.input(y0);
        let error = w * x - y;
        let loss = error * error;
        let adjoints = tape.differentiate(loss, [w]);
        let (w, loss) = (w.symbol(), loss.symbol());
        let network = tape.into_network();
        let run = network.forward(&network.parameters(), []);
        let engine = run.backward(loss).of(w).scalar();
        let recorded = run.of(adjoints.pairs()[0].1).scalar();

        // The dual twin records the identical primal spelling; the
        // derivative is read off the loss payload's tangent after an
        // ordinary forward run.
        let (dual_network, [dual_loss]) = Tape::record(|tape| {
            let w = tape.parameter(Dual::var(w0));
            let x = tape.input(Dual::value(x0));
            let y = tape.input(Dual::value(y0));
            let error = w * x - y;
            [error * error].detach()
        });
        let dual_run = dual_network.forward(&dual_network.parameters(), []);
        let slope = dual_run.of(dual_loss).scalar();

        assert_eq!(
            slope.tangent.to_bits(),
            engine.to_bits(),
            "the dual tangent must be the engine gradient, bit for bit"
        );
        assert_eq!(
            slope.tangent.to_bits(),
            recorded.to_bits(),
            "the dual tangent must be the recorded gradient, bit for bit"
        );
        println!(
            "d/dw (w*{x0} - {y0})^2 at w = {w0}: {} on all three routes",
            slope.tangent
        );
    }
    println!("forward mode by payload, reverse mode by scan: one answer");
}

use std::ops::{Add, Div, Mul, Neg, Sub};
use std::sync::atomic::{AtomicUsize, Ordering};

use topos::{Differentiable, Element, Elementary, GemmTask, Numerics, Tape, Tensor, reference};

/// How many gemm tasks the seam has been offered, across all threads.
static OFFERED: AtomicUsize = AtomicUsize::new(0);

/// An `f64` in audit mode: arithmetic delegates, and the acceleration
/// seam counts its offers before answering with the reference kernel,
/// so the example can prove the hook was consulted.
#[derive(Debug, Clone, Copy, PartialEq)]
struct Audited(f64);

impl Add for Audited {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        Audited(self.0 + rhs.0)
    }
}

impl Sub for Audited {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        Audited(self.0 - rhs.0)
    }
}

impl Mul for Audited {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self {
        Audited(self.0 * rhs.0)
    }
}

impl Div for Audited {
    type Output = Self;
    fn div(self, rhs: Self) -> Self {
        Audited(self.0 / rhs.0)
    }
}

impl Neg for Audited {
    type Output = Self;
    fn neg(self) -> Self {
        Audited(-self.0)
    }
}

impl Differentiable for Audited {
    type Accumulator = Self;

    fn promote(&self) -> Self {
        *self
    }

    fn demote(accumulated: Self) -> Self {
        accumulated
    }

    fn zero() -> Self {
        Audited(0.0)
    }

    fn one() -> Self {
        Audited(1.0)
    }

    fn from_count(count: usize) -> Self {
        Audited(count as f64)
    }

    fn is_count(&self, count: usize) -> bool {
        self.0 == count as f64
    }
}

impl Elementary for Audited {
    fn exp(&self) -> Self {
        Audited(Elementary::exp(&self.0))
    }

    fn ln(&self) -> Self {
        Audited(Elementary::ln(&self.0))
    }

    fn sqrt(&self) -> Self {
        Audited(Elementary::sqrt(&self.0))
    }

    fn tanh(&self) -> Self {
        Audited(Elementary::tanh(&self.0))
    }

    fn sin(&self) -> Self {
        Audited(Elementary::sin(&self.0))
    }

    fn cos(&self) -> Self {
        Audited(Elementary::cos(&self.0))
    }

    fn log1p(&self) -> Self {
        Audited(Elementary::log1p(&self.0))
    }

    fn expm1(&self) -> Self {
        Audited(Elementary::expm1(&self.0))
    }

    fn erf(&self) -> Self {
        Audited(Elementary::erf(&self.0))
    }

    fn erf_derivative(&self) -> Self {
        Audited(Elementary::erf_derivative(&self.0))
    }

    fn powf(&self, exponent: Self) -> Self {
        Audited(Elementary::powf(&self.0, exponent.0))
    }

    fn maximum(&self, other: &Self) -> Self {
        Audited(self.0.max(other.0))
    }

    fn step(&self, threshold: &Self) -> Self {
        Audited(if self.0 >= threshold.0 { 1.0 } else { 0.0 })
    }

    /// The acceleration seam: count the offer, then answer with the
    /// published reference kernel — the same bits the built-in path
    /// computes, which is what makes the differential test below
    /// meaningful.
    fn gemm(task: &GemmTask<'_, Self>) -> Option<Vec<Self>> {
        OFFERED.fetch_add(1, Ordering::Relaxed);
        Some(reference::multiply(task))
    }
}

impl Element for Audited {}

fn main() {
    // The differential test every in-crate element passes: the hook's
    // product must be bit-identical to the built-in slice path. The
    // `Exact` posture pins the built-in path to the reference bits, so
    // the comparison is against the oracle, not a backend.
    let elements: Vec<f64> = (0..64)
        .map(|at| ((at * 7 % 23) as f64 - 11.0) / 4.0)
        .collect();
    let audited_left = Tensor::new(
        [4, 16],
        elements.iter().map(|&e| Audited(e)).collect::<Vec<_>>(),
    );
    let audited_right = Tensor::new(
        [16, 4],
        elements
            .iter()
            .map(|&e| Audited(e * 0.5))
            .collect::<Vec<_>>(),
    );
    let product = audited_left.matmul(&audited_right);

    let left = Tensor::new([4, 16], elements.clone());
    let right = Tensor::new(
        [16, 4],
        elements.iter().map(|&e| e * 0.5).collect::<Vec<_>>(),
    );
    let oracle = Numerics::exactly(|| left.matmul(&right));

    let bits: Vec<u64> = product.to_vec().iter().map(|e| e.0.to_bits()).collect();
    let oracle_bits: Vec<u64> = oracle.to_vec().iter().map(|e| e.to_bits()).collect();
    assert_eq!(bits, oracle_bits);
    println!(
        "gemm hook consulted {} times; product matches the reference bit for bit",
        OFFERED.load(Ordering::Relaxed)
    );

    // The crate-root regression, recorded over the new element: the
    // whole training loop runs on `Audited` without one extra line.
    let tape: Tape<Audited> = Tape::new();
    let w = tape.parameter(Audited(0.0));
    let x = tape.input(Audited(0.0));
    let y = tape.input(Audited(0.0));
    let error = w * x - y;
    let loss = error * error;
    let (w, x, y, loss) = (w.symbol(), x.symbol(), y.symbol(), loss.symbol());
    let network = tape.into_network();
    let mut parameters = network.parameters();

    let samples = [(1.0, 2.0), (2.0, 4.0), (3.0, 6.0)];
    for step in 0..100 {
        let (sample_x, sample_y) = samples[step % samples.len()];
        let run = network.forward(
            &parameters,
            [(x, Audited(sample_x).into()), (y, Audited(sample_y).into())],
        );
        let gradients = run.backward(loss).parameters(&parameters);
        parameters = parameters.step(&gradients, |w, g| {
            w.clone() - g.clone() * Tensor::from(Audited(0.02))
        });
    }
    let learned = parameters.of(w).scalar().0;
    assert!((learned - 2.0).abs() < 1e-6);
    println!("trained w = {learned:.6} on the audited element (target 2)");
}

use std::time::Instant;

use topos::{Backend, Shape, Tape, Tensor, init};

fn main() {
    let mut accelerated = false;
    for backend in Backend::ALL {
        let status = backend.status();
        accelerated |= status.is_ok();
        match status {
            Ok(()) => println!("{backend:?}: ready"),
            Err(reason) => println!("{backend:?}: {reason}"),
        }
    }

    // On the built-in paths alone, the full workload would take
    // minutes per step; an eightfold shrink keeps the run honest
    // and short while measuring the same code path.
    let scale = if accelerated { 1 } else { 8 };
    let batch_len = 1024 / scale;
    let feature_len = 1024 / scale;
    let hidden_len = 2048 / scale;
    let class_len = 512 / scale;
    let step_flops = (2 * batch_len * feature_len * hidden_len
        + 2 * batch_len * hidden_len * class_len) as f64
        * 3.0;

    let tape: Tape<f32> = Tape::new();
    let mut initializer = init::kaiming(7);
    let hidden_weights = tape.parameter(initializer(&Shape::new([feature_len, hidden_len])));
    let output_weights = tape.parameter(initializer(&Shape::new([hidden_len, class_len])));
    let features = tape.input(Tensor::filled([batch_len, feature_len], 0.0));
    let targets = tape.input(Tensor::filled([batch_len, class_len], 0.0));

    let hidden = features.matmul(hidden_weights).tanh();
    let prediction = hidden.matmul(output_weights);
    let error = prediction - targets;
    let loss = (error * error).sum();

    let (features, loss) = (features.symbol(), loss.symbol());
    let network = tape.into_network();
    let mut parameters = network.parameters();
    let learning_rate = Tensor::new([], [0.0001_f32]);

    let mut batch = init::normal(11, 1.0);
    let batch = batch(&Shape::new([batch_len, feature_len]));

    // Phase one: the raw product, the number the backend ladder is
    // about. One `matmul` at this size is a single backend call.
    let side = 2048 / scale;
    let left = initializer(&Shape::new([side, side]));
    let right = initializer(&Shape::new([side, side]));
    let product_flops = 2.0 * (side as f64).powi(3);
    let rounds = if accelerated { 20 } else { 2 };
    let started = Instant::now();
    for _ in 0..rounds {
        std::hint::black_box(left.matmul(&right));
    }
    let per_product = started.elapsed().as_secs_f64() / rounds as f64;
    println!(
        "raw product {side} x {side}: {:.1} GFLOP/s",
        product_flops / per_product / 1e9
    );

    // Phase two: whole training steps — the honest Amdahl statement.
    // Once the products are accelerated, the elementwise operations
    // (2M scalar `tanh` calls above all) own the step time; that gap
    // is the roadmap's next tier, not the backends'. The steps run
    // through an engine-backward plan: `Network::forward` is the
    // exact proving road and never engages a backend, so the plan
    // tier is where a throughput measurement belongs.
    let plan = network.entry([loss]).backward().lower();
    let warmup = 2;
    let steps = 8;
    let mut started = Instant::now();
    for step in 0..warmup + steps {
        if step == warmup {
            started = Instant::now();
        }
        let run = plan.forward(&parameters, [(features, batch.clone())]);
        let gradients = run.backward(loss).parameters(&parameters);
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }
    let elapsed = started.elapsed().as_secs_f64();
    let per_step = elapsed / steps as f64;
    println!(
        "training step, batch {batch_len}, {feature_len} -> {hidden_len} -> {class_len}: \
         {per_step:.3} s/step, {:.1} GFLOP/s of matrix products",
        step_flops / per_step / 1e9
    );

    // Where one step's tasks landed: coverage declares what a backend
    // may serve, the tally reports what each one did.
    let (_, services) = Backend::tallied(|| {
        let run = plan.forward(&parameters, [(features, batch.clone())]);
        run.backward(loss)
    });
    for service in services {
        let server = match service.backend {
            Some(backend) => format!("{backend:?}"),
            None => "reference".to_string(),
        };
        println!(
            "dispatch: {:?} {:?} x{}: {server}",
            service.formula, service.precision, service.count
        );
    }
}

mod chart;
mod corpus;

use std::time::Instant;

use malevich::{Cells, Frame, Plot, Scale};
use topos::{Shape, Tape, Tensor, cross_entropy, init};

use chart::loss_chart;
use corpus::{VOCABULARY_LEN, draw, from_token, load_names, shuffle, training_samples};

/// How many bigram pairs each training step feeds.
const BATCH_LEN: usize = 1024;

fn main() {
    let names = load_names();
    let mut samples = training_samples::<1>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!(
        "loaded {} names, {} bigram pairs",
        names.len(),
        samples.len()
    );

    let tape: Tape<f32> = Tape::new();
    let table = tape.parameter(init::normal(7, 0.01)(&Shape::new([
        VOCABULARY_LEN,
        VOCABULARY_LEN,
    ])));

    // Contexts and targets are one-hot selections fed per run; the
    // defaults only fix the batch shape.
    let contexts = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));

    let logits = table.gather(contexts);
    let loss = cross_entropy(logits, targets);

    let table_symbol = table.symbol();
    let contexts_symbol = contexts.symbol();
    let targets_symbol = targets.symbol();
    let loss_symbol = loss.symbol();
    let network = tape.into_network();
    let recorded_nodes = network.len();

    // A fresh model is roughly uniform over the vocabulary, so the
    // first printed loss should sit near `ln(27) ~ 3.30`; the bigram
    // limit on this corpus is about `2.45`.
    let learning_rate = Tensor::new([], [10.0]);
    let mut parameters = network.parameters();
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..1000 {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let batch = &samples[start..start + BATCH_LEN];
        let batch_contexts: Vec<usize> = batch.iter().map(|&(context, _)| context[0]).collect();
        let batch_targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();

        // Slice the run to the loss it reads.
        let run = network.entry([loss_symbol]).interpret(
            &parameters,
            [
                (
                    contexts_symbol,
                    Tensor::selection(batch_contexts, VOCABULARY_LEN, 1.0),
                ),
                (
                    targets_symbol,
                    Tensor::selection(batch_targets, VOCABULARY_LEN, 1.0),
                ),
            ],
        );
        let batch_loss = run.of(loss_symbol).scalar();
        losses.push(batch_loss);
        if step % 100 == 0 {
            println!("step {step:4}: minibatch loss = {batch_loss:.4}");
        }
        let gradients = run.backward(loss_symbol).parameters(&parameters);
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    assert_eq!(network.len(), recorded_nodes);
    println!("the tape held {recorded_nodes} nodes through every step");
    println!("{}", loss_chart("bigram training", &losses));

    // The trained logits exponentiate into transition probabilities
    // through the composite softmax: reopen the sealed network, record
    // one more expression, and carry the trained state across.
    let tape = network.into_tape();
    let probabilities = tape.resolve(table_symbol).softmax(1).symbol();
    let network = tape.into_network();
    let parameters = parameters.carried(&network);
    // Slice to the freshly recorded softmax: the training expression
    // does not re-run just to render the table.
    let run = network
        .entry([probabilities])
        .interpret(&parameters, std::iter::empty());
    let probabilities = run
        .of(probabilities)
        .as_slice()
        .expect("a computed softmax is contiguous")
        .to_vec();

    // The trained table as a picture — the differentiable mirror of the
    // classic makemore count matrix. Cell extents shift by half a cell
    // so each cell centers on its token's band index, and the frame
    // grows to give every one of the 27 rows its own terminal row.
    let span = (-0.5, VOCABULARY_LEN as f64 - 0.5);
    let letters = (0..VOCABULARY_LEN).map(|token| String::from(from_token(token)));
    let mut frame = Frame::detect();
    frame.height = VOCABULARY_LEN + 6;
    println!(
        "{}",
        Plot::new()
            .layer(Cells::matrix(VOCABULARY_LEN, &probabilities[..]).extents(span, span))
            .colorbar()
            .x_scale(Scale::bands(letters))
            .title("bigram transition probabilities")
            .x_label("next token")
            .y_label("current token")
            .render_best(&frame)
    );

    println!("sampled names:");
    let mut state: u64 = 7;
    for _ in 0..10 {
        let mut token = 0;
        let mut name = String::new();
        loop {
            let row = &probabilities[token * VOCABULARY_LEN..(token + 1) * VOCABULARY_LEN];
            token = draw(row, &mut state);
            if token == 0 {
                break;
            }
            name.push(from_token(token));
        }
        println!("  {name}");
    }
}

mod chart;
mod corpus;

use std::time::Instant;

use topos::{Shape, Tape, Tensor, Value, cross_entropy, init};

use chart::loss_chart;
use corpus::{VOCABULARY_LEN, draw, from_token, load_names, shuffle, training_samples};

/// How many characters of history the model sees before predicting the
/// next one.
const CONTEXT_LEN: usize = 3;

/// How many dimensions the character embedding space has.
const EMBED_DIM: usize = 10;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 100;

/// How many samples each training step feeds.
const BATCH_LEN: usize = 64;

/// The model's parameters as recorded proxies: what the `Mlp` facade
/// would hold for us, laid out by hand.
struct Model<'tape> {
    embeddings: Value<'tape, f32>,
    hidden_weights: Value<'tape, f32>,
    hidden_bias: Value<'tape, f32>,
    output_weights: Value<'tape, f32>,
    output_bias: Value<'tape, f32>,
}

impl<'tape> Model<'tape> {
    /// Allocates the parameters on `tape`: an embedding table, one
    /// tanh hidden layer, and an affine output layer, Xavier-scaled
    /// with zero biases.
    fn new(tape: &'tape Tape<f32>) -> Self {
        let mut weights = init::xavier(7);
        Self {
            embeddings: tape.parameter(init::normal(8, 1.0)(&Shape::new([
                VOCABULARY_LEN,
                EMBED_DIM,
            ]))),
            hidden_weights: tape
                .parameter(weights(&Shape::new([CONTEXT_LEN * EMBED_DIM, HIDDEN_LEN]))),
            hidden_bias: tape.parameter(weights(&Shape::new([HIDDEN_LEN]))),
            output_weights: tape.parameter(weights(&Shape::new([HIDDEN_LEN, VOCABULARY_LEN]))),
            output_bias: tape.parameter(weights(&Shape::new([VOCABULARY_LEN]))),
        }
    }

    /// Records the model's expression over `contexts` — a one-hot
    /// `[rows * CONTEXT_LEN, vocab]` selection — and returns the
    /// `[rows, vocab]` logits: embed, flatten the context window,
    /// squash, and score.
    fn express(&self, contexts: Value<'tape, f32>, rows: usize) -> Value<'tape, f32> {
        let embedded = self
            .embeddings
            .gather(contexts)
            .reshape([rows, CONTEXT_LEN * EMBED_DIM]);
        let product = embedded.matmul(self.hidden_weights);
        let hidden = (product + self.hidden_bias.broadcast_along_like(0, product)).tanh();
        let product = hidden.matmul(self.output_weights);
        product + self.output_bias.broadcast_along_like(0, product)
    }
}

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!("loaded {} names, {} samples", names.len(), samples.len());

    let tape = Tape::new();
    let model = Model::new(&tape);

    // The training expression, batch-shaped: contexts and targets are
    // one-hot selections fed per run, the defaults only fix the shapes.
    let contexts = tape.input(Tensor::selection(
        vec![0; BATCH_LEN * CONTEXT_LEN],
        VOCABULARY_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));
    let loss = cross_entropy(model.express(contexts, BATCH_LEN), targets);

    // The sampling twin: the same parameters expressed over a single
    // context row, with the composite softmax on top.
    let sample_context = tape.input(Tensor::selection(vec![0; CONTEXT_LEN], VOCABULARY_LEN, 1.0));
    let sample_probabilities = model.express(sample_context, 1).softmax(1);

    let (contexts, targets, loss, sample_context, sample_probabilities) = (
        contexts.symbol(),
        targets.symbol(),
        loss.symbol(),
        sample_context.symbol(),
        sample_probabilities.symbol(),
    );
    let recorded_nodes = tape.len();
    let network = tape.into_network();
    let mut parameters = network.parameters();

    // A fresh model is roughly uniform over the vocabulary, so the
    // first printed loss should sit near `ln(27) ~ 3.30`; the goal is
    // to push below the bigram limit of ~2.45.
    let fast = Tensor::new([], [0.1]);
    let slow = Tensor::new([], [0.01]);
    let mut window_loss = 0.0;
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..5000 {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let batch = &samples[start..start + BATCH_LEN];
        let batch_contexts: Vec<usize> = batch
            .iter()
            .flat_map(|(context, _)| context.iter().copied())
            .collect();
        let batch_targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();

        // Slice the run to the loss: the sampling twin on the same
        // tape is skipped during training.
        let run = network.entry([loss]).interpret(
            &parameters,
            [
                (
                    contexts,
                    Tensor::selection(batch_contexts, VOCABULARY_LEN, 1.0),
                ),
                (
                    targets,
                    Tensor::selection(batch_targets, VOCABULARY_LEN, 1.0),
                ),
            ],
        );
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        if step == 0 {
            println!(
                "step 0: minibatch loss = {batch_loss:.4} (a uniform model costs ln 27 ~ 3.30)"
            );
        }
        window_loss += batch_loss;
        if (step + 1) % 500 == 0 {
            println!(
                "steps {:4}..{:4}: mean minibatch loss = {:.4}",
                step + 1 - 500,
                step + 1,
                window_loss / 500.0
            );
            window_loss = 0.0;
        }
        let gradients = run.backward(loss).parameters(&parameters);
        let learning_rate = if step < 4000 { &fast } else { &slow };
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    assert_eq!(network.len(), recorded_nodes);
    println!("the tape held {recorded_nodes} nodes through every step");
    println!("{}", loss_chart("mlp training", &losses));

    println!("sampled names:");
    let mut state: u64 = 7;
    for _ in 0..10 {
        let mut window = [0usize; CONTEXT_LEN];
        let mut name = String::new();
        loop {
            let run = network.entry([sample_probabilities]).interpret(
                &parameters,
                [(
                    sample_context,
                    Tensor::selection(window.to_vec(), VOCABULARY_LEN, 1.0),
                )],
            );
            let row = run.of(sample_probabilities).to_vec();
            let token = draw(&row, &mut state);
            if token == 0 {
                break;
            }
            name.push(from_token(token));
            window.rotate_left(1);
            window[CONTEXT_LEN - 1] = token;
        }
        println!("  {name}");
    }
}

mod chart;
mod corpus;

use std::time::Instant;

use topos::{Activation, Mlp, Module, Shape, Tape, Tensor, cross_entropy, init};

use chart::loss_chart;
use corpus::{VOCABULARY_LEN, draw, from_token, load_names, shuffle, training_samples};

/// How many characters of history the model sees before predicting the
/// next one.
const CONTEXT_LEN: usize = 3;

/// How many dimensions the character embedding space has.
const EMBED_DIM: usize = 10;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 100;

/// How many samples each training step feeds.
const BATCH_LEN: usize = 64;

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!("loaded {} names, {} samples", names.len(), samples.len());

    let tape: Tape<f32> = Tape::new();

    // The embedding table stays a plain parameter; the facade covers
    // the dense layers only. The allocation order and seeds match
    // `makemore_mlp`, so both examples start from identical weights.
    let embeddings = tape.parameter(init::normal(8, 1.0)(&Shape::new([
        VOCABULARY_LEN,
        EMBED_DIM,
    ])));
    let mlp = Mlp::new(
        &tape,
        &[CONTEXT_LEN * EMBED_DIM, HIDDEN_LEN, VOCABULARY_LEN],
        Activation::Tanh,
        init::xavier(7),
    );

    // The training expression, batch-shaped: contexts and targets are
    // one-hot selections fed per run, the defaults only fix the shapes.
    let contexts = tape.input(Tensor::selection(
        vec![0; BATCH_LEN * CONTEXT_LEN],
        VOCABULARY_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));
    let embedded = embeddings
        .gather(contexts)
        .reshape([BATCH_LEN, CONTEXT_LEN * EMBED_DIM]);
    let loss = cross_entropy(mlp.express(embedded), targets);

    // The sampling twin: the same parameters expressed over a single
    // context row, with the composite softmax on top.
    let sample_context = tape.input(Tensor::selection(vec![0; CONTEXT_LEN], VOCABULARY_LEN, 1.0));
    let sample_embedded = embeddings
        .gather(sample_context)
        .reshape([1, CONTEXT_LEN * EMBED_DIM]);
    let sample_probabilities = mlp.express(sample_embedded).softmax(1);

    let (contexts, targets, loss, sample_context, sample_probabilities) = (
        contexts.symbol(),
        targets.symbol(),
        loss.symbol(),
        sample_context.symbol(),
        sample_probabilities.symbol(),
    );
    let recorded_nodes = tape.len();
    let network = tape.into_network();
    let mut parameters = network.parameters();

    // A fresh model is roughly uniform over the vocabulary, so the
    // first printed loss should sit near `ln(27) ~ 3.30`; the goal is
    // to push below the bigram limit of ~2.45.
    let fast = Tensor::new([], [0.1]);
    let slow = Tensor::new([], [0.01]);
    let mut window_loss = 0.0;
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..5000 {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let batch = &samples[start..start + BATCH_LEN];
        let batch_contexts: Vec<usize> = batch
            .iter()
            .flat_map(|(context, _)| context.iter().copied())
            .collect();
        let batch_targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();

        // Slice the run to the loss: the sampling twin on the same
        // tape is skipped during training.
        let run = network.entry([loss]).interpret(
            &parameters,
            [
                (
                    contexts,
                    Tensor::selection(batch_contexts, VOCABULARY_LEN, 1.0),
                ),
                (
                    targets,
                    Tensor::selection(batch_targets, VOCABULARY_LEN, 1.0),
                ),
            ],
        );
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        if step == 0 {
            println!(
                "step 0: minibatch loss = {batch_loss:.4} (a uniform model costs ln 27 ~ 3.30)"
            );
        }
        window_loss += batch_loss;
        if (step + 1) % 500 == 0 {
            println!(
                "steps {:4}..{:4}: mean minibatch loss = {:.4}",
                step + 1 - 500,
                step + 1,
                window_loss / 500.0
            );
            window_loss = 0.0;
        }
        let gradients = run.backward(loss).parameters(&parameters);
        let learning_rate = if step < 4000 { &fast } else { &slow };
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    assert_eq!(network.len(), recorded_nodes);
    println!("the tape held {recorded_nodes} nodes through every step");
    println!("{}", loss_chart("mlp (facade) training", &losses));

    println!("sampled names:");
    let mut state: u64 = 7;
    for _ in 0..10 {
        let mut window = [0usize; CONTEXT_LEN];
        let mut name = String::new();
        loop {
            let run = network.entry([sample_probabilities]).interpret(
                &parameters,
                [(
                    sample_context,
                    Tensor::selection(window.to_vec(), VOCABULARY_LEN, 1.0),
                )],
            );
            let row = run.of(sample_probabilities).to_vec();
            let token = draw(&row, &mut state);
            if token == 0 {
                break;
            }
            name.push(from_token(token));
            window.rotate_left(1);
            window[CONTEXT_LEN - 1] = token;
        }
        println!("  {name}");
    }
}

mod chart;
mod corpus;

use std::time::Instant;

use topos::{Shape, Tape, Tensor, Value, cross_entropy, init};

use chart::loss_chart;
use corpus::{VOCABULARY_LEN, draw, from_token, load_names, shuffle, training_samples};

/// How many characters of history the model sees before predicting the
/// next one.
const CONTEXT_LEN: usize = 3;

/// How many dimensions the character embedding space has.
const EMBED_DIM: usize = 10;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 100;

/// How many samples each training step feeds.
const BATCH_LEN: usize = 64;

/// The model's parameters as recorded proxies, laid out exactly as in
/// `makemore_mlp` so the two examples share their seeds.
struct Model<'tape> {
    embeddings: Value<'tape, f32>,
    hidden_weights: Value<'tape, f32>,
    hidden_bias: Value<'tape, f32>,
    output_weights: Value<'tape, f32>,
    output_bias: Value<'tape, f32>,
}

impl<'tape> Model<'tape> {
    /// Allocates the parameters on `tape`: an embedding table, one
    /// tanh hidden layer, and an affine output layer, Xavier-scaled
    /// with zero biases.
    fn new(tape: &'tape Tape<f32>) -> Self {
        let mut weights = init::xavier(7);
        Self {
            embeddings: tape.parameter(init::normal(8, 1.0)(&Shape::new([
                VOCABULARY_LEN,
                EMBED_DIM,
            ]))),
            hidden_weights: tape
                .parameter(weights(&Shape::new([CONTEXT_LEN * EMBED_DIM, HIDDEN_LEN]))),
            hidden_bias: tape.parameter(weights(&Shape::new([HIDDEN_LEN]))),
            output_weights: tape.parameter(weights(&Shape::new([HIDDEN_LEN, VOCABULARY_LEN]))),
            output_bias: tape.parameter(weights(&Shape::new([VOCABULARY_LEN]))),
        }
    }

    /// Returns the parameters in a fixed order, for pairing with their
    /// recorded gradients.
    fn parameters(&self) -> [Value<'tape, f32>; 5] {
        [
            self.embeddings,
            self.hidden_weights,
            self.hidden_bias,
            self.output_weights,
            self.output_bias,
        ]
    }

    /// Records the model's expression over `contexts` — a one-hot
    /// `[rows * CONTEXT_LEN, vocab]` selection — and returns the
    /// `[rows, vocab]` logits: embed, flatten the context window,
    /// squash, and score.
    fn express(&self, contexts: Value<'tape, f32>, rows: usize) -> Value<'tape, f32> {
        let embedded = self
            .embeddings
            .gather(contexts)
            .reshape([rows, CONTEXT_LEN * EMBED_DIM]);
        let product = embedded.matmul(self.hidden_weights);
        let hidden = (product + self.hidden_bias.broadcast_along_like(0, product)).tanh();
        let product = hidden.matmul(self.output_weights);
        product + self.output_bias.broadcast_along_like(0, product)
    }
}

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!("loaded {} names, {} samples", names.len(), samples.len());

    let tape = Tape::new();
    let model = Model::new(&tape);

    // The training expression, batch-shaped: contexts and targets are
    // one-hot selections fed per run, the defaults only fix the shapes.
    let contexts = tape.input(Tensor::selection(
        vec![0; BATCH_LEN * CONTEXT_LEN],
        VOCABULARY_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));
    let loss = cross_entropy(model.express(contexts, BATCH_LEN), targets);

    // The sampling twin: the same parameters expressed over a single
    // context row, with the composite softmax on top.
    let sample_context = tape.input(Tensor::selection(vec![0; CONTEXT_LEN], VOCABULARY_LEN, 1.0));
    let sample_probabilities = model.express(sample_context, 1).softmax(1);

    let (contexts, targets, loss, sample_context, sample_probabilities) = (
        contexts.symbol(),
        targets.symbol(),
        loss.symbol(),
        sample_context.symbol(),
        sample_probabilities.symbol(),
    );

    // The whole difference from `makemore_mlp` is these lines: the
    // chain rule is appended to the tape once, and one forward-only
    // plan serves loss and gradients together for every step.
    // The sampling twin is not among the targets, so the plan's
    // dead-node elimination drops it from every training step.
    let forward_nodes = tape.len();
    let parameter_symbols = model.parameters().map(|parameter| parameter.symbol());
    let adjoints = tape.differentiate(loss, parameter_symbols);
    println!(
        "recorded the chain rule: {} forward nodes + {} gradient nodes",
        forward_nodes,
        tape.len() - forward_nodes
    );
    let network = tape.into_network();
    let mut parameters = network.parameters();
    let plan = network.entry(adjoints.roots()).lower();

    // A fresh model is roughly uniform over the vocabulary, so the
    // first printed loss should sit near `ln(27) ~ 3.30`; matched
    // seeds make every printed number equal `makemore_mlp`'s.
    let fast = Tensor::new([], [0.1]);
    let slow = Tensor::new([], [0.01]);
    let mut window_loss = 0.0;
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..5000 {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let batch = &samples[start..start + BATCH_LEN];
        let batch_contexts: Vec<usize> = batch
            .iter()
            .flat_map(|(context, _)| context.iter().copied())
            .collect();
        let batch_targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();

        let run = plan.forward(
            &parameters,
            [
                (
                    contexts,
                    Tensor::selection(batch_contexts, VOCABULARY_LEN, 1.0),
                ),
                (
                    targets,
                    Tensor::selection(batch_targets, VOCABULARY_LEN, 1.0),
                ),
            ],
        );
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        if step == 0 {
            println!(
                "step 0: minibatch loss = {batch_loss:.4} (a uniform model costs ln 27 ~ 3.30)"
            );
        }
        window_loss += batch_loss;
        if (step + 1) % 500 == 0 {
            println!(
                "steps {:4}..{:4}: mean minibatch loss = {:.4}",
                step + 1 - 500,
                step + 1,
                window_loss / 500.0
            );
            window_loss = 0.0;
        }
        // The gradients were computed by the same plan run that
        // computed the loss; assembling the update direction is a
        // read, not a backward pass.
        let gradients = run.recorded_gradients(&adjoints);
        let learning_rate = if step < 4000 { &fast } else { &slow };
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s, one plan run per step, no backward pass",
        losses.len(),
        training.elapsed().as_secs_f64()
    );
    println!("{}", loss_chart("compiled mlp training", &losses));

    println!("sampled names:");
    let mut state: u64 = 7;
    for _ in 0..10 {
        let mut window = [0usize; CONTEXT_LEN];
        let mut name = String::new();
        loop {
            let run = network.entry([sample_probabilities]).interpret(
                &parameters,
                [(
                    sample_context,
                    Tensor::selection(window.to_vec(), VOCABULARY_LEN, 1.0),
                )],
            );
            let row = run.of(sample_probabilities).to_vec();
            let token = draw(&row, &mut state);
            if token == 0 {
                break;
            }
            name.push(from_token(token));
            window.rotate_left(1);
            window[CONTEXT_LEN - 1] = token;
        }
        println!("  {name}");
    }
}

// The shared corpus module also serves the sampling examples; this act
// never samples, so its `draw`/`from_token` stay unused here.
#[allow(dead_code)]
mod corpus;

use std::time::Instant;

use malevich::stat::Window;
use malevich::{Frame, Line, Plot, Rule};
use topos::{Adam, Optimizer, Sgd, Shape, Tape, Tensor, Value, cross_entropy, init};

use corpus::{VOCABULARY_LEN, load_names, shuffle, training_samples};

/// How many characters of history the model sees before predicting the
/// next one.
const CONTEXT_LEN: usize = 3;

/// How many dimensions the character embedding space has.
const EMBED_DIM: usize = 10;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 100;

/// How many samples each training step feeds.
const BATCH_LEN: usize = 64;

/// How many training steps each optimizer runs.
const STEPS: usize = 5000;

/// The corpus's bigram limit, the line both curves aim below.
const BIGRAM_LIMIT: f64 = 2.45;

/// The model's parameters as recorded proxies, laid out exactly as in
/// `makemore_mlp_compiled` so the runs share their seeds.
struct Model<'tape> {
    embeddings: Value<'tape, f32>,
    hidden_weights: Value<'tape, f32>,
    hidden_bias: Value<'tape, f32>,
    output_weights: Value<'tape, f32>,
    output_bias: Value<'tape, f32>,
}

impl<'tape> Model<'tape> {
    /// Allocates the parameters on `tape`: an embedding table, one
    /// tanh hidden layer, and an affine output layer, Xavier-scaled
    /// with zero biases.
    fn new(tape: &'tape Tape<f32>) -> Self {
        let mut weights = init::xavier(7);
        Self {
            embeddings: tape.parameter(init::normal(8, 1.0)(&Shape::new([
                VOCABULARY_LEN,
                EMBED_DIM,
            ]))),
            hidden_weights: tape
                .parameter(weights(&Shape::new([CONTEXT_LEN * EMBED_DIM, HIDDEN_LEN]))),
            hidden_bias: tape.parameter(weights(&Shape::new([HIDDEN_LEN]))),
            output_weights: tape.parameter(weights(&Shape::new([HIDDEN_LEN, VOCABULARY_LEN]))),
            output_bias: tape.parameter(weights(&Shape::new([VOCABULARY_LEN]))),
        }
    }

    /// Returns the parameters in a fixed order, for pairing with their
    /// recorded gradients.
    fn parameters(&self) -> [Value<'tape, f32>; 5] {
        [
            self.embeddings,
            self.hidden_weights,
            self.hidden_bias,
            self.output_weights,
            self.output_bias,
        ]
    }

    /// Records the model's expression over `contexts` — a one-hot
    /// `[rows * CONTEXT_LEN, vocab]` selection — and returns the
    /// `[rows, vocab]` logits: embed, flatten the context window,
    /// squash, and score.
    fn express(&self, contexts: Value<'tape, f32>, rows: usize) -> Value<'tape, f32> {
        let embedded = self
            .embeddings
            .gather(contexts)
            .reshape([rows, CONTEXT_LEN * EMBED_DIM]);
        let product = embedded.matmul(self.hidden_weights);
        let hidden = (product + self.hidden_bias.broadcast_along_like(0, product)).tanh();
        let product = hidden.matmul(self.output_weights);
        product + self.output_bias.broadcast_along_like(0, product)
    }
}

/// Trains a fresh model for `STEPS` under `optimizer`, with the
/// per-step learning rate from `learning_rate`, and returns the
/// per-step losses and the wall-clock seconds. Every run rebuilds the
/// same seeds, so trajectories differ only by the optimizer.
fn train(
    samples: &[([usize; CONTEXT_LEN], usize)],
    recorded: bool,
    optimizer: &mut dyn Optimizer<f32>,
    learning_rate: impl Fn(usize) -> Tensor<f32>,
) -> (Vec<f32>, f64) {
    let tape = Tape::new();
    let model = Model::new(&tape);
    let contexts = tape.input(Tensor::selection(
        vec![0; BATCH_LEN * CONTEXT_LEN],
        VOCABULARY_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));
    let loss = cross_entropy(model.express(contexts, BATCH_LEN), targets);

    let (contexts, targets, loss) = (contexts.symbol(), targets.symbol(), loss.symbol());
    let parameter_symbols = model.parameters().map(|parameter| parameter.symbol());

    // The recorded route: the chain rule on the tape, one forward-only
    // plan per run. The engine route skips both and differentiates
    // each interpreter run procedurally, projecting each complete
    // field onto the parameter slots — bit-identical losses, and the
    // same parameter-aligned gradients either way.
    let adjoints = recorded.then(|| tape.differentiate(loss, parameter_symbols));
    let network = tape.into_network();
    let plan = adjoints
        .as_ref()
        .map(|adjoints| network.entry(adjoints.roots()).lower());

    let mut parameters = network.parameters();
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..STEPS {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let batch = &samples[start..start + BATCH_LEN];
        let batch_contexts: Vec<usize> = batch
            .iter()
            .flat_map(|(context, _)| context.iter().copied())
            .collect();
        let batch_targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();
        let feeds = [
            (
                contexts,
                Tensor::selection(batch_contexts, VOCABULARY_LEN, 1.0),
            ),
            (
                targets,
                Tensor::selection(batch_targets, VOCABULARY_LEN, 1.0),
            ),
        ];

        let (batch_loss, gradients) = if let (Some(plan), Some(adjoints)) = (&plan, &adjoints) {
            let run = plan.forward(&parameters, feeds);
            let gradients = run.recorded_gradients(adjoints);
            (run.of(loss).scalar(), gradients)
        } else {
            let run = network.forward(&parameters, feeds);
            (
                run.of(loss).scalar(),
                run.backward(loss).parameters(&parameters),
            )
        };
        losses.push(batch_loss);

        parameters = optimizer.step(&parameters, &gradients, &learning_rate(step));
    }
    (losses, training.elapsed().as_secs_f64())
}

/// Renders both rolling-mean curves on one chart, with the bigram
/// limit as the line to beat.
fn comparison_chart(sgd: &[f32], adam: &[f32]) -> String {
    let window_len = (STEPS / 20).max(2);
    let smooth = |losses: &[f32]| {
        let losses: Vec<f64> = losses.iter().copied().map(f64::from).collect();
        Window::new(window_len).mean(&losses)
    };
    let sgd = smooth(sgd);
    let adam = smooth(adam);
    Plot::new()
        .layer(Line::y(&sgd[..]).label("sgd"))
        .layer(Line::y(&adam[..]).label("adam"))
        .layer(Rule::h(BIGRAM_LIMIT).label("bigram limit"))
        .title("sgd vs adam, rolling mean")
        .x_label("step")
        .y_label("loss")
        .render_best(&Frame::detect())
}

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!("loaded {} names, {} samples", names.len(), samples.len());

    let recorded = match std::env::var("TOPOS_GRADIENTS").as_deref() {
        Ok("engine") => false,
        Ok("recorded") | Err(_) => true,
        Ok(other) => panic!("unknown TOPOS_GRADIENTS {other:?}; use recorded or engine"),
    };
    println!(
        "gradients: {}",
        if recorded {
            "recorded (differentiate + one forward-only plan)"
        } else {
            "engine (interpreter backward)"
        }
    );

    // SGD keeps `makemore_mlp_compiled`'s schedule, so its losses
    // reproduce that example bit for bit; Adam runs at one flat rate —
    // the adaptive moments replace the hand-tuned decay.
    let fast = Tensor::new([], [0.1_f32]);
    let slow = Tensor::new([], [0.01_f32]);
    let (sgd_losses, sgd_seconds) = train(&samples, recorded, &mut Sgd, |step| {
        if step < 4000 {
            fast.clone()
        } else {
            slow.clone()
        }
    });

    let mut adam = Adam::new(
        Tensor::new([], [0.9_f32]),
        Tensor::new([], [0.999_f32]),
        Tensor::new([], [1e-8_f32]),
    );
    let flat = Tensor::new([], [0.005_f32]);
    let (adam_losses, adam_seconds) = train(&samples, recorded, &mut adam, |_| flat.clone());

    for (label, losses, seconds) in [
        ("sgd", &sgd_losses, sgd_seconds),
        ("adam", &adam_losses, adam_seconds),
    ] {
        let window: f32 = losses[STEPS - 500..].iter().sum::<f32>() / 500.0;
        println!(
            "{label:5} {STEPS} steps in {seconds:.3}s ({:.2} ms/step), last-500 mean loss {window:.4}",
            seconds * 1000.0 / STEPS as f64
        );
    }
    println!("{}", comparison_chart(&sgd_losses, &adam_losses));
}

mod chart;
mod corpus;

use std::time::Instant;

use rayon::prelude::*;

use topos::{Activation, Mlp, Module, Parameters, Shape, Tape, Tensor, cross_entropy, init};

use chart::loss_chart;
use corpus::{VOCABULARY_LEN, draw, from_token, load_names, shuffle, training_samples};

/// How many characters of history the model sees before predicting the
/// next one.
const CONTEXT_LEN: usize = 3;

/// How many dimensions the character embedding space has.
const EMBED_DIM: usize = 10;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 100;

/// How many concurrent shards each training step fans out.
///
/// The count is fixed rather than detected — the shard partition
/// decides the arithmetic, so a detected count would make runs differ
/// across machines. Rayon still schedules the fixed shards over every
/// core it has. The count trades per-run fixed cost (favoring fewer,
/// larger shards) against load balancing over heterogeneous cores
/// (favoring more, smaller ones); eight-by-eight measured fastest on an
/// eight-performance-core machine, where sixteen-by-four raised
/// utilization but paid more in run overhead than it recovered.
const SHARD_COUNT: usize = 8;

/// How many samples each shard carries; the effective batch is
/// `SHARD_COUNT * SHARD_LEN`, matching the serial examples' 64.
const SHARD_LEN: usize = 8;

/// Sums shard gradients as a pairwise tree whose shape depends only on
/// the shard count: the reduction runs its pairs concurrently and
/// finishes in logarithmic depth, while the tree — not the scheduler —
/// decides the order of additions, keeping the result deterministic.
fn tree_sum(mut layer: Vec<Parameters<f32>>) -> Parameters<f32> {
    while layer.len() > 1 {
        layer = layer
            .par_chunks(2)
            .map(|pair| match pair {
                [left, right] => left + right,
                [single] => single.clone(),
                _ => unreachable!("chunks of two hold one or two tables"),
            })
            .collect();
    }
    layer.into_iter().next().expect("at least one shard ran")
}

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!("loaded {} names, {} samples", names.len(), samples.len());

    let tape: Tape<f32> = Tape::new();

    // The same model as the serial examples, recorded at shard shape:
    // the batch size is baked into the graph, so the parallel plan is
    // one shard-shaped expression run once per shard, not a wider one.
    let embeddings = tape.parameter(init::normal(8, 1.0)(&Shape::new([
        VOCABULARY_LEN,
        EMBED_DIM,
    ])));
    let mlp = Mlp::new(
        &tape,
        &[CONTEXT_LEN * EMBED_DIM, HIDDEN_LEN, VOCABULARY_LEN],
        Activation::Tanh,
        init::xavier(7),
    );

    let contexts = tape.input(Tensor::selection(
        vec![0; SHARD_LEN * CONTEXT_LEN],
        VOCABULARY_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; SHARD_LEN], VOCABULARY_LEN, 1.0));
    let embedded = embeddings
        .gather(contexts)
        .reshape([SHARD_LEN, CONTEXT_LEN * EMBED_DIM]);
    let loss = cross_entropy(mlp.express(embedded), targets);

    // The sampling twin: the same parameters expressed over a single
    // context row, with the composite softmax on top.
    let sample_context = tape.input(Tensor::selection(vec![0; CONTEXT_LEN], VOCABULARY_LEN, 1.0));
    let sample_embedded = embeddings
        .gather(sample_context)
        .reshape([1, CONTEXT_LEN * EMBED_DIM]);
    let sample_probabilities = mlp.express(sample_embedded).softmax(1);

    let contexts_symbol = contexts.symbol();
    let targets_symbol = targets.symbol();
    let loss_symbol = loss.symbol();
    let sample_context_symbol = sample_context.symbol();
    let sample_probabilities_symbol = sample_probabilities.symbol();
    let network = tape.into_network();
    let recorded_nodes = network.len();

    let batch_len = SHARD_COUNT * SHARD_LEN;
    let shard_inverse = Tensor::new([], [1.0 / SHARD_COUNT as f32]);
    let fast = Tensor::new([], [0.1]);
    let slow = Tensor::new([], [0.01]);
    let mut parameters = network.parameters();
    let mut window_loss = 0.0;
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..5000 {
        let start = (step * batch_len) % (samples.len() - batch_len);
        let batch = &samples[start..start + batch_len];

        // Fan out: one immutable forward and backward run per shard,
        // all reading the same network and parameter state.
        let shard_results: Vec<(f32, Parameters<f32>)> = (0..SHARD_COUNT)
            .into_par_iter()
            .map(|shard| {
                let rows = &batch[shard * SHARD_LEN..(shard + 1) * SHARD_LEN];
                let shard_contexts: Vec<usize> = rows
                    .iter()
                    .flat_map(|(context, _)| context.iter().copied())
                    .collect();
                let shard_targets: Vec<usize> = rows.iter().map(|&(_, next)| next).collect();

                // Slice the run to the loss: the sampling twin on the
                // same tape is skipped during training.
                let run = network.entry([loss_symbol]).interpret(
                    &parameters,
                    [
                        (
                            contexts_symbol,
                            Tensor::selection(shard_contexts, VOCABULARY_LEN, 1.0),
                        ),
                        (
                            targets_symbol,
                            Tensor::selection(shard_targets, VOCABULARY_LEN, 1.0),
                        ),
                    ],
                );
                let shard_loss = run.of(loss_symbol).scalar();
                (
                    shard_loss,
                    run.backward(loss_symbol).parameters(&parameters),
                )
            })
            .collect();

        let (shard_losses, shard_gradients): (Vec<f32>, Vec<Parameters<f32>>) =
            shard_results.into_iter().unzip();
        let batch_loss = shard_losses.iter().sum::<f32>() / SHARD_COUNT as f32;
        let gradients = tree_sum(shard_gradients)
            .map(|gradient| gradient.clone() * shard_inverse.broadcast_like(gradient));

        losses.push(batch_loss);
        if step == 0 {
            println!(
                "step 0: minibatch loss = {batch_loss:.4} (a uniform model costs ln 27 ~ 3.30)"
            );
        }
        window_loss += batch_loss;
        if (step + 1) % 500 == 0 {
            println!(
                "steps {:4}..{:4}: mean minibatch loss = {:.4}",
                step + 1 - 500,
                step + 1,
                window_loss / 500.0
            );
            window_loss = 0.0;
        }
        let learning_rate = if step < 4000 { &fast } else { &slow };
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }
    println!(
        "trained {} steps on {SHARD_COUNT} shards in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    assert_eq!(network.len(), recorded_nodes);
    println!("the tape held {recorded_nodes} nodes through every step");
    println!("{}", loss_chart("mlp (data parallel) training", &losses));

    println!("sampled names:");
    let mut state: u64 = 7;
    for _ in 0..10 {
        let mut window = [0usize; CONTEXT_LEN];
        let mut name = String::new();
        loop {
            let run = network.entry([sample_probabilities_symbol]).interpret(
                &parameters,
                [(
                    sample_context_symbol,
                    Tensor::selection(window.to_vec(), VOCABULARY_LEN, 1.0),
                )],
            );
            let row = run.of(sample_probabilities_symbol).to_vec();
            let token = draw(&row, &mut state);
            if token == 0 {
                break;
            }
            name.push(from_token(token));
            window.rotate_left(1);
            window[CONTEXT_LEN - 1] = token;
        }
        println!("  {name}");
    }
}

mod chart;
mod corpus;

use std::time::Instant;

use topos::{BatchNorm, Shape, Tape, Tensor, Value, cross_entropy, init};

use chart::loss_chart;
use corpus::{VOCABULARY_LEN, draw, from_token, load_names, shuffle, training_samples};

/// How many characters of history the model sees before predicting the
/// next one.
const CONTEXT_LEN: usize = 3;

/// How many dimensions the character embedding space has.
const EMBED_DIM: usize = 10;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 100;

/// How many samples each training step feeds.
const BATCH_LEN: usize = 64;

/// The exponential-moving-average momentum of the running statistics:
/// each step keeps `1 - MOMENTUM` of the estimate and takes `MOMENTUM`
/// of the batch statistic.
const MOMENTUM: f32 = 0.01;

/// The model's parameters as recorded proxies: the part-2 layout with
/// the hidden bias replaced by a batch-normalization stage.
struct Model<'tape> {
    embeddings: Value<'tape, f32>,
    hidden_weights: Value<'tape, f32>,
    norm: BatchNorm<f32>,
    output_weights: Value<'tape, f32>,
    output_bias: Value<'tape, f32>,
}

impl<'tape> Model<'tape> {
    /// Allocates the parameters on `tape`: the embedding table, the
    /// bias-free hidden layer, the norm's scale and shift (ones and
    /// zeros, the standard start), and the affine output.
    fn new(tape: &'tape Tape<f32>) -> Self {
        let mut weights = init::xavier(7);
        Self {
            embeddings: tape.parameter(init::normal(8, 1.0)(&Shape::new([
                VOCABULARY_LEN,
                EMBED_DIM,
            ]))),
            hidden_weights: tape
                .parameter(weights(&Shape::new([CONTEXT_LEN * EMBED_DIM, HIDDEN_LEN]))),
            norm: BatchNorm::new(
                tape,
                Tensor::filled([HIDDEN_LEN], 1.0),
                Tensor::filled([HIDDEN_LEN], 0.0),
                Tensor::filled([], 1e-5),
            ),
            output_weights: tape.parameter(weights(&Shape::new([HIDDEN_LEN, VOCABULARY_LEN]))),
            output_bias: tape.parameter(weights(&Shape::new([VOCABULARY_LEN]))),
        }
    }

    /// Records the shared head of both expressions: embed, flatten,
    /// and the bias-free hidden preactivation.
    fn preactivation(&self, contexts: Value<'tape, f32>, rows: usize) -> Value<'tape, f32> {
        self.embeddings
            .gather(contexts)
            .reshape([rows, CONTEXT_LEN * EMBED_DIM])
            .matmul(self.hidden_weights)
    }

    /// Records the shared tail: squash the normalized preactivation
    /// and score.
    fn logits(&self, normalized: Value<'tape, f32>) -> Value<'tape, f32> {
        let hidden = normalized.tanh();
        let product = hidden.matmul(self.output_weights);
        product + self.output_bias.broadcast_along_like(0, product)
    }
}

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!("loaded {} names, {} samples", names.len(), samples.len());

    let tape = Tape::new();
    let model = Model::new(&tape);

    // The training expression: the norm standardizes each feature by
    // the minibatch's own statistics, and hands them back for the
    // running estimates.
    let contexts = tape.input(Tensor::selection(
        vec![0; BATCH_LEN * CONTEXT_LEN],
        VOCABULARY_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));
    let normalization = model.norm.express(model.preactivation(contexts, BATCH_LEN));
    let loss = cross_entropy(model.logits(normalization.output), targets);

    // The sampling twin normalizes by running estimates fed per run:
    // the norm's inference mode, one recorded expression for every
    // state of the estimates.
    let sample_context = tape.input(Tensor::selection(vec![0; CONTEXT_LEN], VOCABULARY_LEN, 1.0));
    let running_mean = tape.input(Tensor::filled([HIDDEN_LEN], 0.0));
    let running_variance = tape.input(Tensor::filled([HIDDEN_LEN], 1.0));
    let sample_normalized = model.norm.express_with(
        model.preactivation(sample_context, 1),
        running_mean,
        running_variance,
    );
    let sample_probabilities = model.logits(sample_normalized).softmax(1);

    let (contexts, targets, loss) = (contexts.symbol(), targets.symbol(), loss.symbol());
    let (mean, variance) = (normalization.mean.symbol(), normalization.variance.symbol());
    let (sample_context, running_mean, running_variance, sample_probabilities) = (
        sample_context.symbol(),
        running_mean.symbol(),
        running_variance.symbol(),
        sample_probabilities.symbol(),
    );
    let recorded_nodes = tape.len();
    let network = tape.into_network();
    let mut parameters = network.parameters();

    // Entry once: the training plan keeps the batch statistics
    // readable — the keep-set naming exactly what the loop reads —
    // and the sampling plan is forward-only.
    let training_plan = network
        .entry([loss])
        .observe([mean, variance])
        .backward()
        .lower();
    let sampling_plan = network.entry([sample_probabilities]).lower();

    // The running estimates: loop-owned payloads, never engine state.
    let mut mean_estimate = Tensor::filled([HIDDEN_LEN], 0.0_f32);
    let mut variance_estimate = Tensor::filled([HIDDEN_LEN], 1.0_f32);
    let keep = Tensor::filled([HIDDEN_LEN], 1.0 - MOMENTUM);
    let take = Tensor::filled([HIDDEN_LEN], MOMENTUM);

    let fast = Tensor::new([], [0.1]);
    let slow = Tensor::new([], [0.01]);
    let mut window_loss = 0.0;
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..5000 {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let batch = &samples[start..start + BATCH_LEN];
        let batch_contexts: Vec<usize> = batch
            .iter()
            .flat_map(|(context, _)| context.iter().copied())
            .collect();
        let batch_targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();

        let run = training_plan.forward(
            &parameters,
            [
                (
                    contexts,
                    Tensor::selection(batch_contexts, VOCABULARY_LEN, 1.0),
                ),
                (
                    targets,
                    Tensor::selection(batch_targets, VOCABULARY_LEN, 1.0),
                ),
            ],
        );
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        if step == 0 {
            println!(
                "step 0: minibatch loss = {batch_loss:.4} (a uniform model costs ln 27 ~ 3.30)"
            );
        }
        window_loss += batch_loss;
        if (step + 1) % 500 == 0 {
            println!(
                "steps {:4}..{:4}: mean minibatch loss = {:.4}",
                step + 1 - 500,
                step + 1,
                window_loss / 500.0
            );
            window_loss = 0.0;
        }

        // The running estimates fold in this batch's statistics, read
        // through the keep-set: payload arithmetic in loop land.
        let batch_mean = run.of(mean).clone();
        let batch_variance = run.of(variance).clone();
        mean_estimate = mean_estimate * keep.clone() + batch_mean * take.clone();
        variance_estimate = variance_estimate * keep.clone() + batch_variance * take.clone();

        let gradients = run.backward(loss).parameters(&parameters);
        let learning_rate = if step < 4000 { &fast } else { &slow };
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    assert_eq!(network.len(), recorded_nodes);
    println!("the tape held {recorded_nodes} nodes through every step");
    println!("{}", loss_chart("mlp + batchnorm training", &losses));

    println!("sampled names (running statistics fed per draw):");
    let mut state: u64 = 7;
    for _ in 0..10 {
        let mut window = [0usize; CONTEXT_LEN];
        let mut name = String::new();
        loop {
            let run = sampling_plan.forward(
                &parameters,
                [
                    (
                        sample_context,
                        Tensor::selection(window.to_vec(), VOCABULARY_LEN, 1.0),
                    ),
                    (running_mean, mean_estimate.clone()),
                    (running_variance, variance_estimate.clone()),
                ],
            );
            let row = run.of(sample_probabilities).to_vec();
            let token = draw(&row, &mut state);
            if token == 0 {
                break;
            }
            name.push(from_token(token));
            window.rotate_left(1);
            window[CONTEXT_LEN - 1] = token;
        }
        println!("  {name}");
    }
}

mod chart;
#[allow(dead_code)]
mod corpus;

use std::time::Instant;

use malevich::{Color, Frame, Plot, Text};
use topos::{Activation, Mlp, Module, Shape, Tape, Tensor, cross_entropy, init};

use chart::loss_chart;
use corpus::{VOCABULARY_LEN, from_token, load_names, shuffle, training_samples};

/// How many characters of history the model sees before predicting the
/// next one.
const CONTEXT_LEN: usize = 3;

/// How many dimensions the character embedding space has: two, so the
/// space is the page.
const EMBED_DIM: usize = 2;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 100;

/// How many samples each training step feeds.
const BATCH_LEN: usize = 64;

/// Spreads `points` so no two share a cell of a `columns` by `rows`
/// grid laid over their bounding box: a point whose cell is taken moves
/// to the nearest free cell within a small ring. A later text mark
/// would otherwise overwrite an earlier one where the space clusters —
/// which is exactly where the interesting letters are — so the chart
/// accepts a slight positional lie to keep every letter visible. A flat
/// axis widens to a unit span so a degenerate cloud still spreads.
fn spread(points: &mut [(f64, f64)], columns: usize, rows: usize) {
    let mut x_min = f64::INFINITY;
    let mut x_max = f64::NEG_INFINITY;
    let mut y_min = f64::INFINITY;
    let mut y_max = f64::NEG_INFINITY;
    for &(x, y) in points.iter() {
        x_min = x_min.min(x);
        x_max = x_max.max(x);
        y_min = y_min.min(y);
        y_max = y_max.max(y);
    }
    let x_span = if x_max > x_min { x_max - x_min } else { 1.0 };
    let y_span = if y_max > y_min { y_max - y_min } else { 1.0 };

    let mut taken = vec![vec![false; columns]; rows];
    for point in points.iter_mut() {
        let column = ((point.0 - x_min) / x_span * (columns - 1) as f64).round() as isize;
        let row = ((point.1 - y_min) / y_span * (rows - 1) as f64).round() as isize;
        'placed: for radius in 0..=3_isize {
            for row_offset in -radius..=radius {
                for column_offset in -radius..=radius {
                    if row_offset.abs().max(column_offset.abs()) != radius {
                        continue;
                    }
                    let target_row = row + row_offset;
                    let target_column = column + column_offset;
                    if target_row < 0
                        || target_row >= rows as isize
                        || target_column < 0
                        || target_column >= columns as isize
                    {
                        continue;
                    }
                    if !taken[target_row as usize][target_column as usize] {
                        taken[target_row as usize][target_column as usize] = true;
                        point.0 = x_min + target_column as f64 / (columns - 1) as f64 * x_span;
                        point.1 = y_min + target_row as f64 / (rows - 1) as f64 * y_span;
                        break 'placed;
                    }
                }
            }
        }
    }
}

/// Renders the `[vocab, 2]` embedding `table` as a terminal map, one
/// text mark per token: vowels highlighted, the padding token dimmed,
/// consonants in the default foreground.
fn embedding_chart(title: &str, table: &Tensor<f32>) -> String {
    // The map wants more rows than the default third-of-terminal strip:
    // every letter is one cell, so vertical resolution is legibility.
    let mut frame = Frame::detect();
    frame.height = frame.height.max(24);

    let elements = table.to_vec();
    let mut points: Vec<(f64, f64)> = (0..VOCABULARY_LEN)
        .map(|token| {
            (
                f64::from(elements[token * EMBED_DIM]),
                f64::from(elements[token * EMBED_DIM + 1]),
            )
        })
        .collect();
    // The spread grid stands in for the plot area's cells, held a bit
    // coarser than the frame since malevich widens the domains to its
    // ticks and spends some cells on axis furniture.
    spread(
        &mut points,
        frame.width.saturating_sub(12).max(20),
        frame.height.saturating_sub(6).max(10),
    );

    // The domains pad one spread cell beyond the bounding box: a text
    // mark anchored exactly on the domain's right or bottom edge is
    // clipped away (malevich 1.11.1), and edge letters read better with
    // a margin anyway.
    let mut x_bounds = (f64::INFINITY, f64::NEG_INFINITY);
    let mut y_bounds = (f64::INFINITY, f64::NEG_INFINITY);
    for &(x, y) in &points {
        x_bounds = (x_bounds.0.min(x), x_bounds.1.max(x));
        y_bounds = (y_bounds.0.min(y), y_bounds.1.max(y));
    }
    let x_margin = (x_bounds.1 - x_bounds.0).max(1.0) / 40.0;
    let y_margin = (y_bounds.1 - y_bounds.0).max(1.0) / 16.0;

    let mut plot = Plot::new()
        .title(title)
        .x_domain(x_bounds.0 - x_margin, x_bounds.1 + x_margin)
        .y_domain(y_bounds.0 - y_margin, y_bounds.1 + y_margin);
    for (token, &(x, y)) in points.iter().enumerate() {
        let letter = from_token(token);
        let color = match letter {
            'a' | 'e' | 'i' | 'o' | 'u' => Color::BrightCyan,
            '.' => Color::BrightBlack,
            _ => Color::Default,
        };
        plot = plot.layer(Text::at(x, y, String::from(letter)).color(color));
    }
    plot.render_best(&frame)
}

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!("loaded {} names, {} samples", names.len(), samples.len());

    let tape: Tape<f32> = Tape::new();
    let embeddings = tape.parameter(init::normal(8, 1.0)(&Shape::new([
        VOCABULARY_LEN,
        EMBED_DIM,
    ])));
    let mlp = Mlp::new(
        &tape,
        &[CONTEXT_LEN * EMBED_DIM, HIDDEN_LEN, VOCABULARY_LEN],
        Activation::Tanh,
        init::xavier(7),
    );

    let contexts = tape.input(Tensor::selection(
        vec![0; BATCH_LEN * CONTEXT_LEN],
        VOCABULARY_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));
    let embedded = embeddings
        .gather(contexts)
        .reshape([BATCH_LEN, CONTEXT_LEN * EMBED_DIM]);
    let loss = cross_entropy(mlp.express(embedded), targets);

    println!(
        "{}",
        embedding_chart(
            "embedding space before training (the seeded blob)",
            &embeddings.payload().unwrap(),
        )
    );

    let (embeddings, contexts, targets, loss) = (
        embeddings.symbol(),
        contexts.symbol(),
        targets.symbol(),
        loss.symbol(),
    );
    let network = tape.into_network();
    let mut parameters = network.parameters();

    // The two-dimensional bottleneck lands near 2.4 where ten
    // dimensions reach 2.25: the price of a plottable space.
    let fast = Tensor::new([], [0.1]);
    let slow = Tensor::new([], [0.01]);
    let mut window_loss = 0.0;
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..5000 {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let batch = &samples[start..start + BATCH_LEN];
        let batch_contexts: Vec<usize> = batch
            .iter()
            .flat_map(|(context, _)| context.iter().copied())
            .collect();
        let batch_targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();

        // Slice the run to the loss it reads.
        let run = network.entry([loss]).interpret(
            &parameters,
            [
                (
                    contexts,
                    Tensor::selection(batch_contexts, VOCABULARY_LEN, 1.0),
                ),
                (
                    targets,
                    Tensor::selection(batch_targets, VOCABULARY_LEN, 1.0),
                ),
            ],
        );
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        window_loss += batch_loss;
        if (step + 1) % 1000 == 0 {
            println!(
                "steps {:4}..{:4}: mean minibatch loss = {:.4}",
                step + 1 - 1000,
                step + 1,
                window_loss / 1000.0
            );
            window_loss = 0.0;
        }
        let gradients = run.backward(loss).parameters(&parameters);
        let learning_rate = if step < 4000 { &fast } else { &slow };
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );
    println!("{}", loss_chart("embedding map training", &losses));

    let table = parameters.of(embeddings);
    println!(
        "{}",
        embedding_chart("embedding space after training (vowels highlighted)", table)
    );
}

mod chart;
mod corpus;

use std::time::Instant;

use topos::{
    Dropout, Embedding, Module, RmsNorm, Shape, Tape, Tensor, Value, concat, cross_entropy, init,
    scaled_dot_product,
};

use chart::loss_chart;
use corpus::{VOCABULARY_LEN, draw, from_token, load_names, shuffle, training_samples};

/// How many characters of history the model attends over before
/// predicting the next one.
const CONTEXT_LEN: usize = 8;

/// How many dimensions the residual stream has.
const EMBED_DIM: usize = 32;

/// How many attention heads split the stream.
const HEAD_COUNT: usize = 2;

/// How many dimensions each head reads and writes.
const HEAD_DIM: usize = EMBED_DIM / HEAD_COUNT;

/// How many neurons the feed-forward hidden layer has.
const HIDDEN_LEN: usize = 4 * EMBED_DIM;

/// How many samples each training step packs into the attention row.
const BATCH_LEN: usize = 32;

/// How many token rows the packed training batch holds.
const PACKED_LEN: usize = BATCH_LEN * CONTEXT_LEN;

/// The dropout keep probability: masks carry `1 / KEEP` or zero, and
/// only the training expression is ever fed — the sampling twin runs
/// on the all-ones identity default.
const KEEP: f64 = 0.9;

/// One attention head's projections, each `[EMBED_DIM, HEAD_DIM]`.
struct Head<'tape> {
    query: Value<'tape, f32>,
    key: Value<'tape, f32>,
    value: Value<'tape, f32>,
}

/// The model's parameters as recorded proxies: embeddings, one
/// pre-norm attention block, and the language-model head.
struct Model<'tape> {
    embeddings: Embedding<f32>,
    positions: Embedding<f32>,
    heads: Vec<Head<'tape>>,
    projection: Value<'tape, f32>,
    attention_norm: RmsNorm<f32>,
    hidden_weights: Value<'tape, f32>,
    output_weights: Value<'tape, f32>,
    hidden_norm: RmsNorm<f32>,
    final_norm: RmsNorm<f32>,
    logit_weights: Value<'tape, f32>,
    logit_bias: Value<'tape, f32>,
    scale: Value<'tape, f32>,
}

impl<'tape> Model<'tape> {
    /// Allocates the parameters on `tape`: embedding tables for
    /// characters and positions, the heads' projections, the block's
    /// two norms and feed-forward, the final norm, and the affine
    /// logit head. The attention scale `1 / sqrt(HEAD_DIM)` rides
    /// along as a single-value leaf.
    fn new(tape: &'tape Tape<f32>) -> Self {
        let mut weights = init::xavier(7);
        let ones = Tensor::filled([EMBED_DIM], 1.0);
        let epsilon = Tensor::filled([], 1e-5);
        Self {
            embeddings: Embedding::new(
                tape,
                init::normal(8, 1.0)(&Shape::new([VOCABULARY_LEN, EMBED_DIM])),
            ),
            positions: Embedding::new(
                tape,
                init::normal(9, 1.0)(&Shape::new([CONTEXT_LEN, EMBED_DIM])),
            ),
            heads: (0..HEAD_COUNT)
                .map(|_| Head {
                    query: tape.parameter(weights(&Shape::new([EMBED_DIM, HEAD_DIM]))),
                    key: tape.parameter(weights(&Shape::new([EMBED_DIM, HEAD_DIM]))),
                    value: tape.parameter(weights(&Shape::new([EMBED_DIM, HEAD_DIM]))),
                })
                .collect(),
            projection: tape.parameter(weights(&Shape::new([EMBED_DIM, EMBED_DIM]))),
            attention_norm: RmsNorm::new(tape, ones.clone(), epsilon.clone()),
            hidden_weights: tape.parameter(weights(&Shape::new([EMBED_DIM, HIDDEN_LEN]))),
            output_weights: tape.parameter(weights(&Shape::new([HIDDEN_LEN, EMBED_DIM]))),
            hidden_norm: RmsNorm::new(tape, ones.clone(), epsilon.clone()),
            final_norm: RmsNorm::new(tape, ones, epsilon),
            logit_weights: tape.parameter(weights(&Shape::new([EMBED_DIM, VOCABULARY_LEN]))),
            logit_bias: tape.parameter(weights(&Shape::new([VOCABULARY_LEN]))),
            scale: tape.leaf(Tensor::filled([], 1.0 / (HEAD_DIM as f32).sqrt())),
        }
    }

    /// Records the block over one packed token row and returns the
    /// normalized residual stream: embeddings in, attention and
    /// feed-forward added on, final norm out.
    ///
    /// `tokens` and `positions` are one-hot selections over the packed
    /// row; `mask` is the additive `0 / -inf` block-causal leaf shaped
    /// `[rows, rows]`.
    fn states(
        &self,
        _tape: &'tape Tape<f32>,
        tokens: Value<'tape, f32>,
        positions: Value<'tape, f32>,
        mask: Value<'tape, f32>,
        dropouts: &[Dropout<f32>; 2],
    ) -> Value<'tape, f32> {
        let stream = self.embeddings.express(tokens) + self.positions.express(positions);

        // Pre-norm attention: every head attends over the same
        // normalized stream, and `concat` joins the head outputs along
        // the feature axis.
        let normalized = self.attention_norm.express(stream);
        let heads: Vec<Value<'tape, f32>> = self
            .heads
            .iter()
            .map(|head| {
                let query = normalized.matmul(head.query);
                let key = normalized.matmul(head.key);
                let value = normalized.matmul(head.value);
                scaled_dot_product(query, key, value, mask, self.scale)
            })
            .collect();
        let stream = stream + dropouts[0].express(concat(&heads, 1).matmul(self.projection));

        // Pre-norm feed-forward onto the residual stream.
        let normalized = self.hidden_norm.express(stream);
        let stream = stream
            + dropouts[1].express(
                normalized
                    .matmul(self.hidden_weights)
                    .relu()
                    .matmul(self.output_weights),
            );

        self.final_norm.express(stream)
    }

    /// Records the logits of the rows picked by the one-hot
    /// `extraction` — the last context position of each packed sample.
    fn logits(
        &self,
        states: Value<'tape, f32>,
        extraction: Value<'tape, f32>,
    ) -> Value<'tape, f32> {
        let product = states.gather(extraction).matmul(self.logit_weights);
        product + self.logit_bias.broadcast_along_like(0, product)
    }
}

/// Returns the additive attention mask for `samples` packed windows of
/// `CONTEXT_LEN` tokens: zero where the key's row is in the same
/// sample and not after the query's row, negative infinity elsewhere.
fn block_causal_mask(samples: usize) -> Tensor<f32> {
    let rows = samples * CONTEXT_LEN;
    let mut elements = Vec::with_capacity(rows * rows);
    for query in 0..rows {
        for key in 0..rows {
            let same_sample = query / CONTEXT_LEN == key / CONTEXT_LEN;
            let causal = key % CONTEXT_LEN <= query % CONTEXT_LEN;
            elements.push(if same_sample && causal {
                0.0
            } else {
                f32::NEG_INFINITY
            });
        }
    }
    Tensor::new([rows, rows], elements)
}

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!("loaded {} names, {} samples", names.len(), samples.len());

    let tape = Tape::new();
    let model = Model::new(&tape);

    // The training expression over the packed batch: per-sample
    // prediction rows come back through the fixed one-hot extraction.
    let tokens = tape.input(Tensor::selection(vec![0; PACKED_LEN], VOCABULARY_LEN, 1.0));
    let positions = tape.leaf(Tensor::selection(
        (0..PACKED_LEN)
            .map(|row| row % CONTEXT_LEN)
            .collect::<Vec<_>>(),
        CONTEXT_LEN,
        1.0,
    ));
    let mask = tape.leaf(block_causal_mask(BATCH_LEN));
    let extraction = tape.leaf(Tensor::selection(
        (0..BATCH_LEN)
            .map(|sample| sample * CONTEXT_LEN + CONTEXT_LEN - 1)
            .collect::<Vec<_>>(),
        PACKED_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));
    // Dropout on both residual writes; the mask inputs default to
    // ones, so only the fed training steps ever drop anything.
    let dropouts = [
        Dropout::new(&tape, [PACKED_LEN, EMBED_DIM]),
        Dropout::new(&tape, [PACKED_LEN, EMBED_DIM]),
    ];
    let states = model.states(&tape, tokens, positions, mask, &dropouts);
    let loss = cross_entropy(model.logits(states, extraction), targets);

    // The sampling twin is the same expression over one window.
    let sample_tokens = tape.input(Tensor::selection(vec![0; CONTEXT_LEN], VOCABULARY_LEN, 1.0));
    let sample_positions = tape.leaf(Tensor::selection(
        (0..CONTEXT_LEN).collect::<Vec<_>>(),
        CONTEXT_LEN,
        1.0,
    ));
    let sample_mask = tape.leaf(block_causal_mask(1));
    let sample_extraction = tape.leaf(Tensor::selection(vec![CONTEXT_LEN - 1], CONTEXT_LEN, 1.0));
    // The twin's dropouts are never fed: the identity default is the
    // inference mode, with no flag and no second formula.
    let sample_dropouts = [
        Dropout::new(&tape, [CONTEXT_LEN, EMBED_DIM]),
        Dropout::new(&tape, [CONTEXT_LEN, EMBED_DIM]),
    ];
    let sample_states = model.states(
        &tape,
        sample_tokens,
        sample_positions,
        sample_mask,
        &sample_dropouts,
    );
    let sample_probabilities = model.logits(sample_states, sample_extraction).softmax(1);

    let (tokens, targets, loss, sample_tokens, sample_probabilities) = (
        tokens.symbol(),
        targets.symbol(),
        loss.symbol(),
        sample_tokens.symbol(),
        sample_probabilities.symbol(),
    );
    let recorded_nodes = tape.len();
    let network = tape.into_network();
    let mut parameters = network.parameters();

    // Entry once: training keeps only the loss, sampling is
    // forward-only.
    let training_plan = network.entry([loss]).backward().lower();
    let sampling_plan = network.entry([sample_probabilities]).lower();

    // The seeded mask factory: two draws per step, one per residual
    // write, deterministic in `(seed, step)` so runs replay bitwise.
    let mut dropout_masks = init::dropout::<f32>(13, KEEP);
    let mask_shape = Shape::new([PACKED_LEN, EMBED_DIM]);

    let fast = Tensor::new([], [0.1]);
    let slow = Tensor::new([], [0.01]);
    let mut window_loss = 0.0;
    let mut losses = Vec::new();
    let training = Instant::now();
    for step in 0..5000 {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let batch = &samples[start..start + BATCH_LEN];
        let batch_tokens: Vec<usize> = batch
            .iter()
            .flat_map(|(context, _)| context.iter().copied())
            .collect();
        let batch_targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();

        let run = training_plan.forward(
            &parameters,
            [
                (tokens, Tensor::selection(batch_tokens, VOCABULARY_LEN, 1.0)),
                (
                    targets,
                    Tensor::selection(batch_targets, VOCABULARY_LEN, 1.0),
                ),
                (dropouts[0].mask(), dropout_masks(&mask_shape)),
                (dropouts[1].mask(), dropout_masks(&mask_shape)),
            ],
        );
        let batch_loss = run.of(loss).scalar();
        losses.push(batch_loss);
        if step == 0 {
            println!(
                "step 0: minibatch loss = {batch_loss:.4} (a uniform model costs ln 27 ~ 3.30)"
            );
        }
        window_loss += batch_loss;
        if (step + 1) % 500 == 0 {
            println!(
                "steps {:4}..{:4}: mean minibatch loss = {:.4}",
                step + 1 - 500,
                step + 1,
                window_loss / 500.0
            );
            window_loss = 0.0;
        }

        let gradients = run.backward(loss).parameters(&parameters);
        let learning_rate = if step < 4000 { &fast } else { &slow };
        parameters = parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * learning_rate.broadcast_like(gradient)
        });
    }

    println!(
        "trained {} steps in {:.3}s",
        losses.len(),
        training.elapsed().as_secs_f64()
    );

    assert_eq!(network.len(), recorded_nodes);
    println!("the tape held {recorded_nodes} nodes through every step");
    println!("{}", loss_chart("transformer training", &losses));

    println!("sampled names:");
    let mut state: u64 = 7;
    for _ in 0..10 {
        let mut window = [0usize; CONTEXT_LEN];
        let mut name = String::new();
        loop {
            let run = sampling_plan.forward(
                &parameters,
                [(
                    sample_tokens,
                    Tensor::selection(window.to_vec(), VOCABULARY_LEN, 1.0),
                )],
            );
            let row = run.of(sample_probabilities).to_vec();
            let token = draw(&row, &mut state);
            if token == 0 {
                break;
            }
            name.push(from_token(token));
            window.rotate_left(1);
            window[CONTEXT_LEN - 1] = token;
        }
        println!("  {name}");
    }
}

// The shared corpus module also serves the sampling examples; this
// act never samples, so its `draw`/`from_token` stay unused here.
#[allow(dead_code)]
mod corpus;

use std::io::{Read, Write};
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
use std::time::Instant;

use topos::{Shape, Tape, Tensor, Value, cross_entropy, init};

use corpus::{VOCABULARY_LEN, load_names, shuffle, training_samples};

/// How many characters of history the model sees before predicting the
/// next one.
const CONTEXT_LEN: usize = 3;

/// How many dimensions the character embedding space has.
const EMBED_DIM: usize = 10;

/// How many neurons the tanh hidden layer has.
const HIDDEN_LEN: usize = 100;

/// How many samples each training step feeds.
const BATCH_LEN: usize = 64;

/// How many training steps each route runs.
const STEPS: usize = 5000;

/// The initial parameter payloads, minted identically to
/// `makemore_mlp_compiled`'s `Model::new` (same factories, same call
/// order, same seeds), so both routes and that example share their
/// trajectory. Order is the recording order and therefore the
/// emitted module's leading-argument order.
fn initial_parameters() -> [Tensor<f32>; 5] {
    let mut weights = init::xavier(7);
    [
        init::normal(8, 1.0)(&Shape::new([VOCABULARY_LEN, EMBED_DIM])),
        weights(&Shape::new([CONTEXT_LEN * EMBED_DIM, HIDDEN_LEN])),
        weights(&Shape::new([HIDDEN_LEN])),
        weights(&Shape::new([HIDDEN_LEN, VOCABULARY_LEN])),
        weights(&Shape::new([VOCABULARY_LEN])),
    ]
}

/// Records the model's expression over `contexts` and returns the
/// `[rows, vocab]` logits: embed, flatten the context window, squash,
/// and score. `parameters` arrive in `initial_parameters` order.
fn express<'tape>(
    parameters: &[Value<'tape, f32>; 5],
    contexts: Value<'tape, f32>,
    rows: usize,
) -> Value<'tape, f32> {
    let [
        embeddings,
        hidden_weights,
        hidden_bias,
        output_weights,
        output_bias,
    ] = *parameters;
    let embedded = embeddings
        .gather(contexts)
        .reshape([rows, CONTEXT_LEN * EMBED_DIM]);
    let product = embedded.matmul(hidden_weights);
    let hidden = (product + hidden_bias.broadcast_along_like(0, product)).tanh();
    let product = hidden.matmul(output_weights);
    product + output_bias.broadcast_along_like(0, product)
}

/// One-hot payloads for a batch slice, contexts then targets.
fn batch_payloads(batch: &[([usize; CONTEXT_LEN], usize)]) -> (Tensor<f32>, Tensor<f32>) {
    let contexts: Vec<usize> = batch
        .iter()
        .flat_map(|(context, _)| context.iter().copied())
        .collect();
    let targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();
    (
        Tensor::selection(contexts, VOCABULARY_LEN, 1.0),
        Tensor::selection(targets, VOCABULARY_LEN, 1.0),
    )
}

/// The resident XLA server over `tools/serve-stablehlo-xla.py`:
/// compile once, then one request per training step. Every argument
/// is dynamic — parameters change every step.
struct XlaServer {
    child: Child,
    requests: ChildStdin,
    responses: ChildStdout,
    response_len: usize,
}

impl XlaServer {
    fn new(module: &str, dynamic_shapes: &[&[usize]], response_len: usize) -> Self {
        let directory = std::env::temp_dir().join("topos-makemore-e2");
        std::fs::create_dir_all(&directory).expect("the staging directory creates");
        let module_path = directory.join("joint-step.mlir");
        let static_path = directory.join("static.bin");
        let manifest_path = directory.join("manifest.json");
        std::fs::write(&module_path, module).expect("the module writes");
        std::fs::write(&static_path, []).expect("the empty static file writes");
        let shapes: Vec<String> = dynamic_shapes
            .iter()
            .map(|shape| {
                let extents: Vec<String> = shape.iter().map(usize::to_string).collect();
                format!("[{}]", extents.join(", "))
            })
            .collect();
        std::fs::write(
            &manifest_path,
            format!("{{\"dynamic\": [{}]}}", shapes.join(", ")),
        )
        .expect("the manifest writes");

        let python = std::env::var("TOPOS_XLA_PYTHON").unwrap_or_else(|_| "python3".to_string());
        let mut command: Vec<String> = python.split_whitespace().map(str::to_string).collect();
        command.push("tools/serve-stablehlo-xla.py".to_string());
        let mut child = Command::new(&command[0])
            .args(&command[1..])
            .arg(&module_path)
            .arg(&static_path)
            .arg(&manifest_path)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .spawn()
            .expect("the serving process starts; is `jax` installed for it?");
        let requests = child.stdin.take().expect("the server's input pipes");
        let responses = child.stdout.take().expect("the server's output pipes");
        Self {
            child,
            requests,
            responses,
            response_len,
        }
    }

    /// Executes one joint step: sends every argument's elements in
    /// order, reads back the results in the pinned order.
    fn step(&mut self, arguments: impl Iterator<Item = f32>) -> Vec<f32> {
        let request: Vec<u8> = arguments.flat_map(f32::to_le_bytes).collect();
        self.requests
            .write_all(&request)
            .expect("the request writes");
        self.requests.flush().expect("the request flushes");
        let mut response = vec![0u8; 4 * self.response_len];
        self.responses
            .read_exact(&mut response)
            .expect("the server answers; see its standard error");
        response
            .chunks_exact(4)
            .map(|chunk| f32::from_le_bytes(chunk.try_into().expect("four bytes")))
            .collect()
    }
}

impl Drop for XlaServer {
    fn drop(&mut self) {
        let _ = self.child.kill();
        let _ = self.child.wait();
    }
}

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);
    println!("loaded {} names, {} samples", names.len(), samples.len());

    // Record the expression and the chain rule once. The network's
    // only job here is to be the spec: after `emit_stablehlo`, the
    // emitted route never touches it again.
    let tape = Tape::new();
    let initial = initial_parameters();
    let parameters = initial.clone().map(|payload| tape.parameter(payload));
    let contexts = tape.input(Tensor::selection(
        vec![0; BATCH_LEN * CONTEXT_LEN],
        VOCABULARY_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));
    let loss = cross_entropy(express(&parameters, contexts, BATCH_LEN), targets);

    let (contexts, targets, loss) = (contexts.symbol(), targets.symbol(), loss.symbol());
    let parameter_symbols = parameters.map(|parameter| parameter.symbol());
    // Result order is declared: emission returns the request's roots
    // in request order, so the gradients emit in parameter order with
    // no aliasing ceremony.
    let adjoints = tape.differentiate(loss, parameter_symbols);
    let network = tape.into_network();
    let plan = network.entry(adjoints.roots()).lower();
    let module = plan.emit_stablehlo().expect("the joint step emits");
    println!(
        "emitted the joint step: {} nodes, {} bytes of StableHLO",
        plan.len(),
        module.len()
    );

    let fast = Tensor::new([], [0.1_f32]);
    let slow = Tensor::new([], [0.01_f32]);
    let learning_rate = |step: usize| if step < 4000 { &fast } else { &slow };

    // Route one, the oracle: the in-crate recorded plan, bit-identical
    // to `makemore_mlp_compiled`.
    let mut oracle_parameters = network.parameters();
    let mut oracle_losses = Vec::new();
    let oracle_clock = Instant::now();
    for step in 0..STEPS {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let (batch_contexts, batch_targets) = batch_payloads(&samples[start..start + BATCH_LEN]);
        let run = plan.forward(
            &oracle_parameters,
            [(contexts, batch_contexts), (targets, batch_targets)],
        );
        oracle_losses.push(run.of(loss).scalar());
        let gradients = run.recorded_gradients(&adjoints);
        let rate = learning_rate(step);
        oracle_parameters = oracle_parameters.step(&gradients, |parameter, gradient| {
            parameter.clone() - gradient.clone() * rate.broadcast_like(gradient)
        });
    }
    let oracle_seconds = oracle_clock.elapsed().as_secs_f64();

    // Route two, abroad: the same module through XLA, parameters held
    // host-side as plain payloads and updated with payload arithmetic.
    let shapes: Vec<Vec<usize>> = initial
        .iter()
        .map(|tensor| tensor.shape().axes().to_vec())
        .chain([
            vec![BATCH_LEN * CONTEXT_LEN, VOCABULARY_LEN],
            vec![BATCH_LEN, VOCABULARY_LEN],
        ])
        .collect();
    let dynamic_shapes: Vec<&[usize]> = shapes.iter().map(Vec::as_slice).collect();
    let gradient_volumes: Vec<usize> = initial
        .iter()
        .map(|tensor| tensor.shape().volume())
        .collect();
    let response_len = 1 + gradient_volumes.iter().sum::<usize>();
    let mut server = XlaServer::new(&module, &dynamic_shapes, response_len);

    let mut hosted: Vec<Tensor<f32>> = initial.to_vec();
    let mut emitted_losses = Vec::new();
    let emitted_clock = Instant::now();
    for step in 0..STEPS {
        let start = (step * BATCH_LEN) % (samples.len() - BATCH_LEN);
        let (batch_contexts, batch_targets) = batch_payloads(&samples[start..start + BATCH_LEN]);
        let response = server.step(
            hosted
                .iter()
                .flat_map(|tensor| tensor.to_vec())
                .chain(batch_contexts.to_vec())
                .chain(batch_targets.to_vec()),
        );
        emitted_losses.push(response[0]);
        let mut at = 1;
        let rate = learning_rate(step);
        for (parameter, &volume) in hosted.iter_mut().zip(&gradient_volumes) {
            let gradient = Tensor::new(
                parameter.shape().clone(),
                response[at..at + volume].to_vec(),
            );
            at += volume;
            *parameter = parameter.clone() - gradient.clone() * rate.broadcast_like(&gradient);
        }
    }
    let emitted_seconds = emitted_clock.elapsed().as_secs_f64();

    // Tier 1, per step: the first loss is one forward on identical
    // payloads — XLA may reassociate, never diverge.
    let first_drift = (f64::from(oracle_losses[0]) - f64::from(emitted_losses[0])).abs()
        / f64::from(oracle_losses[0]);
    println!(
        "step 0: oracle {:.6}, emitted {:.6} (relative drift {first_drift:.2e})",
        oracle_losses[0], emitted_losses[0]
    );
    assert!(
        first_drift < 1e-3,
        "the emitted first step left the oracle envelope"
    );

    // Tier 2, trajectory: windowed means over the whole schedule;
    // per-step drift compounds legitimately, so the page shows both.
    println!("           oracle   emitted");
    for window in 0..STEPS / 500 {
        let range = window * 500..(window + 1) * 500;
        let mean = |losses: &[f32]| losses[range.clone()].iter().sum::<f32>() / 500.0;
        println!(
            "steps {:4}..{:4}: {:.4}   {:.4}",
            range.start,
            range.end,
            mean(&oracle_losses),
            mean(&emitted_losses)
        );
    }
    println!(
        "oracle  (in-crate plan): {STEPS} steps in {oracle_seconds:.3}s ({:.2} ms/step)",
        oracle_seconds * 1000.0 / STEPS as f64
    );
    println!(
        "emitted (XLA joint step): {STEPS} steps in {emitted_seconds:.3}s ({:.2} ms/step)",
        emitted_seconds * 1000.0 / STEPS as f64
    );
}

// The shared corpus module also serves the sampling examples; this
// harness never samples, so most of it stays unused here.
#[allow(dead_code)]
mod corpus;

use std::io::{Read, Write};
use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio};
use std::time::Instant;

use topos::{
    Module, RmsNorm, Shape, Symbol, Tape, Tensor, Value, concat, cross_entropy, init, stack,
};

use corpus::{VOCABULARY_LEN, load_names, shuffle, training_samples};

/// How many characters of history the model attends over.
const CONTEXT_LEN: usize = 8;

/// How many dimensions the residual stream has.
const EMBED_DIM: usize = 32;

/// How many attention heads split the stream.
const HEAD_COUNT: usize = 2;

/// How many dimensions each head reads and writes.
const HEAD_DIM: usize = EMBED_DIM / HEAD_COUNT;

/// How many neurons the feed-forward hidden layer has.
const HIDDEN_LEN: usize = 4 * EMBED_DIM;

/// How many samples each training step packs into the attention row.
const BATCH_LEN: usize = 32;

/// How many token rows the packed training batch holds.
const PACKED_LEN: usize = BATCH_LEN * CONTEXT_LEN;

/// How many timed requests each module serves, after warmup.
const TIMED_STEPS: usize = 300;

/// How many warmup requests precede the timing.
const WARMUP_STEPS: usize = 20;

/// The parameter payloads in recording order — the emitted module's
/// leading arguments. Identical to the transformer act's seeds.
fn parameter_payloads() -> Vec<Tensor<f32>> {
    let mut weights = init::xavier(7);
    let ones = Tensor::filled([EMBED_DIM], 1.0);
    let mut payloads = vec![
        init::normal(8, 1.0)(&Shape::new([VOCABULARY_LEN, EMBED_DIM])),
        init::normal(9, 1.0)(&Shape::new([CONTEXT_LEN, EMBED_DIM])),
    ];
    for _ in 0..HEAD_COUNT {
        payloads.push(weights(&Shape::new([EMBED_DIM, HEAD_DIM])));
        payloads.push(weights(&Shape::new([EMBED_DIM, HEAD_DIM])));
        payloads.push(weights(&Shape::new([EMBED_DIM, HEAD_DIM])));
    }
    payloads.push(weights(&Shape::new([EMBED_DIM, EMBED_DIM])));
    payloads.push(ones.clone());
    payloads.push(weights(&Shape::new([EMBED_DIM, HIDDEN_LEN])));
    payloads.push(weights(&Shape::new([HIDDEN_LEN, EMBED_DIM])));
    payloads.push(ones.clone());
    payloads.push(ones);
    payloads.push(weights(&Shape::new([EMBED_DIM, VOCABULARY_LEN])));
    payloads.push(weights(&Shape::new([VOCABULARY_LEN])));
    payloads
}

/// Returns the additive block-causal mask for the packed batch.
fn block_causal_mask(samples: usize) -> Tensor<f32> {
    let rows = samples * CONTEXT_LEN;
    let mut elements = Vec::with_capacity(rows * rows);
    for query in 0..rows {
        for key in 0..rows {
            let same_sample = query / CONTEXT_LEN == key / CONTEXT_LEN;
            let causal = key % CONTEXT_LEN <= query % CONTEXT_LEN;
            elements.push(if same_sample && causal {
                0.0
            } else {
                f32::NEG_INFINITY
            });
        }
    }
    Tensor::new([rows, rows], elements)
}

/// One emitted formulation of the training step: the module text,
/// the tape statistics, and the oracle's first-step loss.
struct Formulation {
    module: String,
    forward_nodes: usize,
    total_nodes: usize,
    oracle_loss: f32,
}

/// Records the transformer training step — head loop or batched
/// attention per `batched` — differentiates it, pins the result
/// order, compiles, and emits. Parameters are handed in so both
/// formulations share their payloads exactly.
fn recorded(
    payloads: &[Tensor<f32>],
    feeds: &(Tensor<f32>, Tensor<f32>),
    batched: bool,
) -> Formulation {
    let tape = Tape::new();
    let mut supply = payloads.iter().cloned();
    let mut next = move || supply.next().expect("payloads cover the model");

    let embeddings = tape.parameter(next());
    let positions_table = tape.parameter(next());
    let heads: Vec<[Value<'_, f32>; 3]> = (0..HEAD_COUNT)
        .map(|_| {
            [
                tape.parameter(next()),
                tape.parameter(next()),
                tape.parameter(next()),
            ]
        })
        .collect();
    let projection = tape.parameter(next());
    let attention_norm = RmsNorm::new(&tape, next(), Tensor::filled([], 1e-5));
    let hidden_weights = tape.parameter(next());
    let output_weights = tape.parameter(next());
    let hidden_norm = RmsNorm::new(&tape, next(), Tensor::filled([], 1e-5));
    let final_norm = RmsNorm::new(&tape, next(), Tensor::filled([], 1e-5));
    let logit_weights = tape.parameter(next());
    let logit_bias = tape.parameter(next());
    let scale = tape.leaf(Tensor::filled([], 1.0 / (HEAD_DIM as f32).sqrt()));

    let tokens = tape.input(Tensor::selection(vec![0; PACKED_LEN], VOCABULARY_LEN, 1.0));
    let positions = tape.leaf(Tensor::selection(
        (0..PACKED_LEN)
            .map(|row| row % CONTEXT_LEN)
            .collect::<Vec<_>>(),
        CONTEXT_LEN,
        1.0,
    ));
    let mask = tape.leaf(block_causal_mask(BATCH_LEN));
    let extraction = tape.leaf(Tensor::selection(
        (0..BATCH_LEN)
            .map(|sample| sample * CONTEXT_LEN + CONTEXT_LEN - 1)
            .collect::<Vec<_>>(),
        PACKED_LEN,
        1.0,
    ));
    let targets = tape.input(Tensor::selection(vec![0; BATCH_LEN], VOCABULARY_LEN, 1.0));

    let stream = embeddings.gather(tokens) + positions_table.gather(positions);
    let normalized = attention_norm.express(stream);
    let attended = if batched {
        let project = |slot: usize| {
            let slices: Vec<_> = heads
                .iter()
                .map(|head| normalized.matmul(head[slot]))
                .collect();
            stack(&slices, 0)
        };
        let scores = project(0).matmul(project(1).permute([0, 2, 1]));
        let scaled = scores * scale.broadcast_like(scores);
        let weights = (scaled + mask.broadcast_along_like(0, scaled)).softmax(2);
        let context = weights.matmul(project(2));
        context.permute([1, 0, 2]).reshape([PACKED_LEN, EMBED_DIM])
    } else {
        let outputs: Vec<_> = heads
            .iter()
            .map(|head| {
                let scores = normalized
                    .matmul(head[0])
                    .matmul(normalized.matmul(head[1]).transpose());
                let scaled = scores * scale.broadcast_like(scores);
                let weights = (scaled + mask).softmax(1);
                weights.matmul(normalized.matmul(head[2]))
            })
            .collect();
        concat(&outputs, 1)
    };
    let stream = stream + attended.matmul(projection);
    let normalized = hidden_norm.express(stream);
    let stream = stream
        + normalized
            .matmul(hidden_weights)
            .relu()
            .matmul(output_weights);
    let states = final_norm.express(stream);
    let product = states.gather(extraction).matmul(logit_weights);
    let loss = cross_entropy(
        product + logit_bias.broadcast_along_like(0, product),
        targets,
    );

    let forward_nodes = tape.len();
    let wrt: Vec<Symbol> = std::iter::once(embeddings)
        .chain([positions_table])
        .chain(heads.iter().flatten().copied())
        .chain([
            projection,
            hidden_weights,
            output_weights,
            logit_weights,
            logit_bias,
        ])
        .map(|value| value.symbol())
        .collect();
    // Result order is declared: emission returns the request's roots
    // in request order, no aliasing needed.
    let adjoints = tape.differentiate(loss, wrt);
    let (tokens, targets, loss) = (tokens.symbol(), targets.symbol(), loss.symbol());
    let network = tape.into_network();
    let parameters = network.parameters();
    let plan = network.entry(adjoints.roots()).lower();
    let run = plan.forward(
        &parameters,
        [(tokens, feeds.0.clone()), (targets, feeds.1.clone())],
    );
    Formulation {
        module: plan.emit_stablehlo().expect("the joint step emits"),
        forward_nodes,
        total_nodes: network.len(),
        oracle_loss: run.of(loss).scalar(),
    }
}

/// Serves one module and times repeated identical requests.
fn timed(module: &str, name: &str, request: &[u8], response_len: usize) -> (f32, f64) {
    let directory = std::env::temp_dir().join("topos-attention-grading");
    std::fs::create_dir_all(&directory).expect("the staging directory creates");
    let module_path = directory.join(format!("{name}.mlir"));
    let static_path = directory.join("static.bin");
    let manifest_path = directory.join("manifest.json");
    std::fs::write(&module_path, module).expect("the module writes");
    std::fs::write(&static_path, []).expect("the empty static file writes");
    let mut shapes: Vec<String> = parameter_payloads()
        .iter()
        .map(|payload| {
            let extents: Vec<String> = payload
                .shape()
                .axes()
                .iter()
                .map(usize::to_string)
                .collect();
            format!("[{}]", extents.join(", "))
        })
        .collect();
    shapes.push(format!("[{PACKED_LEN}, {VOCABULARY_LEN}]"));
    shapes.push(format!("[{BATCH_LEN}, {VOCABULARY_LEN}]"));
    std::fs::write(
        &manifest_path,
        format!("{{\"dynamic\": [{}]}}", shapes.join(", ")),
    )
    .expect("the manifest writes");

    let python = std::env::var("TOPOS_XLA_PYTHON").unwrap_or_else(|_| "python3".to_string());
    let mut command: Vec<String> = python.split_whitespace().map(str::to_string).collect();
    command.push("tools/serve-stablehlo-xla.py".to_string());
    let mut child: Child = Command::new(&command[0])
        .args(&command[1..])
        .arg(&module_path)
        .arg(&static_path)
        .arg(&manifest_path)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .spawn()
        .expect("the serving process starts; is `jax` installed for it?");
    let mut requests: ChildStdin = child.stdin.take().expect("the server's input pipes");
    let mut responses: ChildStdout = child.stdout.take().expect("the server's output pipes");

    let mut step = |request: &[u8]| -> Vec<f32> {
        requests.write_all(request).expect("the request writes");
        requests.flush().expect("the request flushes");
        let mut response = vec![0u8; 4 * response_len];
        responses
            .read_exact(&mut response)
            .expect("the server answers; see its standard error");
        response
            .chunks_exact(4)
            .map(|chunk| f32::from_le_bytes(chunk.try_into().expect("four bytes")))
            .collect()
    };

    let first = step(request)[0];
    for _ in 0..WARMUP_STEPS {
        step(request);
    }
    let clock = Instant::now();
    for _ in 0..TIMED_STEPS {
        step(request);
    }
    let seconds = clock.elapsed().as_secs_f64();
    let _ = child.kill();
    let _ = child.wait();
    (first, seconds)
}

fn main() {
    let names = load_names();
    let mut samples = training_samples::<CONTEXT_LEN>(&names);
    let mut shuffle_state: u64 = 9;
    shuffle(&mut samples, &mut shuffle_state);

    // One fixed batch: identical feeds for the oracle and every
    // server request, so the two modules answer the same question.
    let batch = &samples[..BATCH_LEN];
    let tokens: Vec<usize> = batch
        .iter()
        .flat_map(|(context, _)| context.iter().copied())
        .collect();
    let targets: Vec<usize> = batch.iter().map(|&(_, next)| next).collect();
    let feeds = (
        Tensor::selection(tokens, VOCABULARY_LEN, 1.0_f32),
        Tensor::selection(targets, VOCABULARY_LEN, 1.0_f32),
    );

    let payloads = parameter_payloads();
    let request: Vec<u8> = payloads
        .iter()
        .flat_map(|payload| payload.to_vec())
        .chain(feeds.0.to_vec())
        .chain(feeds.1.to_vec())
        .flat_map(f32::to_le_bytes)
        .collect();

    for (name, batched) in [("head-loop", false), ("batched", true)] {
        let formulation = recorded(&payloads, &feeds, batched);
        let response_len = 1 + payloads
            .iter()
            .enumerate()
            .filter(|&(index, _)| !is_norm_gain(index))
            .map(|(_, payload)| payload.shape().volume())
            .sum::<usize>();
        let dots = formulation.module.matches("dot_general").count();
        let (first, seconds) = timed(&formulation.module, name, &request, response_len);
        let drift = (f64::from(formulation.oracle_loss) - f64::from(first)).abs()
            / f64::from(formulation.oracle_loss);
        // A red row on a flaky target is a finding, not a crash:
        // `TOPOS_ENVELOPE=report` keeps timing while naming the
        // deviation, the manifest posture for jax-metal.
        if std::env::var("TOPOS_ENVELOPE").as_deref() == Ok("report") {
            if drift >= 1e-3 {
                println!("{name}: RED ROW — left the oracle envelope, drift {drift:.1e}");
            }
        } else {
            assert!(drift < 1e-3, "{name} left the oracle envelope: {drift:e}");
        }
        println!(
            "{name:9} tape {} -> {} nodes, {} module lines, {dots} dot_general; \
             first loss {first:.6} (oracle drift {drift:.1e}); \
             {:.3} ms/step over {TIMED_STEPS} steps",
            formulation.forward_nodes,
            formulation.total_nodes,
            formulation.module.lines().count(),
            seconds * 1000.0 / TIMED_STEPS as f64,
        );
    }
}

/// Returns whether the payload at `index` is a norm gain, which stays
/// out of the `wrt` set and therefore out of the response.
fn is_norm_gain(index: usize) -> bool {
    let after_heads = 2 + 3 * HEAD_COUNT;
    index == after_heads + 1 || index == after_heads + 4 || index == after_heads + 5
}