#!/usr/bin/env python3
"""Train small manual transformers on deterministic generated CPU tasks.

Examples:

    python examples/pytorch/build_transformer.py seq2seq --task reverse --quick
    python examples/pytorch/build_transformer.py decoder-only --steps 100 --quick

The first task maps a source sequence to either a copy or a reversal.  The
second task predicts a repeating ascending token pattern.  Both use separate
held-out data and write a checkpoint only when ``--checkpoint`` is supplied.
"""

from __future__ import annotations

import argparse
from dataclasses import dataclass
from pathlib import Path

import torch
from torch import nn

from transformer_parts import DecoderOnlyTransformer, Seq2SeqTransformer


PAD_ID = 0
BOS_ID = 1
EOS_ID = 2
FIRST_DATA_ID = 3


@dataclass(frozen=True)
class SequenceBatch:
    """Padded IDs; ``targets`` use PAD_ID where loss must be ignored."""

    inputs: torch.Tensor
    decoder_inputs: torch.Tensor
    targets: torch.Tensor


def make_seq2seq_dataset(
    count: int,
    *,
    seed: int,
    task: str = "reverse",
    variable_length: bool = True,
    min_length: int = 3,
    max_length: int = 7,
    vocab_size: int = 16,
) -> SequenceBatch:
    """Create padded copy/reverse examples with independent deterministic RNG."""
    if task not in {"copy", "reverse"}:
        raise ValueError("task must be 'copy' or 'reverse'")
    if not (1 <= min_length <= max_length):
        raise ValueError("expected 1 <= min_length <= max_length")
    if vocab_size <= FIRST_DATA_ID:
        raise ValueError("vocab_size needs at least one ordinary data token")

    generator = torch.Generator(device="cpu").manual_seed(seed)
    source = torch.full((count, max_length), PAD_ID, dtype=torch.long)
    decoder_input = torch.full((count, max_length + 1), PAD_ID, dtype=torch.long)
    target = torch.full((count, max_length + 1), PAD_ID, dtype=torch.long)
    for row in range(count):
        if variable_length:
            length = int(
                torch.randint(min_length, max_length + 1, (), generator=generator).item()
            )
        else:
            length = max_length
        values = torch.randint(
            FIRST_DATA_ID, vocab_size, (length,), generator=generator, dtype=torch.long
        )
        answer = values if task == "copy" else values.flip(0)
        source[row, :length] = values
        decoder_input[row, 0] = BOS_ID
        decoder_input[row, 1 : length + 1] = answer
        target[row, :length] = answer
        target[row, length] = EOS_ID
    return SequenceBatch(source, decoder_input, target)


def make_pattern_dataset(
    count: int,
    *,
    seed: int,
    pattern_length: int = 8,
    vocab_size: int = 16,
) -> SequenceBatch:
    """Create cycles from independently sampled starts in a finite token set."""
    if pattern_length < 2:
        raise ValueError("pattern_length must be at least 2")
    data_token_count = vocab_size - FIRST_DATA_ID
    if data_token_count < 2:
        raise ValueError("vocab_size needs at least two ordinary data tokens")
    generator = torch.Generator(device="cpu").manual_seed(seed)
    starts = torch.randint(0, data_token_count, (count,), generator=generator)
    offsets = torch.arange(pattern_length)
    values = (starts[:, None] + offsets[None, :]) % data_token_count + FIRST_DATA_ID
    inputs = torch.cat(
        (torch.full((count, 1), BOS_ID, dtype=torch.long), values), dim=1
    )
    targets = torch.cat(
        (values, torch.full((count, 1), EOS_ID, dtype=torch.long)), dim=1
    )
    # SequenceBatch keeps one layout for both tasks. Decoder-only code uses
    # decoder_inputs and ignores inputs.
    return SequenceBatch(values, inputs, targets)


