"""Training loop. Usage: python -m src.train [model] [seed] [--flag ...]"""
import csv
import json
import math
import os
import random
import sys

import numpy as np
import torch
import torch.nn as nn

from src.config import Config, parse_args
from src.data import build_examples, decode_tokens, get_splits, make_batch
from src.model_api import build_model


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


def resolve_device(device_str: str) -> torch.device:
    if device_str == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(device_str)


def make_gpu_dataset(examples: list, cfg: Config, device: torch.device, max_cache_size: int = 10_000) -> dict[str, torch.Tensor | int | bool]:
    """Create static tensors in VRAM for small/medium datasets. For large datasets (>max_cache_size),
    signals that chunked CPU streaming should be used to prevent VRAM OOM."""
    from src.data import _global_lengths
    in_max, out_max = _global_lengths(cfg)
    if not examples:
        return {
            "x": torch.empty((0, in_max), dtype=torch.long, device=device),
            "y": torch.empty((0, out_max), dtype=torch.long, device=device),
            "y_in": torch.empty((0, out_max), dtype=torch.long, device=device),
            "y_mask": torch.empty((0, out_max), dtype=torch.bool, device=device),
            "loss_mask": torch.empty((0, out_max), dtype=torch.bool, device=device),
            "size": 0,
            "is_cached": True,
        }
    if len(examples) > max_cache_size:
        return {"size": len(examples), "is_cached": False}

    batch = make_batch(examples, cfg)
    return {
        "x": batch["x"].to(device, non_blocking=True),
        "y": batch["y"].to(device, non_blocking=True),
        "y_in": batch["y_in"].to(device, non_blocking=True),
        "y_mask": batch["y_mask"].to(device, non_blocking=True),
        "loss_mask": batch["loss_mask"].to(device, non_blocking=True),
        "size": len(examples),
        "is_cached": True,
    }


@torch.no_grad()
def evaluate_gpu(model, gpu_ds: dict[str, torch.Tensor | int | bool], examples: list, cfg: Config, device: torch.device):
    """Batched evaluation: operates directly on VRAM tensors if cached, or streams in eval_batch_size chunks."""
    model.eval()
    N = gpu_ds["size"]
    if N == 0:
        return 0.0, 0.0, []

    eval_bs = getattr(cfg, "eval_batch_size", 512)
    use_amp = getattr(cfg, "use_amp", True) and device.type == "cuda"
    token_correct = 0
    token_total = 0
    em_correct = 0
    per_example = []
    is_cached = gpu_ds.get("is_cached", False)

    for bi in range(0, N, eval_bs):
        chunk_examples = examples[bi: bi + eval_bs]
        if is_cached:
            x = gpu_ds["x"][bi: bi + eval_bs]
            y = gpu_ds["y"][bi: bi + eval_bs]
            y_in_eval = gpu_ds["y_in"][bi: bi + eval_bs]
            mask = gpu_ds["loss_mask"][bi: bi + eval_bs]
        else:
            batch = make_batch(chunk_examples, cfg)
            x = batch["x"].to(device)
            y = batch["y"].to(device)
            y_in_eval = batch["y_in"].to(device)
            mask = batch["loss_mask"].to(device)

        with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
            out = model(x, y_in_eval)
            logits = out["logits"]
        pred_tok = logits.argmax(-1)
        token_correct += int((pred_tok[mask] == y[mask]).sum().item())
        token_total += int(mask.sum().item())

        B = x.shape[0]
        from src.data import _global_lengths
        _, out_max = _global_lengths(cfg)
        max_len = max(getattr(cfg, "max_out_len", 6), out_max)

        y_in = torch.full((B, 1), cfg.eos_id, dtype=torch.long, device=device)
        with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
            for _ in range(max_len):
                o = model(x, y_in)
                nxt = o["logits"][:, -1].argmax(-1)
                y_in = torch.cat([y_in, nxt[:, None]], dim=1)
        gen = y_in[:, 1:].cpu()

        for i, ex in enumerate(chunk_examples):
            xrow, target = ex[0], ex[1]
            pred = decode_tokens(gen[i].tolist(), cfg)
            target_n = decode_tokens(target, cfg)
            ok = (pred == target_n) and (pred != -1)
            em_correct += int(ok)
            per_example.append((tuple(xrow), target_n, pred, ok))

    model.train()
    return token_correct / max(1, token_total), em_correct / N, per_example


@torch.no_grad()
def evaluate(model, examples, cfg: Config, device: torch.device | None = None):
    """Token accuracy + exact-match accuracy + per-example detail (backward compatible wrapper)."""
    if device is None:
        try:
            device = next(model.parameters()).device
        except StopIteration:
            device = torch.device("cpu")
    gpu_ds = make_gpu_dataset(examples, cfg, device)
    return evaluate_gpu(model, gpu_ds, examples, cfg, device)


def log_example_rows(log_examples, val_per_example):
    """Format the fixed 10 logged val inputs as '42:ok' strings, in fixed order."""
    by_x = {x: (target, pred, ok) for x, target, pred, ok in val_per_example}
    rows = []
    for x, _target in log_examples:
        t, p, ok = by_x.get(tuple(x), (-1, -1, False))
        label = "".join(map(str, x))  # "42" in both vocab modes
        rows.append(f"{label}:{'ok' if ok else f'{p}~{t}'}")
    return ";".join(rows)


