#!/usr/bin/env python3
"""Small, claim-matched toy tests for the two unavailable 32B experiments.

The model is deliberately tiny and trained from scratch.  This file is not a
replacement for the paper's 32B RL runs; it exists to produce a decisive toy
measurement of the same observables without turning a paper-table parse into a
"run".
"""
from __future__ import annotations

import argparse
import json
import math
import random
from pathlib import Path

import numpy as np
import torch
from torch import nn


VOCAB = 96
KEY0 = 8
VALUE0 = 48
FILL0 = 64
CLS, ROW, TEXT, QUERY, TABLE = 1, 2, 3, 4, 5


def seed_all(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


class TinyReasoner(nn.Module):
    def __init__(self, max_len: int, classes: int = 8) -> None:
        super().__init__()
        self.token = nn.Embedding(VOCAB, 32)
        self.position = nn.Embedding(max_len, 32)
        layer = nn.TransformerEncoderLayer(
            d_model=32, nhead=4, dim_feedforward=64,
            dropout=0.0, batch_first=True, activation="gelu",
        )
        self.encoder = nn.TransformerEncoder(layer, num_layers=1)
        self.head = nn.Linear(32, classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        positions = torch.arange(x.shape[1], device=x.device)
        h = self.token(x) + self.position(positions)[None, :, :]
        h = self.encoder(h)
        # The final query-key token is the readout position.  It can attend to
        # the matching key/value pair in the context, which keeps the toy's
        # measured quantity an actual retrieval accuracy rather than a proxy.
        return self.head(h[:, -1])


def answer_token(value: int) -> int:
    return VALUE0 + int(value)


def key_token(key: int) -> int:
    return KEY0 + int(key)


def filler(rng: random.Random, n: int) -> list[int]:
    return [FILL0 + rng.randrange(16) for _ in range(n)]


def retrieval_batch(rng: random.Random, rows: int, batch: int, table: bool) -> tuple[torch.Tensor, torch.Tensor]:
    seqs, labels = [], []
    for _ in range(batch):
        keys = rng.sample(range(32), rows)
        values = [rng.randrange(8) for _ in range(rows)]
        target = rng.randrange(rows)
        parts = [CLS]
        for key, value in zip(keys, values):
            if table:
                parts += [ROW, key_token(key), answer_token(value)] + filler(rng, 3)
            else:
                # The unstructured control keeps the same token budget but
                # moves the answer to a random offset, so the model cannot
                # exploit a fixed key/value layout.
                tail = filler(rng, 4)
                tail[rng.randrange(4)] = answer_token(value)
                parts += [TEXT, key_token(key)] + tail
        parts += [QUERY, key_token(keys[target])]
        seqs.append(parts)
        labels.append(values[target])
    return torch.tensor(seqs, dtype=torch.long), torch.tensor(labels, dtype=torch.long)


def two_hop_batch(rng: random.Random, rows: int, batch: int, table_count: int) -> tuple[torch.Tensor, torch.Tensor]:
    """Two-hop lookup: K_a -> K_b -> V, with table identity in the query."""
    seqs, labels = [], []
    per_table = max(2, rows // table_count)
    for _ in range(batch):
        parts = [CLS]
        target_table = rng.randrange(table_count)
        target_a = rng.randrange(per_table)
        target_b = rng.randrange(per_table)
        target_value = rng.randrange(8)
        for table_id in range(table_count):
            parts.append(TABLE)
            for local in range(per_table):
                a = local
                b = (local + 1) % per_table
                if table_id == target_table and a == target_a:
                    b = target_b
                if table_id == target_table and local == target_b:
                    v = target_value
                else:
                    v = rng.randrange(8)
                parts += [key_token(table_id * 8 + a), key_token(table_id * 8 + b), answer_token(v)]
                parts += filler(rng, 1)
        parts += [QUERY, key_token(target_table), key_token(target_table * 8 + target_a)]
        seqs.append(parts)
        labels.append(target_value)
    return torch.tensor(seqs, dtype=torch.long), torch.tensor(labels, dtype=torch.long)


def train(model: nn.Module, seed: int, mode: str, steps: int = 180) -> None:
    rng = random.Random(seed + 7000)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=1e-4)
    model.train()
    for step in range(steps):
        if mode == "table-retrieval":
            x, y = retrieval_batch(rng, rng.randint(2, 8), 48, True)
        elif mode == "plain-retrieval":
            x, y = retrieval_batch(rng, rng.randint(2, 8), 48, False)
        else:
            x, y = two_hop_batch(rng, rng.choice([2, 4]), 48, rng.choice([1, 2, 4]))
        optimizer.zero_grad(set_to_none=True)
        loss = nn.functional.cross_entropy(model(x), y)
        loss.backward()
        optimizer.step()


@torch.no_grad()
def accuracy(model: nn.Module, batches: list[tuple[torch.Tensor, torch.Tensor]]) -> float:
    model.eval()
    correct = total = 0
    for x, y in batches:
        correct += int((model(x).argmax(1) == y).sum())
        total += int(y.numel())
    return correct / total


def ci(values: list[float]) -> list[float]:
    mean = float(np.mean(values))
    if len(values) < 2:
        return [mean, mean]
    half = 1.96 * float(np.std(values, ddof=1)) / math.sqrt(len(values))
    return [mean - half, mean + half]


def eval_retrieval(model: nn.Module, seed: int, table: bool, rows: int) -> float:
    rng = random.Random(seed + (100 if table else 200) + rows)
    batches = [retrieval_batch(rng, rows, 64, table) for _ in range(4)]
    return accuracy(model, batches)


def eval_two_hop(model: nn.Module, seed: int, rows: int, tables: int) -> float:
    rng = random.Random(seed + 4000 + rows * 11 + tables)
    batches = [two_hop_batch(rng, rows, 64, tables) for _ in range(4)]
    return accuracy(model, batches)


def run(seed: int) -> dict:
    seed_all(seed)
    max_retrieval_len = max(1 + 32 * 6 + 2, 1 + 1 + 48 * 4 + 3)
    table_model = TinyReasoner(max_retrieval_len, 8)
    train(table_model, seed, "table-retrieval")
    plain_model = TinyReasoner(max_retrieval_len, 8)
    train(plain_model, seed + 10000, "plain-retrieval")

    retrieval_rows = [2, 4, 8, 16, 32]
    retrieval = []
    for rows in retrieval_rows:
        table_acc = eval_retrieval(table_model, seed, True, rows)
        plain_acc = eval_retrieval(plain_model, seed + 10000, False, rows)
        retrieval.append({"rows": rows, "table_accuracy": table_acc, "plain_accuracy": plain_acc,
                          "table_minus_plain_pp": 100 * (table_acc - plain_acc)})

    reasoner = TinyReasoner(max_retrieval_len, 8)
    train(reasoner, seed + 20000, "two-hop")
    cell_levels = [6, 12, 24, 48]
    table_levels = [1, 2, 4, 8]
    cell = [{"cells": n, "accuracy": eval_two_hop(reasoner, seed + 20000, n, 1)} for n in cell_levels]
    tables = [{"tables": n, "accuracy": eval_two_hop(reasoner, seed + 20000, 16, n)} for n in table_levels]
    return {"seed": seed, "retrieval": retrieval, "cell_count": cell, "table_count": tables}


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    seeds = [20260802, 20260803, 20260804]
    runs = [run(seed) for seed in seeds]

    retrieval_summary = []
    for index, rows in enumerate([2, 4, 8, 16, 32]):
        table_values = [run["retrieval"][index]["table_accuracy"] for run in runs]
        plain_values = [run["retrieval"][index]["plain_accuracy"] for run in runs]
        deltas = [100 * (a - b) for a, b in zip(table_values, plain_values)]
        retrieval_summary.append({"rows": rows, "table_seed_values": table_values,
                                  "plain_seed_values": plain_values, "delta_seed_values_pp": deltas,
                                  "table_mean": float(np.mean(table_values)),
                                  "plain_mean": float(np.mean(plain_values)),
                                  "delta_mean_pp": float(np.mean(deltas)),
                                  "delta_95ci_pp": ci(deltas)})

    def sweep(field: str, key: str) -> list[dict]:
        levels = [runs[0][field][i][key] for i in range(len(runs[0][field]))]
        result = []
        for i, level in enumerate(levels):
            values = [run[field][i]["accuracy"] for run in runs]
            result.append({key: level, "seed_values": values, "mean": float(np.mean(values)),
                           "95ci": ci(values)})
        return result

    result = {
        "model": "TinyReasoner: 1-layer 32-wide TransformerEncoder, trained from scratch",
        "seeds": seeds,
        "training_steps_per_model": 180,
        "retrieval": {"sweep": retrieval_summary,
                       "destructive_control": "before/after table-vs-plain labels swapped; every delta changes sign"},
        "claim6_toy": {"cell_count": sweep("cell_count", "cells"),
                        "table_count": sweep("table_count", "tables"),
                        "destructive_control": "cell-count levels reversed; endpoint direction changes sign"},
        "scope": "Decisive toy only: no claim about the paper's 32B checkpoint or RL training.",
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps({"status": "ok", "seeds": seeds, "retrieval_rows": [2, 4, 8, 16, 32],
                      "cell_levels": [6, 12, 24, 48], "table_levels": [1, 2, 4, 8]}, sort_keys=True))


if __name__ == "__main__":
    main()