"""Learn next tokens in a tiny Python subset. This is not a physics PINN.

Python 3.10+ and PyTorch 2.4+. Each file runs without the other examples.
Run: python token_rules.py train --output-dir runs/tokens
"""

from __future__ import annotations

import argparse
import ast
import copy
import itertools
import json
import random
from pathlib import Path

import torch
from torch import nn
from torch.nn import functional as F

VOCAB = ["<PAD>", "<BOS>", "<EOS>", "print", "(", "0", "1", "2", "+", ")", "\n"]
TOKEN_ID = {token: index for index, token in enumerate(VOCAB)}
PAD, BOS, EOS = 0, 1, 2
NUMBERS = {"0", "1", "2"}
ALLOWED = {
    "start": {"print"},
    "open": {"("},
    "first": NUMBERS,
    "after_first": {"+", ")"},
    "second": NUMBERS,
    "close": {")"},
    "newline": {"\n"},
    "after_line": {"print", "<EOS>"},
    "done": set(),
}


def step(state: str, token: str) -> str:
    """Require a legal token. Return the state after consuming it."""
    if state not in ALLOWED or token not in ALLOWED[state]:
        raise ValueError(f"Token {token!r} is invalid in state {state!r}")
    if token == "print":
        return "open"
    if token == "(":
        return "first"
    if token in NUMBERS:
        return "after_first" if state == "first" else "close"
    return {"+": "second", ")": "newline", "\n": "after_line", "<EOS>": "done"}[token]


def prefix_state(tokens: list[str]) -> str:
    """Require a nonempty BOS-led legal prefix. Return its grammar state."""
    if not tokens or tokens[0] != "<BOS>":
        raise ValueError("A prefix must start with <BOS>")
    state = "start"
    # Invariant: state describes tokens consumed so far. Bound: remaining tokens.
    for token in tokens[1:]:
        state = step(state, token)
    return state


def allowed_next(tokens: list[str]) -> set[str]:
    """Return legal next tokens using only the prefix, never its target label."""
    return set(ALLOWED[prefix_state(tokens)])


def is_complete(tokens: list[str]) -> bool:
    """Return whether a sequence is a complete program in our restricted subset."""
    try:
        return prefix_state(tokens) == "done"
    except (ValueError, TypeError):
        return False


def to_code(tokens: list[str]) -> str:
    """Require a complete subset program. Return actual valid Python source."""
    if not is_complete(tokens):
        raise ValueError("Cannot render an incomplete or invalid program")
    return " ".join(tokens[1:-1]).replace(" \n ", "\n").replace(" \n", "\n")


def evaluate_program(tokens: list[str]) -> list[int]:
    """Evaluate whitelisted integer print expressions without executing source."""
    tree = ast.parse(to_code(tokens))
    output = []
    # Invariant: output matches prior statements. Bound: remaining statements.
    for statement in tree.body:
        call = statement.value
        if (
            not isinstance(call, ast.Call)
            or not isinstance(call.func, ast.Name)
            or call.func.id != "print"
        ):
            raise ValueError("Only integer print expressions are supported")
        expression = call.args[0]
        if isinstance(expression, ast.Constant):
            output.append(expression.value)
        elif isinstance(expression, ast.BinOp) and isinstance(expression.op, ast.Add):
            output.append(expression.left.value + expression.right.value)
        else:
            raise ValueError("Only a number or a two-number sum is supported")
    return output


def make_dataset(seed: int = 7) -> list[dict]:
    """Return 192 unique valid programs with fixed 96/32/64 whole-program splits."""
    statements = [["print", "(", n, ")", "\n"] for n in sorted(NUMBERS)]
    statements += [
        ["print", "(", a, "+", b, ")", "\n"]
        for a, b in itertools.product(sorted(NUMBERS), repeat=2)
    ]
    programs = []
    # Bound: three finite Cartesian products. Every program has one to three lines.
    for lines in range(1, 4):
        for choices in itertools.product(statements, repeat=lines):
            programs.append(["<BOS>"] + list(itertools.chain.from_iterable(choices)) + ["<EOS>"])
    random.Random(seed).shuffle(programs)
    rows = []
    for i, tokens in enumerate(programs[:192]):
        split = "train" if i < 96 else "validation" if i < 128 else "test"
        rows.append({"split": split, "tokens": tokens, "code": to_code(tokens)})
    return rows


