///
/// The most atomic way to train and run inference for a GPT in pure, dependency-free Zig.
/// This file is the complete algorithm.
/// Everything else is just efficiency.
///
/// @karpathy (original Python), @xenova (JavaScript port), Zig port
///
const std = @import("std");
const math = std.math;
const Allocator = std.mem.Allocator;
const Io = std.Io;
const ArrayList = std.ArrayList;

// ─── Mersenne Twister 19937 (matches Python's random module exactly) ───

var mt_state: [624]u32 = undefined;
var mt_idx: u32 = 625;
var gauss_next: ?f64 = null;

fn mtSeed(n: u64) void {
    var key_buf: [4]u32 = undefined;
    var key_len: usize = 0;
    var v = n;
    if (v == 0) {
        key_buf[0] = 0;
        key_len = 1;
    } else {
        while (v > 0) : (v = v / 0x100000000) {
            key_buf[key_len] = @as(u32, @intCast(v & 0xFFFFFFFF));
            key_len += 1;
        }
    }
    const key = key_buf[0..key_len];

    mt_state[0] = 19650218;
    for (1..624) |idx| {
        const prev = mt_state[idx - 1];
        const xor = prev ^ (prev >> 30);
        const product: u32 = @truncate(@as(u64, 1812433253) *% @as(u64, xor));
        mt_state[idx] = product +% @as(u32, @intCast(idx));
    }
    var i: usize = 1;
    var j: usize = 0;
    var k: usize = @max(624, key.len);
    while (k > 0) : (k -= 1) {
        if (i >= 624) {
            mt_state[0] = mt_state[623];
            i = 1;
        }
        if (j >= key.len) j = 0;
        const prev1 = mt_state[i - 1];
        const xor1: u32 = @truncate(@as(u64, 1664525) *% @as(u64, prev1 ^ (prev1 >> 30)));
        mt_state[i] = (mt_state[i] ^ xor1) +% key[j] +% @as(u32, @intCast(j));
        i += 1;
        j += 1;
    }
    k = 623;
    while (k > 0) : (k -= 1) {
        if (i >= 624) {
            mt_state[0] = mt_state[623];
            i = 1;
        }
        const prev2 = mt_state[i - 1];
        const xor2: u32 = @truncate(@as(u64, 1566083941) *% @as(u64, prev2 ^ (prev2 >> 30)));
        mt_state[i] = (mt_state[i] ^ xor2) -% @as(u32, @intCast(i));
        i += 1;
    }
    mt_state[0] = 0x80000000;
    mt_idx = 624;
    gauss_next = null;
}

fn mtInt32() u32 {
    if (mt_idx >= 624) {
        for (0..624) |k| {
            const y = (mt_state[k] & 0x80000000) | (mt_state[(k + 1) % 624] & 0x7FFFFFFF);
            mt_state[k] = mt_state[(k + 397) % 624] ^ (y >> 1) ^ (if (y & 1 != 0) @as(u32, 0x9908B0DF) else @as(u32, 0));
        }
        mt_idx = 0;
    }
    var y = mt_state[mt_idx];
    mt_idx += 1;
    y ^= y >> 11;
    y ^= (y << 7) & 0x9D2C5680;
    y ^= (y << 15) & 0xEFC60000;
    y ^= y >> 18;
    return y;
}

fn mtRandom() f64 {
    const a: f64 = @floatFromInt(mtInt32() >> 5);
    const b: f64 = @floatFromInt(mtInt32() >> 6);
    return (a * 67108864.0 + b) / 9007199254740992.0;
}

fn mtGauss(mu: f64, sigma: f64) f64 {
    var z: f64 = undefined;
    if (gauss_next) |gn| {
        z = gn;
        gauss_next = null;
    } else {
        const x2pi = mtRandom() * 2.0 * math.pi;
        const g2rad = @sqrt(-2.0 * @log(1.0 - mtRandom()));
        z = @cos(x2pi) * g2rad;
        gauss_next = @sin(x2pi) * g2rad;
    }
    return mu + z * sigma;
}