def masked_cross_entropy(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    """Average cross-entropy over non-padding targets only."""
    if logits.shape[:-1] != targets.shape:
        raise ValueError("logit and target sequence shapes do not match")
    if not bool(targets.ne(PAD_ID).any()):
        raise ValueError("cannot compute a loss for an all-padding target")
    return nn.functional.cross_entropy(
        logits.reshape(-1, logits.shape[-1]),
        targets.reshape(-1),
        ignore_index=PAD_ID,
    )


def token_and_sequence_accuracy(
    logits: torch.Tensor, targets: torch.Tensor
) -> tuple[float, float]:
    """Return masked token accuracy and whole-sequence accuracy."""
    predicted = logits.argmax(dim=-1)
    valid = targets.ne(PAD_ID)
    token_accuracy = (predicted.eq(targets) & valid).sum().item() / valid.sum().item()
    correct_or_padding = predicted.eq(targets) | ~valid
    sequence_accuracy = correct_or_padding.all(dim=1).float().mean().item()
    return token_accuracy, sequence_accuracy


def most_common_baseline(train_targets: torch.Tensor, heldout_targets: torch.Tensor) -> float:
    """Accuracy when every non-padding position gets one frequent train token."""
    valid_train = train_targets[train_targets.ne(PAD_ID)]
    most_common = int(torch.bincount(valid_train).argmax().item())
    valid_heldout = heldout_targets.ne(PAD_ID)
    return (
        heldout_targets.eq(most_common).logical_and(valid_heldout).sum().item()
        / valid_heldout.sum().item()
    )


def minibatch_indices(
    count: int, batch_size: int, steps: int, *, seed: int
) -> list[torch.Tensor]:
    generator = torch.Generator(device="cpu").manual_seed(seed)
    return [torch.randint(count, (batch_size,), generator=generator) for _ in range(steps)]


def _generated_matches(generated: torch.Tensor, targets: torch.Tensor) -> float:
    """Compare greedy sequences after removing BOS and trailing padding."""
    matches = 0
    for generated_row, target_row in zip(generated.tolist(), targets.tolist()):
        generated_answer = generated_row[1:]
        expected = [token for token in target_row if token != PAD_ID]
        if EOS_ID in generated_answer:
            generated_answer = generated_answer[: generated_answer.index(EOS_ID) + 1]
        if generated_answer == expected:
            matches += 1
    return matches / targets.shape[0]


def _save_and_verify(
    model: nn.Module,
    checkpoint_path: Path,
    config: dict[str, int | float],
    example_ids: torch.Tensor,
    *,
    decoder_only: bool,
    source_ids: torch.Tensor | None = None,
) -> None:
    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"model_state_dict": model.state_dict(), "config": config}, checkpoint_path)
    model_class = DecoderOnlyTransformer if decoder_only else Seq2SeqTransformer
    restored = model_class(**config)
    payload = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    restored.load_state_dict(payload["model_state_dict"])
    restored.eval()
    model.eval()
    with torch.no_grad():
        if decoder_only:
            original = model(example_ids)
            reloaded = restored(example_ids)
        else:
            if source_ids is None:
                raise ValueError("source_ids are required for seq2seq verification")
            original = model(source_ids, example_ids)
            reloaded = restored(source_ids, example_ids)
    if not torch.equal(original, reloaded):
        raise AssertionError("checkpoint reload changed model output")