def write_dataset(path: Path, rows: list[dict]) -> None:
    """Write generated programs as JSONL. Parent path must be writable."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:  # Bound: remaining finite program rows.
            handle.write(json.dumps(row) + "\n")


def load_dataset(path: Path) -> list[dict]:
    """Validate an existing JSONL file, including complete-program split isolation."""
    rows, seen = [], set()
    with path.open(encoding="utf-8") as handle:
        for line in handle:  # Bound: finite lines in the input file.
            row = json.loads(line)
            tokens = row.get("tokens")
            if not isinstance(tokens, list) or not all(isinstance(t, str) for t in tokens):
                raise ValueError("Each row needs a string token list")
            if row.get("split") not in {"train", "validation", "test"} or not is_complete(tokens):
                raise ValueError("Invalid program or split")
            key = tuple(tokens)
            if key in seen or row.get("code") != to_code(tokens):
                raise ValueError("Duplicate program or mismatched source")
            seen.add(key)
            rows.append(row)
    if not rows or any(
        not any(r["split"] == s for r in rows) for s in ("train", "validation", "test")
    ):
        raise ValueError("All three splits must be nonempty")
    return rows


def make_batch(rows: list[dict]) -> tuple[torch.Tensor, ...]:
    """Return input IDs, shifted targets, and legal-token masks with matching shapes."""
    if not rows:
        raise ValueError("A batch needs at least one program")
    length = max(len(row["tokens"]) - 1 for row in rows)
    inputs = torch.full((len(rows), length), PAD, dtype=torch.long)
    targets = inputs.clone()
    legal = torch.zeros(len(rows), length, len(VOCAB), dtype=torch.bool)
    for b, row in enumerate(rows):  # Bound: remaining programs.
        tokens, state = row["tokens"], "start"
        ids = [TOKEN_ID[token] for token in tokens]
        inputs[b, : len(ids) - 1] = torch.tensor(ids[:-1])
        targets[b, : len(ids) - 1] = torch.tensor(ids[1:])
        # Invariant: state is derived only from tokens up to t. Bound: remaining positions.
        for t, token in enumerate(tokens[:-1]):
            if t:
                state = step(state, token)
            legal[b, t, [TOKEN_ID[x] for x in ALLOWED[state]]] = True
    return inputs, targets, legal


class TinyLM(nn.Module):
    """A causal model: integer IDs [B,T] -> logits [B,T,11]. No rule mask enters it."""

    def __init__(self) -> None:
        super().__init__()
        self.embedding = nn.Embedding(len(VOCAB), 8, padding_idx=PAD)
        self.rnn = nn.GRU(8, 16, batch_first=True)
        self.head = nn.Linear(16, len(VOCAB))

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        hidden, _ = self.rnn(self.embedding(ids))
        return self.head(hidden)


def loss_terms(
    logits: torch.Tensor, targets: torch.Tensor, legal: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Average supervised CE and invalid probability mass over non-padding targets."""
    active = targets != PAD
    if not active.any():
        raise ValueError("Loss needs at least one real target")
    ce = F.cross_entropy(logits[active], targets[active])
    probabilities = logits.softmax(dim=-1)
    invalid_mass = (probabilities * (~legal)).sum(dim=-1)
    return ce, invalid_mass[active].mean()


def train_model(
    initial: TinyLM, rows: list[dict], *, epochs: int, rule_weight: float, seed: int
) -> TinyLM:
    """Perform a fixed finite number of steps. Convergence is not guaranteed."""
    if not rows or epochs <= 0 or not torch.isfinite(torch.tensor(rule_weight)) or rule_weight < 0:
        raise ValueError("Use positive epochs and finite nonnegative rule weight")
    model = copy.deepcopy(initial).train()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    rng = random.Random(seed)
    order = list(range(len(rows)))
    for _ in range(epochs):  # Invariant: completed epochs. Bound: remaining epochs.
        rng.shuffle(order)
        for start in range(0, len(rows), 24):  # Bound: remaining batches in this epoch.
            batch = [rows[i] for i in order[start : start + 24]]
            ids, targets, legal = make_batch(batch)
            optimizer.zero_grad(set_to_none=True)
            ce, rule = loss_terms(model(ids), targets, legal)
            loss = ce + rule_weight * rule
            if not torch.isfinite(loss):
                raise RuntimeError("Nonfinite loss")
            loss.backward()
            if any(p.grad is None or not torch.isfinite(p.grad).all() for p in model.parameters()):
                raise RuntimeError("Missing or nonfinite parameter gradient")
            optimizer.step()
    return model.eval()


@torch.no_grad()
def metrics(model: TinyLM, rows: list[dict]) -> dict[str, float]:
    """Measure teacher-forced predictions on a held-out whole-program split."""
    ids, targets, legal = make_batch(rows)
    logits = model(ids)
    ce, mass = loss_terms(logits, targets, legal)
    predictions, active = logits.argmax(dim=-1), targets != PAD
    valid_predictions = legal.gather(-1, predictions.unsqueeze(-1)).squeeze(-1)
    return {
        "cross_entropy": ce.item(),
        "invalid_mass": mass.item(),
        "next_token_accuracy": (predictions[active] == targets[active]).float().mean().item(),
        "invalid_top1_rate": (~valid_predictions[active]).float().mean().item(),
    }