fn mtShuffle(comptime T: type, arr: []T) void {
    var i: usize = arr.len - 1;
    while (i > 0) : (i -= 1) {
        const k: u32 = 32 - @clz(@as(u32, @intCast(i + 1)));
        const shift: u5 = @intCast(32 -| k); // saturating subtract for the case k=32
        var r: usize = @intCast(mtInt32() >> shift);
        while (r > i) {
            r = @intCast(mtInt32() >> shift);
        }
        std.mem.swap(T, &arr[i], &arr[r]);
    }
}

fn mtChoices(weights: []const f64, alloc: Allocator) usize {
    const n = weights.len;
    const cum = alloc.alloc(f64, n) catch unreachable;
    cum[0] = weights[0];
    for (1..n) |i| {
        cum[i] = cum[i - 1] + weights[i];
    }
    const x = mtRandom() * cum[n - 1];
    var lo: usize = 0;
    var hi: usize = n - 1;
    while (lo < hi) {
        const mid = (lo + hi) >> 1;
        if (x < cum[mid]) {
            hi = mid;
        } else {
            lo = mid + 1;
        }
    }
    return lo;
}

// ─── Value (Autograd) ───

const Value = struct {
    data: f64,
    grad: f64,
    c0: ?*Value,
    c1: ?*Value,
    lg0: f64,
    lg1: f64,
    nch: u2,
    gen: u32,

    fn create(alloc: Allocator, data: f64) *Value {
        const v = alloc.create(Value) catch unreachable;
        v.* = .{ .data = data, .grad = 0, .c0 = null, .c1 = null, .lg0 = 0, .lg1 = 0, .nch = 0, .gen = 0 };
        return v;
    }

    fn bin(alloc: Allocator, data: f64, a: *Value, b: *Value, lg0: f64, lg1: f64) *Value {
        const v = alloc.create(Value) catch unreachable;
        v.* = .{ .data = data, .grad = 0, .c0 = a, .c1 = b, .lg0 = lg0, .lg1 = lg1, .nch = 2, .gen = 0 };
        return v;
    }

    fn una(alloc: Allocator, data: f64, a: *Value, lg0: f64) *Value {
        const v = alloc.create(Value) catch unreachable;
        v.* = .{ .data = data, .grad = 0, .c0 = a, .c1 = null, .lg0 = lg0, .lg1 = 0, .nch = 1, .gen = 0 };
        return v;
    }

    fn addV(self: *Value, other: *Value, alloc: Allocator) *Value {
        return bin(alloc, self.data + other.data, self, other, 1, 1);
    }

    fn addS(self: *Value, s: f64, alloc: Allocator) *Value {
        return una(alloc, self.data + s, self, 1);
    }

    fn mulV(self: *Value, other: *Value, alloc: Allocator) *Value {
        return bin(alloc, self.data * other.data, self, other, other.data, self.data);
    }

    fn mulS(self: *Value, s: f64, alloc: Allocator) *Value {
        return una(alloc, self.data * s, self, s);
    }

    fn powS(self: *Value, p: f64, alloc: Allocator) *Value {
        return una(alloc, math.pow(f64, self.data, p), self, p * math.pow(f64, self.data, p - 1));
    }

    fn logV(self: *Value, alloc: Allocator) *Value {
        return una(alloc, @log(self.data), self, 1.0 / self.data);
    }

    fn expV(self: *Value, alloc: Allocator) *Value {
        const e = @exp(self.data);
        return una(alloc, e, self, e);
    }

    fn relu(self: *Value, alloc: Allocator) *Value {
        return una(alloc, @max(0.0, self.data), self, if (self.data > 0) @as(f64, 1.0) else @as(f64, 0.0));
    }

    fn neg(self: *Value, alloc: Allocator) *Value {
        return una(alloc, -self.data, self, -1.0);
    }

    fn subS(self: *Value, s: f64, alloc: Allocator) *Value {
        return self.addS(-s, alloc);
    }

    fn divV(self: *Value, other: *Value, alloc: Allocator) *Value {
        return self.mulV(other.powS(-1, alloc), alloc);
    }

    fn divS(self: *Value, s: f64, alloc: Allocator) *Value {
        return self.mulS(1.0 / s, alloc);
    }
};

var backward_gen: u32 = 0;