def main() -> None:
    cfg = parse_args()
    set_seed(cfg.seed)
    device = resolve_device(cfg.device)
    use_amp = cfg.use_amp and device.type == "cuda"
    out_dir = os.path.join(cfg.out_dir, cfg.model, f"seed{cfg.seed}")
    csv_path = os.path.join(out_dir, "metrics.csv")
    if os.path.exists(csv_path):
        raise SystemExit(f"REFUSING to rerun in place: {csv_path} exists. Use a fresh --out_dir "
                         f"(reruns would corrupt the CSV and checkpoint provenance).")
    os.makedirs(out_dir, exist_ok=True)
    cfg.save(os.path.join(out_dir, "config.json"))
    meta = {
        "python": sys.version.split()[0],
        "torch": torch.__version__,
        "numpy": np.__version__,
        "device": str(device),
        "use_amp": use_amp,
        "compile_model": cfg.compile_model,
        "torch_threads": torch.get_num_threads(),
        "cmd": sys.argv,
    }
    if device.type == "cuda":
        meta["gpu_name"] = torch.cuda.get_device_name(device)
        meta["gpu_capability"] = list(torch.cuda.get_device_capability(device))
    with open(os.path.join(out_dir, "run_meta.json"), "w") as fh:
        json.dump(meta, fh, indent=2)

    train_in, val_in = get_splits(cfg)
    train_ex = build_examples(train_in, cfg)
    val_ex = build_examples(val_in, cfg)
    train_gpu = make_gpu_dataset(train_ex, cfg, device)
    val_gpu = make_gpu_dataset(val_ex, cfg, device)

    model = build_model(cfg).to(device)
    if cfg.compile_model and hasattr(torch, "compile") and device.type == "cuda":
        try:
            model = torch.compile(model, mode="reduce-overhead")
        except Exception as e:
            print(f"Warning: torch.compile failed ({e}), falling back to uncompiled model")

    print(f"model={cfg.model} params={model.param_count()} device={device} amp={use_amp} "
          f"train={len(train_ex)} val={len(val_ex)} "
          f"vocab={cfg.vocab} eos={cfg.eos_id} pad={cfg.pad_id}")

    opt_kwargs = {"lr": cfg.lr, "weight_decay": cfg.weight_decay}
    if device.type == "cuda":
        try:
            opt = torch.optim.AdamW(model.parameters(), fused=True, **opt_kwargs)
        except Exception:
            opt = torch.optim.AdamW(model.parameters(), **opt_kwargs)
    else:
        opt = torch.optim.AdamW(model.parameters(), **opt_kwargs)
    scaler = torch.amp.GradScaler("cuda", enabled=use_amp)
    ce = nn.CrossEntropyLoss(reduction="none")

    log_examples = sorted(val_ex, key=lambda ex: (len(ex[0]), ex[0]))[: cfg.log_n_examples]

    fieldnames = ["step", "train_loss", "train_token_acc", "train_em", "val_token_acc",
                  "val_em", "mean_halt_steps", "log_examples", "param_count"]
    best_val_em = -1.0
    patience_left = cfg.early_stop_patience
    best_path = os.path.join(out_dir, "best.pt")

    N_train = train_gpu["size"]
    n_batches = math.ceil(N_train / cfg.batch_size) if N_train > 0 else 0
    rng = random.Random(cfg.seed)
    step = 0
    done = False
    first_row = True

    if N_train == 0:
        print("Warning: empty training split; skipping training.")
        torch.save(model.state_dict(), os.path.join(out_dir, "last.pt"))
        return

    def should_eval(cur_step: int) -> bool:
        if cur_step >= cfg.max_train_steps:
            return True
        if getattr(cfg, "eval_schedule", "fixed") in ("adaptive", "logarithmic"):
            if cur_step < 10_000:
                return cur_step % 200 == 0
            elif cur_step < 50_000:
                return cur_step % 1_000 == 0
            elif cur_step < 500_000:
                return cur_step % 5_000 == 0
            else:
                return cur_step % 10_000 == 0
        return cur_step % cfg.eval_every == 0

    def _eval_pass(cur_loss, halt_steps):
        nonlocal best_val_em, patience_left, done, first_row
        train_tok, train_em, _ = evaluate_gpu(model, train_gpu, train_ex, cfg, device)
        val_tok, val_em, val_per = evaluate_gpu(model, val_gpu, val_ex, cfg, device)
        mean_halt = float(halt_steps.mean()) if halt_steps is not None else float("nan")
        row = {
            "step": step, "train_loss": float(cur_loss), "train_token_acc": train_tok,
            "train_em": train_em, "val_token_acc": val_tok, "val_em": val_em,
            "mean_halt_steps": mean_halt,
            "log_examples": log_example_rows(log_examples, val_per),
            "param_count": model.param_count(),
        }
        with open(csv_path, "a", newline="") as fh:
            w = csv.DictWriter(fh, fieldnames=fieldnames)
            if first_row:
                w.writeheader()
                first_row = False
            w.writerow(row)
        if val_em > best_val_em:
            best_val_em = val_em
            torch.save(model.state_dict(), best_path)
        if val_em >= cfg.early_stop_em:
            patience_left -= 1
        else:
            patience_left = cfg.early_stop_patience
        if patience_left <= 0:
            done = True
        print(f"step {step} loss {float(cur_loss):.4f} train_em {train_em:.3f} "
              f"val_em {val_em:.3f} val_tok {val_tok:.3f} halt {mean_halt:.2f}")

    while step < cfg.max_train_steps and not done:
        if train_gpu.get("is_cached", False):
            perm = torch.randperm(N_train, device=device)
            for bi in range(n_batches):
                idx = perm[bi * cfg.batch_size: (bi + 1) * cfg.batch_size]
                x = train_gpu["x"][idx]
                y = train_gpu["y"][idx]
                y_in = train_gpu["y_in"][idx]
                mask = train_gpu["loss_mask"][idx]

                with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
                    out = model(x, y_in)
                    logits = out["logits"]
                    y_safe = y.clamp(max=cfg.vocab - 1)   # CE index guard for pad positions
                    loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape(
                        logits.shape[0], -1) * mask.float()
                    loss_tokens = loss_tokens.sum() / mask.sum().clamp(min=1)
                    if cfg.halting and step >= cfg.halt_warmup_steps:
                        if step >= cfg.halt_ramp_end_steps:
                            lam = cfg.halt_penalty
                        else:
                            frac = (step - cfg.halt_warmup_steps) / max(1, cfg.halt_ramp_end_steps - cfg.halt_warmup_steps)
                            lam = cfg.halt_penalty * frac
                    else:
                        lam = 0.0
                    halt = out["halt_steps"]
                    penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0
                    loss = loss_tokens + penalty

                if cfg.lr_decay:
                    # cosine 1e-3 -> 1e-4 over the full run (E3, locked in Addendum 4)
                    frac = min(1.0, step / cfg.max_train_steps)
                    lr_t = cfg.lr * 0.1 + 0.5 * (cfg.lr - cfg.lr * 0.1) * (1 + math.cos(math.pi * frac))
                    opt.param_groups[0]["lr"] = lr_t

                opt.zero_grad()
                scaler.scale(loss).backward()
                scaler.step(opt)
                scaler.update()

                step += 1
                if should_eval(step):
                    _eval_pass(loss.detach(), halt.detach() if halt is not None else None)
                    if done:
                        break
                if step >= cfg.max_train_steps:
                    break
        else:
            rng.shuffle(train_ex)
            for bi in range(n_batches):
                sl = train_ex[bi * cfg.batch_size: (bi + 1) * cfg.batch_size]
                if not sl:
                    continue
                batch = make_batch(sl, cfg)
                x = batch["x"].to(device)
                y = batch["y"].to(device)
                y_in = batch["y_in"].to(device)
                mask = batch["loss_mask"].to(device)

                with torch.amp.autocast(device_type="cuda", dtype=torch.float16, enabled=use_amp):
                    out = model(x, y_in)
                    logits = out["logits"]
                    y_safe = y.clamp(max=cfg.vocab - 1)   # CE index guard for pad positions
                    loss_tokens = ce(logits.reshape(-1, cfg.vocab), y_safe.reshape(-1)).reshape(
                        logits.shape[0], -1) * mask.float()
                    loss_tokens = loss_tokens.sum() / mask.sum().clamp(min=1)
                    if cfg.halting and step >= cfg.halt_warmup_steps:
                        if step >= cfg.halt_ramp_end_steps:
                            lam = cfg.halt_penalty
                        else:
                            frac = (step - cfg.halt_warmup_steps) / max(1, cfg.halt_ramp_end_steps - cfg.halt_warmup_steps)
                            lam = cfg.halt_penalty * frac
                    else:
                        lam = 0.0
                    halt = out["halt_steps"]
                    penalty = halt.float().mean() * lam if halt is not None and lam > 0 else 0.0
                    loss = loss_tokens + penalty

                if cfg.lr_decay:
                    # cosine 1e-3 -> 1e-4 over the full run (E3, locked in Addendum 4)
                    frac = min(1.0, step / cfg.max_train_steps)
                    lr_t = cfg.lr * 0.1 + 0.5 * (cfg.lr - cfg.lr * 0.1) * (1 + math.cos(math.pi * frac))
                    opt.param_groups[0]["lr"] = lr_t

                opt.zero_grad()
                scaler.scale(loss).backward()
                scaler.step(opt)
                scaler.update()

                step += 1
                if should_eval(step):
                    _eval_pass(loss.detach(), halt.detach() if halt is not None else None)
                    if done:
                        break
                if step >= cfg.max_train_steps:
                    break

    torch.save(model.state_dict(), os.path.join(out_dir, "last.pt"))
    print(f"DONE steps={step} best_val_em={best_val_em:.4f}")


if __name__ == "__main__":
    main()