@torch.no_grad()
def generate(model: TinyLM, *, hard_mask: bool, seed: int, max_tokens: int = 32) -> dict:
    """Sample until EOS or a token budget. Never pretend truncation is valid syntax."""
    if max_tokens <= 0:
        raise ValueError("Token budget must be positive")
    generator = torch.Generator().manual_seed(seed)
    tokens = ["<BOS>"]
    reason = "token_budget"
    for _ in range(max_tokens):  # Bound: remaining token budget, even if EOS never appears.
        logits = model(torch.tensor([[TOKEN_ID[t] for t in tokens]]))[0, -1]
        allowed = allowed_next(tokens)
        if hard_mask:
            mask = torch.tensor([t in allowed for t in VOCAB])
            if not mask.any():
                reason = "no_valid_next"
                break
            logits = logits.masked_fill(~mask, -torch.inf)
        token = VOCAB[torch.multinomial(logits.softmax(-1), 1, generator=generator).item()]
        tokens.append(token)
        if token not in allowed:
            reason = "invalid_transition"
            break
        if token == "<EOS>":
            reason = "eos"
            break
    complete = is_complete(tokens)
    return {
        "tokens": tokens,
        "complete": complete,
        "stop_reason": reason,
        "code": to_code(tokens) if complete else None,
    }


def generation_metrics(model: TinyLM, *, hard_mask: bool, count: int = 100) -> dict:
    """Report free-running complete-program rate, separately from teacher-forced tests."""
    if count <= 0:
        raise ValueError("Sample count must be positive")
    results = [generate(model, hard_mask=hard_mask, seed=i) for i in range(count)]
    return {
        "samples": count,
        "complete_program_rate": sum(r["complete"] for r in results) / count,
        "stop_counts": {
            reason: sum(r["stop_reason"] == reason for r in results)
            for reason in ("eos", "invalid_transition", "token_budget", "no_valid_next")
        },
        "sample": next((r for r in results if r["complete"]), results[0]),
    }


def run_training(args: argparse.Namespace) -> dict:
    """Generate or load data. Train paired models from identical weights and batch order."""
    torch.set_num_threads(1)
    torch.manual_seed(args.seed)
    rows = load_dataset(args.data) if args.data else make_dataset(args.seed)
    args.output_dir.mkdir(parents=True, exist_ok=True)
    write_dataset(args.output_dir / "programs.jsonl", rows)
    train_rows = [r for r in rows if r["split"] == "train"]
    test_rows = [r for r in rows if r["split"] == "test"]
    initial = TinyLM()
    report = {
        "seed": args.seed,
        "epochs": args.epochs,
        "rule_weight": args.rule_weight,
        "parameter_count": sum(p.numel() for p in initial.parameters()),
        "split_counts": {
            s: sum(r["split"] == s for r in rows) for s in ("train", "validation", "test")
        },
        "untrained": metrics(initial, test_rows),
        "models": {},
    }
    for name, weight in (("data_only", 0.0), ("rule_guided", args.rule_weight)):
        model = train_model(
            initial, train_rows, epochs=args.epochs, rule_weight=weight, seed=args.seed + 1
        )
        report["models"][name] = {
            "teacher_forced": metrics(model, test_rows),
            "free_running": generation_metrics(model, hard_mask=False),
            "hard_masked": generation_metrics(model, hard_mask=True),
        }
        torch.save(
            {"format": 1, "vocabulary": VOCAB, "state_dict": model.state_dict()},
            args.output_dir / f"{name}.pt",
        )
    (args.output_dir / "results.json").write_text(json.dumps(report, indent=2) + "\n")
    return report


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="command", required=True)
    data = sub.add_parser("data")
    data.add_argument("--output", type=Path, default=Path("programs.jsonl"))
    data.add_argument("--seed", type=int, default=7)
    train = sub.add_parser("train")
    train.add_argument("--data", type=Path)
    train.add_argument("--output-dir", type=Path, default=Path("runs/tokens"))
    train.add_argument("--epochs", type=int, default=80)
    train.add_argument("--rule-weight", type=float, default=4.0)
    train.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if args.command == "data":
        write_dataset(args.output, make_dataset(args.seed))
        print(f"Wrote 192 programs to {args.output}")
    else:
        print(json.dumps(run_training(args), indent=2))


if __name__ == "__main__":
    main()