fn doBackward(loss: *Value, alloc: Allocator) void {
    backward_gen += 1;
    const gen = backward_gen;

    var topo: ArrayList(*Value) = .empty;
    const Frame = struct { v: *Value, state: u8 };
    var stack: ArrayList(Frame) = .empty;
    stack.append(alloc, .{ .v = loss, .state = 0 }) catch unreachable;

    while (stack.items.len > 0) {
        const frame = &stack.items[stack.items.len - 1];
        if (frame.v.gen == gen and frame.state == 0) {
            _ = stack.pop();
            continue;
        }
        if (frame.state == 0) {
            frame.v.gen = gen;
            if (frame.v.nch >= 1) {
                frame.state = 1;
                stack.append(alloc, .{ .v = frame.v.c0.?, .state = 0 }) catch unreachable;
                continue;
            }
        }
        if (frame.state <= 1) {
            if (frame.v.nch == 2) {
                frame.state = 2;
                stack.append(alloc, .{ .v = frame.v.c1.?, .state = 0 }) catch unreachable;
                continue;
            }
        }
        topo.append(alloc, frame.v) catch unreachable;
        _ = stack.pop();
    }

    loss.grad = 1;
    var i: usize = topo.items.len;
    while (i > 0) {
        i -= 1;
        const v = topo.items[i];
        const g = v.grad;
        if (v.nch >= 1) v.c0.?.grad += v.lg0 * g;
        if (v.nch == 2) v.c1.?.grad += v.lg1 * g;
    }
}

// ─── Helper functions ───

fn sumValues(vals: []*Value, alloc: Allocator) *Value {
    var result = vals[0];
    for (vals[1..]) |v| {
        result = result.addV(v, alloc);
    }
    return result;
}

fn doLinear(x: []*Value, w: []const []*Value, alloc: Allocator) []*Value {
    const nout = w.len;
    const out = alloc.alloc(*Value, nout) catch unreachable;
    for (0..nout) |oi| {
        const row = w[oi];
        var acc = row[0].mulV(x[0], alloc);
        for (1..x.len) |j| {
            acc = acc.addV(row[j].mulV(x[j], alloc), alloc);
        }
        out[oi] = acc;
    }
    return out;
}

fn doSoftmax(logits: []*Value, alloc: Allocator) []*Value {
    var max_val: f64 = -math.inf(f64);
    for (logits) |v| {
        if (v.data > max_val) max_val = v.data;
    }
    const n = logits.len;
    const exps = alloc.alloc(*Value, n) catch unreachable;
    for (0..n) |i| {
        exps[i] = logits[i].subS(max_val, alloc).expV(alloc);
    }
    const total = sumValues(exps, alloc);
    const result = alloc.alloc(*Value, n) catch unreachable;
    for (0..n) |i| {
        result[i] = exps[i].divV(total, alloc);
    }
    return result;
}

fn doRmsnorm(x: []*Value, alloc: Allocator) []*Value {
    const n = x.len;
    const sq = alloc.alloc(*Value, n) catch unreachable;
    for (0..n) |i| {
        sq[i] = x[i].mulV(x[i], alloc);
    }
    const ms = sumValues(sq, alloc).mulS(1.0 / @as(f64, @floatFromInt(n)), alloc);
    const s = ms.addS(1e-5, alloc).powS(-0.5, alloc);
    const result = alloc.alloc(*Value, n) catch unreachable;
    for (0..n) |i| {
        result[i] = x[i].mulV(s, alloc);
    }
    return result;
}

// ─── Model hyperparameters ───
const n_embd = 16;
const n_head = 4;
const n_layer = 1;
const block_size = 16;
const head_dim = n_embd / n_head;
const attn_scale: f64 = 1.0 / @sqrt(@as(f64, head_dim));

const StateDict = struct {
    wte: [][]*Value,
    wpe: [][]*Value,
    lm_head: [][]*Value,
    attn_wq: [n_layer][][]*Value,
    attn_wk: [n_layer][][]*Value,
    attn_wv: [n_layer][][]*Value,
    attn_wo: [n_layer][][]*Value,
    mlp_fc1: [n_layer][][]*Value,
    mlp_fc2: [n_layer][][]*Value,
};