def run_seq2seq(
    *,
    steps: int,
    task: str,
    variable_length: bool,
    quick: bool,
    checkpoint_path: Path | None = None,
) -> dict[str, float]:
    """Train and evaluate an encoder-decoder transformer."""
    torch.manual_seed(2026)
    vocab_size = 16
    # A broad generated training set discourages memorizing a few sequences.
    train_count, heldout_count = (4_096, 128) if quick else (16_384, 512)
    batch_size = 32 if quick else 64
    config: dict[str, int | float] = {
        "vocab_size": vocab_size,
        "d_model": 32 if quick else 64,
        "num_heads": 4,
        "hidden_dim": 64 if quick else 128,
        "num_layers": 1 if quick else 2,
        "dropout": 0.0,
        "padding_id": PAD_ID,
    }
    train = make_seq2seq_dataset(
        train_count,
        seed=10,
        task=task,
        variable_length=variable_length,
        vocab_size=vocab_size,
    )
    heldout = make_seq2seq_dataset(
        heldout_count,
        seed=20,
        task=task,
        variable_length=variable_length,
        vocab_size=vocab_size,
    )
    model = Seq2SeqTransformer(**config)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
    losses: list[float] = []
    model.train()
    for indices in minibatch_indices(train_count, batch_size, steps, seed=30):
        optimizer.zero_grad(set_to_none=True)
        logits = model(train.inputs[indices], train.decoder_inputs[indices])
        loss = masked_cross_entropy(logits, train.targets[indices])
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        losses.append(loss.item())

    model.eval()
    with torch.no_grad():
        heldout_logits = model(heldout.inputs, heldout.decoder_inputs)
        heldout_loss = masked_cross_entropy(heldout_logits, heldout.targets).item()
        token_accuracy, teacher_forced_exact = token_and_sequence_accuracy(
            heldout_logits, heldout.targets
        )
        evaluation_count = min(64, heldout_count)
        generated = model.generate(
            heldout.inputs[:evaluation_count],
            bos_id=BOS_ID,
            eos_id=EOS_ID,
            max_new_tokens=heldout.targets.shape[1],
        )
        greedy_exact = _generated_matches(generated, heldout.targets[:evaluation_count])
    baseline = most_common_baseline(train.targets, heldout.targets)

    print(f"task: seq2seq {task}, variable_length={variable_length}")
    print(f"train source shape: {tuple(train.inputs.shape)}")
    print(f"held-out source shape: {tuple(heldout.inputs.shape)}")
    print(f"first/final training loss: {losses[0]:.4f} / {losses[-1]:.4f}")
    print(f"held-out loss: {heldout_loss:.4f}")
    print(f"most-common-token baseline accuracy: {baseline:.3f}")
    print(f"held-out token accuracy (teacher forced): {token_accuracy:.3f}")
    print(f"held-out exact accuracy (teacher forced): {teacher_forced_exact:.3f}")
    print(f"held-out exact accuracy (greedy decode): {greedy_exact:.3f}")

    if checkpoint_path is not None:
        _save_and_verify(
            model,
            checkpoint_path,
            config,
            heldout.decoder_inputs[:4],
            decoder_only=False,
            source_ids=heldout.inputs[:4],
        )
        print(f"checkpoint saved and verified: {checkpoint_path}")

    return {
        "initial_loss": losses[0],
        "final_loss": losses[-1],
        "heldout_loss": heldout_loss,
        "baseline_accuracy": baseline,
        "token_accuracy": token_accuracy,
        "teacher_forced_exact_accuracy": teacher_forced_exact,
        "greedy_exact_accuracy": greedy_exact,
    }


def run_decoder_only(
    *,
    steps: int,
    quick: bool,
    checkpoint_path: Path | None = None,
) -> dict[str, float]:
    """Train and evaluate a decoder-only transformer on a cyclic pattern."""
    torch.manual_seed(2027)
    vocab_size = 16
    train_count, heldout_count = (256, 96) if quick else (1_024, 256)
    batch_size = 32 if quick else 64
    config: dict[str, int | float] = {
        "vocab_size": vocab_size,
        "d_model": 32 if quick else 64,
        "num_heads": 4,
        "hidden_dim": 64 if quick else 128,
        "num_layers": 1 if quick else 2,
        "dropout": 0.0,
        "padding_id": PAD_ID,
    }
    train = make_pattern_dataset(train_count, seed=40, vocab_size=vocab_size)
    heldout = make_pattern_dataset(heldout_count, seed=50, vocab_size=vocab_size)
    model = DecoderOnlyTransformer(**config)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
    losses: list[float] = []
    model.train()
    for indices in minibatch_indices(train_count, batch_size, steps, seed=60):
        optimizer.zero_grad(set_to_none=True)
        logits = model(train.decoder_inputs[indices])
        loss = masked_cross_entropy(logits, train.targets[indices])
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        losses.append(loss.item())

    model.eval()
    with torch.no_grad():
        heldout_logits = model(heldout.decoder_inputs)
        heldout_loss = masked_cross_entropy(heldout_logits, heldout.targets).item()
        token_accuracy, exact_accuracy = token_and_sequence_accuracy(
            heldout_logits, heldout.targets
        )
        # Prefix BOS + first data token. The cache handles all later positions.
        evaluation_count = min(64, heldout_count)
        prefix = heldout.decoder_inputs[:evaluation_count, :2]
        generated = model.generate_cached(prefix, heldout.targets.shape[1] - 1)
        greedy_exact = _generated_matches(generated, heldout.targets[:evaluation_count])
    baseline = most_common_baseline(train.targets, heldout.targets)

    print("task: decoder-only ascending cyclic pattern")
    print(
        "dataset note: train and held-out rows are sampled independently, but "
        f"there are only {vocab_size - FIRST_DATA_ID} starts, so patterns repeat "
        "across both sets"
    )
    print(
        "generation note: the first target is random, so cached greedy decoding "
        "is given the first data token"
    )
    print(f"train input shape: {tuple(train.decoder_inputs.shape)}")
    print(f"held-out input shape: {tuple(heldout.decoder_inputs.shape)}")
    print(f"first/final training loss: {losses[0]:.4f} / {losses[-1]:.4f}")
    print(f"held-out loss: {heldout_loss:.4f}")
    print(f"most-common-token baseline accuracy: {baseline:.3f}")
    print(f"held-out token accuracy: {token_accuracy:.3f}")
    print(f"held-out exact accuracy (teacher forced): {exact_accuracy:.3f}")
    print(
        "held-out exact accuracy (cached greedy decode, given first data token): "
        f"{greedy_exact:.3f}"
    )

    if checkpoint_path is not None:
        _save_and_verify(
            model,
            checkpoint_path,
            config,
            heldout.decoder_inputs[:4],
            decoder_only=True,
        )
        print(f"checkpoint saved and verified: {checkpoint_path}")

    return {
        "initial_loss": losses[0],
        "final_loss": losses[-1],
        "heldout_loss": heldout_loss,
        "baseline_accuracy": baseline,
        "token_accuracy": token_accuracy,
        "teacher_forced_exact_accuracy": exact_accuracy,
        "greedy_exact_accuracy": greedy_exact,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="build", required=True)

    seq2seq = subparsers.add_parser("seq2seq", help="train an encoder-decoder model")
    seq2seq.add_argument("--task", choices=("copy", "reverse"), default="reverse")
    seq2seq.add_argument(
        "--fixed-length",
        action="store_true",
        help="use length 7 only (the default includes lengths 3 through 7)",
    )
    seq2seq.add_argument("--steps", type=int)
    seq2seq.add_argument("--quick", action="store_true")
    seq2seq.add_argument("--checkpoint", type=Path)

    decoder = subparsers.add_parser("decoder-only", help="train a causal language model")
    decoder.add_argument("--steps", type=int)
    decoder.add_argument("--quick", action="store_true")
    decoder.add_argument("--checkpoint", type=Path)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    torch.set_num_threads(1)
    steps = args.steps if args.steps is not None else (40 if args.quick else 500)
    if steps < 1:
        raise ValueError("--steps must be at least 1")
    if args.build == "seq2seq":
        run_seq2seq(
            steps=steps,
            task=args.task,
            variable_length=not args.fixed_length,
            quick=args.quick,
            checkpoint_path=args.checkpoint,
        )
    else:
        run_decoder_only(
            steps=steps,
            quick=args.quick,
            checkpoint_path=args.checkpoint,
        )


if __name__ == "__main__":
    main()