fn makeMatrix(alloc: Allocator, nout: usize, nin: usize, stdev: f64) [][]*Value {
    const rows = alloc.alloc([]*Value, nout) catch unreachable;
    for (0..nout) |i| {
        const row = alloc.alloc(*Value, nin) catch unreachable;
        for (0..nin) |j| {
            row[j] = Value.create(alloc, mtGauss(0, stdev));
        }
        rows[i] = row;
    }
    return rows;
}

fn doGpt(
    token_id: usize,
    pos_id: usize,
    keys: *[n_layer]ArrayList([]*Value),
    kv_values: *[n_layer]ArrayList([]*Value),
    sd: *const StateDict,
    alloc: Allocator,
) []*Value {
    const tok_emb = sd.wte[token_id];
    const pos_emb = sd.wpe[pos_id];
    var x = alloc.alloc(*Value, n_embd) catch unreachable;
    for (0..n_embd) |i| {
        x[i] = tok_emb[i].addV(pos_emb[i], alloc);
    }
    x = doRmsnorm(x, alloc);

    for (0..n_layer) |li| {
        // 1) Multi-head attention block
        const x_residual = x;
        var xn = doRmsnorm(x, alloc);

        const q = doLinear(xn, sd.attn_wq[li], alloc);
        const k_vec = doLinear(xn, sd.attn_wk[li], alloc);
        const v_vec = doLinear(xn, sd.attn_wv[li], alloc);

        keys[li].append(alloc, k_vec) catch unreachable;
        kv_values[li].append(alloc, v_vec) catch unreachable;

        const x_attn = alloc.alloc(*Value, n_embd) catch unreachable;
        for (0..n_head) |h| {
            const hs = h * head_dim;
            const q_h = q[hs .. hs + head_dim];

            const num_t = keys[li].items.len;
            const attn_logits = alloc.alloc(*Value, num_t) catch unreachable;
            for (0..num_t) |t| {
                const kt = keys[li].items[t];
                const dot = alloc.alloc(*Value, head_dim) catch unreachable;
                for (0..head_dim) |j_| {
                    dot[j_] = q_h[j_].mulV(kt[hs + j_], alloc);
                }
                attn_logits[t] = sumValues(dot, alloc).mulS(attn_scale, alloc);
            }
            const attn_weights = doSoftmax(attn_logits, alloc);

            for (0..head_dim) |j_| {
                const weighted = alloc.alloc(*Value, num_t) catch unreachable;
                for (0..num_t) |t| {
                    weighted[t] = attn_weights[t].mulV(kv_values[li].items[t][hs + j_], alloc);
                }
                x_attn[hs + j_] = sumValues(weighted, alloc);
            }
        }

        xn = doLinear(x_attn, sd.attn_wo[li], alloc);
        // residual connection
        const x_res1 = alloc.alloc(*Value, n_embd) catch unreachable;
        for (0..n_embd) |i| {
            x_res1[i] = xn[i].addV(x_residual[i], alloc);
        }

        // 2) MLP block
        const x_residual2 = x_res1;
        const xn2 = doRmsnorm(x_res1, alloc);
        var fc1_out = doLinear(xn2, sd.mlp_fc1[li], alloc);
        for (0..fc1_out.len) |i| {
            fc1_out[i] = fc1_out[i].relu(alloc);
        }
        const fc2_out = doLinear(fc1_out, sd.mlp_fc2[li], alloc);
        x = alloc.alloc(*Value, n_embd) catch unreachable;
        for (0..n_embd) |i| {
            x[i] = fc2_out[i].addV(x_residual2[i], alloc);
        }
    }

    return doLinear(x, sd.lm_head, alloc);
}

// ─── f64 fast-path for inference (no autograd) ───

fn doLinearF64(x: []const f64, w: []const []*Value, alloc: Allocator) []f64 {
    const nout = w.len;
    const out = alloc.alloc(f64, nout) catch unreachable;
    for (0..nout) |oi| {
        const row = w[oi];
        var acc: f64 = 0;
        for (0..x.len) |j| {
            acc += row[j].data * x[j];
        }
        out[oi] = acc;
    }
    return out;
}

fn doSoftmaxF64(logits: []f64) []f64 {
    var max_val: f64 = -math.inf(f64);
    for (logits) |v| {
        if (v > max_val) max_val = v;
    }
    var total: f64 = 0;
    for (logits) |*v| {
        v.* = @exp(v.* - max_val);
        total += v.*;
    }
    for (logits) |*v| {
        v.* /= total;
    }
    return logits;
}

fn doRmsnormF64(x: []const f64, alloc: Allocator) []f64 {
    const n = x.len;
    var ms: f64 = 0;
    for (x) |xi| {
        ms += xi * xi;
    }
    ms /= @floatFromInt(n);
    const s = 1.0 / @sqrt(ms + 1e-5);
    const result = alloc.alloc(f64, n) catch unreachable;
    for (0..n) |i| {
        result[i] = x[i] * s;
    }
    return result;
}

fn doGptF64(
    token_id: usize,
    pos_id: usize,
    keys: *[n_layer]ArrayList([]f64),
    kv_values: *[n_layer]ArrayList([]f64),
    sd: *const StateDict,
    alloc: Allocator,
) []f64 {
    const tok_emb = sd.wte[token_id];
    const pos_emb = sd.wpe[pos_id];
    var x = alloc.alloc(f64, n_embd) catch unreachable;
    for (0..n_embd) |i| {
        x[i] = tok_emb[i].data + pos_emb[i].data;
    }
    x = doRmsnormF64(x, alloc);

    for (0..n_layer) |li| {
        // 1) Multi-head attention block
        const x_residual = x;
        var xn = doRmsnormF64(x, alloc);

        const q = doLinearF64(xn, sd.attn_wq[li], alloc);
        const k_vec = doLinearF64(xn, sd.attn_wk[li], alloc);
        const v_vec = doLinearF64(xn, sd.attn_wv[li], alloc);

        keys[li].append(alloc, k_vec) catch unreachable;
        kv_values[li].append(alloc, v_vec) catch unreachable;

        const x_attn = alloc.alloc(f64, n_embd) catch unreachable;
        for (0..n_head) |h| {
            const hs = h * head_dim;
            const q_h = q[hs .. hs + head_dim];

            const num_t = keys[li].items.len;
            const attn_logits = alloc.alloc(f64, num_t) catch unreachable;
            for (0..num_t) |t| {
                const kt = keys[li].items[t];
                var dot: f64 = 0;
                for (0..head_dim) |j_| {
                    dot += q_h[j_] * kt[hs + j_];
                }
                attn_logits[t] = dot * attn_scale;
            }
            const attn_weights = doSoftmaxF64(attn_logits);

            for (0..head_dim) |j_| {
                var sum: f64 = 0;
                for (0..num_t) |t| {
                    sum += attn_weights[t] * kv_values[li].items[t][hs + j_];
                }
                x_attn[hs + j_] = sum;
            }
        }

        xn = doLinearF64(x_attn, sd.attn_wo[li], alloc);
        // residual connection
        const x_res1 = alloc.alloc(f64, n_embd) catch unreachable;
        for (0..n_embd) |i| {
            x_res1[i] = xn[i] + x_residual[i];
        }

        // 2) MLP block
        const x_residual2 = x_res1;
        const xn2 = doRmsnormF64(x_res1, alloc);
        const fc1_out = doLinearF64(xn2, sd.mlp_fc1[li], alloc);
        for (fc1_out) |*v| {
            v.* = @max(0.0, v.*);
        }
        const fc2_out = doLinearF64(fc1_out, sd.mlp_fc2[li], alloc);
        x = alloc.alloc(f64, n_embd) catch unreachable;
        for (0..n_embd) |i| {
            x[i] = fc2_out[i] + x_residual2[i];
        }
    }

    return doLinearF64(x, sd.lm_head, alloc);
}

// ─── Main ───
pub fn main(init: std.process.Init) !void {
    const io = init.io;
    const alloc = init.gpa;

    var stdout_buffer: [4096]u8 = undefined;
    var stdout_writer = Io.File.stdout().writer(io, &stdout_buffer);
    const stdout = &stdout_writer.interface;

    // Read input.txt
    const file_data = try Io.Dir.cwd().readFileAlloc(io, "input.txt", alloc, .limited(10 * 1024 * 1024));
    defer alloc.free(file_data);

    // Split into docs
    var doc_list: ArrayList([]const u8) = .empty;
    defer doc_list.deinit(alloc);
    const trimmed_data = std.mem.trim(u8, file_data, &[_]u8{ '\n', '\r', ' ', '\t' });
    var line_iter = std.mem.splitScalar(u8, trimmed_data, '\n');
    while (line_iter.next()) |line| {
        const trimmed = std.mem.trim(u8, line, &[_]u8{ ' ', '\t', '\r' });
        if (trimmed.len > 0) {
            try doc_list.append(alloc, trimmed);
        }
    }
    var docs = doc_list.items;

    mtSeed(42);
    mtShuffle([]const u8, docs);
    try stdout.print("num docs: {d}\n", .{docs.len});

    // Tokenizer: unique characters
    var char_set = std.AutoHashMap(u8, void).init(alloc);
    defer char_set.deinit();
    for (docs) |doc| {
        for (doc) |ch| {
            try char_set.put(ch, {});
        }
    }
    var uchars_list: ArrayList(u8) = .empty;
    defer uchars_list.deinit(alloc);
    var it = char_set.keyIterator();
    while (it.next()) |key| {
        try uchars_list.append(alloc, key.*);
    }
    std.mem.sort(u8, uchars_list.items, {}, std.sort.asc(u8));
    const uchars = uchars_list.items;

    var char_to_id = std.AutoHashMap(u8, usize).init(alloc);
    defer char_to_id.deinit();
    for (uchars, 0..) |ch, i| {
        try char_to_id.put(ch, i);
    }

    const BOS = uchars.len;
    const vocab_size = uchars.len + 1;
    try stdout.print("vocab size: {d}\n", .{vocab_size});

    // Initialize parameters
    const param_alloc = std.heap.page_allocator;

    var sd: StateDict = undefined;
    sd.wte = makeMatrix(param_alloc, vocab_size, n_embd, 0.08);
    sd.wpe = makeMatrix(param_alloc, block_size, n_embd, 0.08);
    sd.lm_head = makeMatrix(param_alloc, vocab_size, n_embd, 0.08);
    for (0..n_layer) |i| {
        sd.attn_wq[i] = makeMatrix(param_alloc, n_embd, n_embd, 0.08);
        sd.attn_wk[i] = makeMatrix(param_alloc, n_embd, n_embd, 0.08);
        sd.attn_wv[i] = makeMatrix(param_alloc, n_embd, n_embd, 0.08);
        sd.attn_wo[i] = makeMatrix(param_alloc, n_embd, n_embd, 0.08);
        sd.mlp_fc1[i] = makeMatrix(param_alloc, 4 * n_embd, n_embd, 0.08);
        sd.mlp_fc2[i] = makeMatrix(param_alloc, n_embd, 4 * n_embd, 0.08);
    }

    // Flatten params
    var params_list: ArrayList(*Value) = .empty;
    const matrices = [_][][]*Value{ sd.wte, sd.wpe, sd.lm_head };
    for (matrices) |mat_| {
        for (mat_) |row| {
            for (row) |p| {
                try params_list.append(param_alloc, p);
            }
        }
    }
    for (0..n_layer) |li| {
        const layer_mats = [_][][]*Value{
            sd.attn_wq[li], sd.attn_wk[li], sd.attn_wv[li], sd.attn_wo[li],
            sd.mlp_fc1[li], sd.mlp_fc2[li],
        };
        for (layer_mats) |mat_| {
            for (mat_) |row| {
                for (row) |p| {
                    try params_list.append(param_alloc, p);
                }
            }
        }
    }
    const params = params_list.items;
    try stdout.print("num params: {d}\n", .{params.len});
    try stdout_writer.flush();

    // Adam optimizer buffers
    const m_buf = try param_alloc.alloc(f64, params.len);
    const v_buf = try param_alloc.alloc(f64, params.len);
    @memset(m_buf, 0);
    @memset(v_buf, 0);

    const learning_rate: f64 = 0.01;
    const beta1: f64 = 0.85;
    const beta2: f64 = 0.99;
    const eps_adam: f64 = 1e-8;
    const num_steps: usize = 1000;

    // Training loop
    for (0..num_steps) |step| {
        const doc = docs[step % docs.len];
        var tokens: ArrayList(usize) = .empty;
        defer tokens.deinit(alloc);
        try tokens.ensureTotalCapacity(alloc, doc.len + 2);
        tokens.appendAssumeCapacity(BOS);
        for (doc) |ch| {
            tokens.appendAssumeCapacity(char_to_id.get(ch).?);
        }
        tokens.appendAssumeCapacity(BOS);
        const n = @min(block_size, tokens.items.len - 1);

        var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
        defer arena.deinit();
        const step_alloc = arena.allocator();

        var keys: [n_layer]ArrayList([]*Value) = undefined;
        var kv_values: [n_layer]ArrayList([]*Value) = undefined;
        for (0..n_layer) |li| {
            keys[li] = .empty;
            kv_values[li] = .empty;
        }

        var losses: ArrayList(*Value) = .empty;
        losses.ensureTotalCapacity(step_alloc, n) catch unreachable;
        for (0..n) |pos_id| {
            const token_id = tokens.items[pos_id];
            const target_id = tokens.items[pos_id + 1];
            const logits = doGpt(token_id, pos_id, &keys, &kv_values, &sd, step_alloc);
            const probs = doSoftmax(logits, step_alloc);
            const loss_t = probs[target_id].logV(step_alloc).neg(step_alloc);
            losses.append(step_alloc, loss_t) catch unreachable;
        }
        const loss = sumValues(losses.items, step_alloc).mulS(1.0 / @as(f64, @floatFromInt(n)), step_alloc);

        doBackward(loss, step_alloc);

        const lr_t = learning_rate * (1.0 - @as(f64, @floatFromInt(step)) / @as(f64, @floatFromInt(num_steps)));
        const step_f: f64 = @floatFromInt(step + 1);
        const bc1 = 1.0 - math.pow(f64, beta1, step_f);
        const bc2 = 1.0 - math.pow(f64, beta2, step_f);
        for (params, 0..) |p, i| {
            m_buf[i] = beta1 * m_buf[i] + (1.0 - beta1) * p.grad;
            v_buf[i] = beta2 * v_buf[i] + (1.0 - beta2) * p.grad * p.grad;
            const m_hat = m_buf[i] / bc1;
            const v_hat = v_buf[i] / bc2;
            p.data -= lr_t * m_hat / (@sqrt(v_hat) + eps_adam);
            p.grad = 0;
        }

        try stdout.print("step {d:>4} / {d:>4} | loss {d:.4}\r", .{ step + 1, num_steps, loss.data });
        try stdout_writer.flush();
    }

    // Inference
    const temperature: f64 = 0.5;
    try stdout.print("\n--- inference (new, hallucinated names) ---\n", .{});
    try stdout_writer.flush();

    for (0..20) |sample_idx| {
        var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
        defer arena.deinit();
        const inf_alloc = arena.allocator();

        var f_keys: [n_layer]ArrayList([]f64) = undefined;
        var f_values: [n_layer]ArrayList([]f64) = undefined;
        for (0..n_layer) |li| {
            f_keys[li] = .empty;
            f_values[li] = .empty;
        }

        var token_id: usize = BOS;
        var sample: ArrayList(u8) = .empty;

        for (0..block_size) |pos_id| {
            const logits = doGptF64(token_id, pos_id, &f_keys, &f_values, &sd, inf_alloc);
            const scaled = inf_alloc.alloc(f64, logits.len) catch unreachable;
            for (0..logits.len) |i| {
                scaled[i] = logits[i] / temperature;
            }
            const probs = doSoftmaxF64(scaled);
            token_id = mtChoices(probs, inf_alloc);
            if (token_id == BOS) break;
            sample.append(inf_alloc, uchars[token_id]) catch unreachable;
        }
        try stdout.print("sample {d:>2}: {s}\n", .{ sample_idx + 1, sample.items });
    }
    try stdout_writer.flush();
}