# -*- coding: utf-8 -*-
"""Train ShotNet trên bb9_synth 150k (BG27 bước 2) — venv CV, GPU RTX 3070.

Kỷ luật BRIEF 27:
  - Val ~2% cắt từ TRAIN theo seed (``val_split_indices``); 5k held-out là
    TEST — script này KHÔNG BAO GIỜ đọc shard heldout.
  - Ngân sách tune: tối đa 3 config chạy trọn (bảng ``CONFIGS`` — mỗi config
    một dòng giả thuyết). Checkpoint chọn theo VAL, không theo held-out.
    BG29 (12/08) mở thêm ĐÚNG 2 config c4a/c4b — cùng kiến trúc/optimizer/
    ngân sách bước với c3, chỉ đổi PHÂN PHỐI train (dropout slot bi mục tiêu
    + reweight lát target). Muốn config thứ ba là phải hỏi Cowork.
  - Log curve CSV + PNG mỗi epoch: loss từng head + 4 chỉ số gate đo trên
    VAL (dphi_med, v0_relerr_med raw, side_acc, vert_acc).
  - Checkpoint tốt nhất theo ``val_score`` vào ``models/<run_name>/`` + spec
    đầy đủ (config, seed, commit, dataset) — nếp ``qfield_20260721``.

``val_score`` = dphi_med/2 + v0_relerr_med/0.10 + (1−side_acc)/0.20 +
(1−vert_acc)/0.20 — mỗi vế = 1.0 khi đúng bar tuyệt đối của gate G-27.3,
nhỏ hơn là tốt hơn; chọn checkpoint thẳng theo thước gate thay vì theo loss
(loss trộn trọng số tay, không phải thước nghiệm thu).

Chạy (venv CV):
  D:\\Khoa luan\\poolcoach-cv-env\\Scripts\\python.exe ^
      poolcoach-rl\\scripts\\broadcast\\train_shotnet.py --config c1
  # smoke pipeline (KHÔNG tính vào ngân sách 3 config):
  ... train_shotnet.py --config c1 --smoke --out-root <scratchpad>

Bẫy quen: đọc shard lần đầu trên ổ D lạnh sẽ ì (npz 1.7GB) — đừng tưởng
treo rồi kill (bài BG26b Bất ngờ 1); console ASCII (cp1252).
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import subprocess
import sys
import time
from dataclasses import asdict, dataclass, field, replace
from pathlib import Path

import numpy as np
import torch

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))

from poolcoach_cv.shotnet import (  # noqa: E402
    BucketBatcher, ShardDataset, ShotNet, ShotNetConfig, SlotDropoutConfig,
    TargetReweight, collate, gate_metrics, shotnet_loss, val_split_indices)

DATA_DIR_DEFAULT = Path(r"D:\Khoa luan\datasets\bb9_synth")


# ------------------------------------------------------------ config

@dataclass
class TrainConfig:
    name: str
    hypothesis: str                      # MỘT dòng vì sao chạy config này
    model: ShotNetConfig = field(default_factory=ShotNetConfig)
    epochs: int = 40
    lr: float = 3e-4
    weight_decay: float = 0.01
    warmup_frac: float = 0.03
    grad_clip: float = 1.0
    token_budget: int = 40_000           # max_len_batch × cỡ_batch (VRAM 8GB)
    max_batch: int = 256
    val_frac: float = 0.02
    split_seed: int = 20260811           # cùng ngày chốt spec dataset
    torch_seed: int = 0
    amp_bf16: bool = True                # Ampere: bf16 autocast, khỏi scaler
    run_date: str = "20260811"           # tiền tố thư mục run (giữ tên cũ
                                         # của c1/c2/c3 tái lập được)
    # BG29 — CHỈ tác động đường train, val vẫn sạch như c3 (set_train_mode)
    slot_dropout: SlotDropoutConfig | None = None
    target_reweight: TargetReweight | None = None


CONFIGS: dict[str, TrainConfig] = {
    "c1": TrainConfig(
        name="c1",
        hypothesis="transformer nho d128x4L (~0.8M params) + loss mac dinh "
                   "du cho bai toan nguoc tren quy dao sach-vua-phai; "
                   "aux ball OFF de lay moc don gian nhat truoc",
    ),
    # c1 do duoc (val, best ep37): dphi 11.7 / v0rel 0.136 / side 0.526 /
    # vert 0.233 — spin KHONG hoc duoc gi (a_pred sap ve hang: std 0.02,
    # corr 0.02), dphi xau dan theo v0 (4.8 cham -> 18.0 nhanh) va xau hon
    # tren identifiable (12.3 vs 5.6 non-ident). c2 = 3 sua co gia thuyet:
    "c2": TrainConfig(
        name="c2",
        hypothesis="(i) them delta-xy per-slot: phi/V0 la ham VAN TOC, c1 "
                   "bat encoder tu dung chuyen dong tu toa do tuyet doi va "
                   "tran o 11.7 deg; (ii) (a,b)/0.4 + delta 0.5 + w_ab 2.0: "
                   "gradient spin c1 chi ~0.05, head sap ve hang - can "
                   "ngang phi de encoder chiu hoc tin hieu sau va cham; "
                   "(iii) aux v0_ball/phi_ball ON (label san, BG26b BN4): "
                   "hoc trang thai bi QUAN SAT DUOC roi suy nguoc ve gay "
                   "(phi_gay = phi_bi + squirt(a)). Co model GIU NGUYEN "
                   "de quy duoc hieu ung; c3 de danh cho scale-up",
        model=ShotNetConfig(use_deltas=True, ab_scale=0.4, huber_ab=0.5,
                            w_ab=2.0, aux_ball=True),
    ),
    # c2 do duoc (val, best ep32): dphi 9.36 / v0rel 0.1275 / side 0.596 /
    # vert 0.290 — ca 3 sua deu can (spin bat dau hoc, side/vert VAN CON
    # LEO tai ep39, train loss con xuong => chua hoi tu). c3 (config cuoi
    # ngan sach) = giu nguyen c2 + 3 sua:
    "c3": TrainConfig(
        name="c3",
        hypothesis="(iv) phi loss vec-MSE co gradient ~sin(delta) nen cu "
                   "vo vong (che dau + cham som, ~90 deg) thong tri; gate "
                   "cham MEDIAN -> Huber tren GOC (delta 5 deg, w 6.0) ep "
                   "khoi giua phan phoi ve duoi bar, ap cho ca aux "
                   "phi_ball; (v) huber_v0 0.25->0.10: cung logic median "
                   "(duoi mu dropout khong duoc keo gradient); (vi) "
                   "d192x6L + 60 epoch: c2 chua hoi tu (side/vert con leo "
                   "o ep39) va spin can dung luong bieu dien hinh hoc sau "
                   "va cham",
        model=ShotNetConfig(d_model=192, n_layers=6, n_heads=6, d_ff=768,
                            use_deltas=True, ab_scale=0.4, huber_ab=0.5,
                            w_ab=2.0, aux_ball=True,
                            phi_loss="ang_huber", phi_huber_deg=5.0,
                            w_phi=6.0, huber_v0=0.10),
        epochs=60,
    ),
    # ---- BG29 (12/08): c3 SẠCH trên synth nhưng sập trên clip thật vì mật
    # độ detection bi mục tiêu (cú 11 = 1.01 det/frame, cú 12 = 3.09, synth
    # 96.4% visibility ≈ 5.3 det/frame). Kiến trúc/optimizer/ngân sách bước
    # GIỮ NGUYÊN c3 — chỉ đổi phân phối train. Ngân sách chốt trước: 2 config.
    "c4a": TrainConfig(
        name="c4a",
        hypothesis="c3 + dropout slot bi muc tieu (keep ~U[0.10, 0.95] tren "
                   "75% cu, cue KHONG dropout): khe sim->real BG28 la PHAN "
                   "PHOI mat do detection chu khong phai kien truc hay input "
                   "path (slot-builder da duoc minh oan bang thi nghiem cach "
                   "ly 30 cu). Net phai hoc tra loi khi thua.",
        model=ShotNetConfig(d_model=192, n_layers=6, n_heads=6, d_ff=768,
                            use_deltas=True, ab_scale=0.4, huber_ab=0.5,
                            w_ab=2.0, aux_ball=True,
                            phi_loss="ang_huber", phi_huber_deg=5.0,
                            w_phi=6.0, huber_v0=0.10),
        epochs=60,
        run_date="20260812",
        slot_dropout=SlotDropoutConfig(),
    ),
    "c4b": TrainConfig(
        name="c4b",
        hypothesis="c4a + reweight loss ve lat target (cham dau >=0.3s, "
                   "13.62% train, w=3.0): thuoc chinh thuc cua vong nay la "
                   "lat >=0.3s, con 74.8% cu synth cham <0.2s la vung mu "
                   "thong tin dang keo gradient di cho khac.",
        model=ShotNetConfig(d_model=192, n_layers=6, n_heads=6, d_ff=768,
                            use_deltas=True, ab_scale=0.4, huber_ab=0.5,
                            w_ab=2.0, aux_ball=True,
                            phi_loss="ang_huber", phi_huber_deg=5.0,
                            w_phi=6.0, huber_v0=0.10),
        epochs=60,
        run_date="20260812",
        slot_dropout=SlotDropoutConfig(),
        target_reweight=TargetReweight(weight=3.0),
    ),
}


# ------------------------------------------------------------ helpers

def _git_commit() -> str:
    try:
        return subprocess.run(["git", "-C", str(ROOT), "rev-parse", "HEAD"],
                              capture_output=True, text=True,
                              check=True).stdout.strip()
    except Exception:
        return "unknown"


def _predictions(out: dict, cfg: ShotNetConfig) -> dict:
    """Chuyển output forward (đang có sẵn, khỏi forward lần hai) → numpy
    prediction cùng format ``ShotNet.predict``."""
    v0 = torch.exp(out["v0_z"]) if cfg.v0_log else out["v0_z"]
    phi = torch.rad2deg(torch.atan2(out["phi_vec"][:, 1],
                                    out["phi_vec"][:, 0])) % 360.0
    return {"v0": v0.float().cpu().numpy(),
            "phi_deg": phi.float().cpu().numpy(),
            "a": (out["ab"][:, 0] * cfg.ab_scale).float().cpu().numpy(),
            "b": (out["ab"][:, 1] * cfg.ab_scale).float().cpu().numpy(),
            "p_ident": torch.sigmoid(out["ident_logit"]).float().cpu().numpy()}


def val_score(m: dict) -> float:
    """Thước chọn checkpoint — 1.0 mỗi vế = đúng bar G-27.3 (docstring)."""
    return (m["dphi_med"] / 2.0 + m["v0_relerr_med"] / 0.10
            + (1.0 - m["side_acc"]) / 0.20 + (1.0 - m["vert_acc"]) / 0.20)


def _to_device(batch: dict, device: str) -> dict:
    return {k: v.to(device, non_blocking=True) for k, v in batch.items()}


LOSS_KEYS = ("total", "v0", "phi", "ab", "ident")


def run_epoch_train(model, ds, batches, cfg: TrainConfig, opt, sched,
                    device: str) -> dict:
    model.train()
    sums = {k: 0.0 for k in LOSS_KEYS}
    n = 0
    for idxs in batches:
        batch = _to_device(collate([ds[int(i)] for i in idxs]), device)
        opt.zero_grad(set_to_none=True)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16,
                            enabled=cfg.amp_bf16 and device == "cuda"):
            out = model(batch["x"], batch["t"], batch["mask"])
            loss = shotnet_loss(out, batch, cfg.model)
        loss["total"].backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.grad_clip)
        opt.step()
        sched.step()
        bs = len(idxs)
        n += bs
        for k in LOSS_KEYS:
            sums[k] += float(loss[k].detach()) * bs
    return {k: v / max(n, 1) for k, v in sums.items()}


@torch.no_grad()
def run_val(model, ds, batches, cfg: TrainConfig, device: str) -> dict:
    model.eval()
    sums = {k: 0.0 for k in LOSS_KEYS}
    n = 0
    preds, labels = [], []
    for idxs in batches:
        batch = _to_device(collate([ds[int(i)] for i in idxs]), device)
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16,
                            enabled=cfg.amp_bf16 and device == "cuda"):
            out = model(batch["x"], batch["t"], batch["mask"])
        out = {k: v.float() for k, v in out.items()}
        loss = shotnet_loss(out, batch, cfg.model)
        bs = len(idxs)
        n += bs
        for k in LOSS_KEYS:
            sums[k] += float(loss[k]) * bs
        preds.append(_predictions(out, cfg.model))
        labels.append({k: batch[k].cpu().numpy()
                       for k in ("label_v0", "label_phi", "label_a",
                                 "label_b", "identifiable")})
    pred = {k: np.concatenate([p[k] for p in preds]) for k in preds[0]}
    lab = {k: np.concatenate([q[k] for q in labels]) for k in labels[0]}
    m = gate_metrics(pred, lab)
    m.update({f"loss_{k}": v / max(n, 1) for k, v in sums.items()})
    m["score"] = val_score(m)
    return m


def draw_curves(csv_path: Path, png_path: Path) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    rows = list(csv.DictReader(open(csv_path, encoding="utf-8")))
    ep = [int(r["epoch"]) for r in rows]

    def col(name):
        return [float(r[name]) for r in rows]

    fig, axes = plt.subplots(2, 2, figsize=(12, 8))
    ax = axes[0, 0]
    ax.plot(ep, col("train_total"), label="train")
    ax.plot(ep, col("val_loss_total"), label="val")
    ax.set_title("loss total")
    ax.legend()
    ax = axes[0, 1]
    ax.semilogy(ep, col("val_dphi_med"), label="|dphi| med (deg)")
    ax.axhline(2.0, color="r", ls="--", lw=0.8, label="bar 2 deg")
    ax.set_title("val |dphi| median")
    ax.legend()
    ax = axes[1, 0]
    ax.semilogy(ep, col("val_v0_relerr_med"), label="V0 rel err med (raw)")
    ax.axhline(0.10, color="r", ls="--", lw=0.8, label="bar 10%")
    ax.set_title("val V0 rel err median")
    ax.legend()
    ax = axes[1, 1]
    ax.plot(ep, col("val_side_acc"), label="side acc")
    ax.plot(ep, col("val_vert_acc"), label="vert acc")
    ax.axhline(0.80, color="r", ls="--", lw=0.8, label="bar 80%")
    ax.set_ylim(0, 1.02)
    ax.set_title("val spin sign acc (identifiable)")
    ax.legend()
    for a in axes.flat:
        a.set_xlabel("epoch")
        a.grid(alpha=0.3)
    fig.tight_layout()
    fig.savefig(png_path, dpi=110)
    plt.close(fig)


# ------------------------------------------------------------ main

def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="c1", choices=sorted(CONFIGS))
    ap.add_argument("--data", default=str(DATA_DIR_DEFAULT))
    ap.add_argument("--out-root", default=str(ROOT / "models"))
    ap.add_argument("--run-name", default="")
    ap.add_argument("--epochs", type=int, default=0, help="override epochs")
    ap.add_argument("--smoke", action="store_true",
                    help="1 shard + 2 epoch — kiem pipeline, KHONG tinh "
                         "vao ngan sach 3 config")
    args = ap.parse_args()

    cfg = CONFIGS[args.config]
    if args.epochs:
        cfg = replace(cfg, epochs=args.epochs)
    if args.smoke:
        cfg = replace(cfg, epochs=min(cfg.epochs, 2))
    run_name = args.run_name or (
        f"shotnet_{cfg.run_date}_{cfg.name}"
        + ("_smoke" if args.smoke else ""))
    out_dir = Path(args.out_root) / run_name
    out_dir.mkdir(parents=True, exist_ok=True)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"run {run_name} | config {cfg.name} | device {device} "
          f"({torch.cuda.get_device_name(0) if device == 'cuda' else 'cpu'})")
    print(f"hypothesis: {cfg.hypothesis}")

    shards = sorted(Path(args.data).glob("train_*.npz"))
    if not shards:
        sys.exit("khong thay shard train_*.npz")
    if args.smoke:
        shards = shards[:1]
    print(f"nap {len(shards)} shard train vao RAM (lan dau o D lanh se i "
          f"vai phut - KHONG phai treo)...")
    t0 = time.time()
    ds = ShardDataset(shards, deltas=cfg.model.use_deltas,
                      augment=cfg.slot_dropout,
                      reweight=cfg.target_reweight)
    print(f"  {len(ds)} cu / {time.time() - t0:.0f}s | "
          f"frames p50 {int(np.median(ds.lengths))} max {ds.lengths.max()}")
    if cfg.slot_dropout:
        sd = cfg.slot_dropout
        print(f"augment slot dropout: p_apply {sd.p_apply} keep "
              f"U[{sd.keep_min}, {sd.keep_max}] seed {sd.seed} "
              f"(TRAIN only -- val/eval sach nhu c3)")
    if cfg.target_reweight:
        rw = cfg.target_reweight
        print(f"reweight lat target: w {rw.weight} cho cham dau >= "
              f"{rw.t_first_s}s (TRAIN only)")

    tr_idx, va_idx = val_split_indices(len(ds), cfg.val_frac, cfg.split_seed)
    print(f"split: train {len(tr_idx)} / val {len(va_idx)} "
          f"(seed {cfg.split_seed}, frac {cfg.val_frac})")
    tr_batcher = BucketBatcher(ds.lengths, tr_idx, cfg.token_budget,
                               cfg.max_batch, seed=cfg.split_seed)
    va_batches = BucketBatcher(ds.lengths, va_idx, cfg.token_budget,
                               cfg.max_batch,
                               seed=cfg.split_seed).epoch_batches(
                                   0, shuffle=False)

    torch.manual_seed(cfg.torch_seed)
    np.random.seed(cfg.torch_seed)
    model = ShotNet(cfg.model).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"model {n_params} params")

    opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr,
                            weight_decay=cfg.weight_decay)
    steps_per_epoch = len(tr_batcher.epoch_batches(0))
    total_steps = max(steps_per_epoch * cfg.epochs, 1)
    warmup = max(int(total_steps * cfg.warmup_frac), 1)

    def lr_lambda(step):
        if step < warmup:
            return (step + 1) / warmup
        p = (step - warmup) / max(total_steps - warmup, 1)
        return 0.5 * (1.0 + math.cos(math.pi * min(p, 1.0)))

    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)

    csv_path = out_dir / "curves.csv"
    png_path = out_dir / "curves.png"
    fields = (["epoch", "lr", "elapsed_s"]
              + [f"train_{k}" for k in LOSS_KEYS]
              + [f"val_loss_{k}" for k in LOSS_KEYS]
              + ["val_dphi_med", "val_v0_relerr_med", "val_side_acc",
                 "val_vert_acc", "val_ident_acc", "val_score", "is_best"])
    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        csv.DictWriter(f, fieldnames=fields).writeheader()

    best = {"score": float("inf"), "epoch": -1, "metrics": None}
    t_start = time.time()
    for epoch in range(cfg.epochs):
        # augmentation/reweight CHỈ bật ở nửa train của epoch — val phải
        # thấy dữ liệu y hệt c3, nếu không val_score không so được và
        # checkpoint sẽ chọn theo một thước khác (BG29 bước 2.1)
        ds.set_train_mode(True, epoch)
        tr = run_epoch_train(model, ds, tr_batcher.epoch_batches(epoch),
                             cfg, opt, sched, device)
        ds.set_train_mode(False)
        vm = run_val(model, ds, va_batches, cfg, device)
        is_best = vm["score"] < best["score"]
        if is_best:
            best = {"score": vm["score"], "epoch": epoch, "metrics": vm}
            torch.save({"state_dict": model.state_dict(),
                        "model_config": asdict(cfg.model),
                        "train_config": {k: v for k, v in asdict(cfg).items()
                                         if k != "model"},
                        "epoch": epoch, "val_metrics": vm,
                        "commit": _git_commit()},
                       out_dir / "best.pt")
        row = {"epoch": epoch, "lr": f"{sched.get_last_lr()[0]:.3e}",
               "elapsed_s": round(time.time() - t_start, 1),
               **{f"train_{k}": round(v, 5) for k, v in tr.items()},
               **{f"val_loss_{k}": round(vm[f"loss_{k}"], 5)
                  for k in LOSS_KEYS},
               "val_dphi_med": round(vm["dphi_med"], 3),
               "val_v0_relerr_med": round(vm["v0_relerr_med"], 4),
               "val_side_acc": round(vm["side_acc"], 4),
               "val_vert_acc": round(vm["vert_acc"], 4),
               "val_ident_acc": round(vm["ident_acc"], 4),
               "val_score": round(vm["score"], 4), "is_best": int(is_best)}
        with open(csv_path, "a", newline="", encoding="utf-8") as f:
            csv.DictWriter(f, fieldnames=fields).writerow(row)
        draw_curves(csv_path, png_path)
        print(f"ep {epoch:3d} | train {tr['total']:.4f} | "
              f"val {vm['loss_total']:.4f} | dphi {vm['dphi_med']:6.2f} | "
              f"v0rel {vm['v0_relerr_med']:.4f} | "
              f"side {vm['side_acc']:.3f} | vert {vm['vert_acc']:.3f} | "
              f"score {vm['score']:.3f}{' *best*' if is_best else ''}",
              flush=True)

    elapsed_min = (time.time() - t_start) / 60.0
    data_spec = json.loads((Path(args.data) / "spec.json")
                           .read_text(encoding="utf-8"))
    d = cfg.run_date
    spec = {
        "run_name": run_name, "created": f"{d[:4]}-{d[4:6]}-{d[6:]}",
        "commit": _git_commit(),
        "dataset": {"dir": str(args.data),
                    "data_commit": data_spec.get("commit"),
                    "n_train_total": len(ds),
                    "shards": [p.name for p in shards]},
        "split": {"rule": "val cat tu TRAIN theo permutation seed; "
                          "5k held-out KHONG dung khi train/tune",
                  "seed": cfg.split_seed, "val_frac": cfg.val_frac,
                  "n_train": int(len(tr_idx)), "n_val": int(len(va_idx))},
        "config": asdict(cfg), "n_params": n_params,
        "torch": torch.__version__,
        "gpu": (torch.cuda.get_device_name(0) if device == "cuda" else "cpu"),
        "epochs_run": cfg.epochs, "elapsed_min": round(elapsed_min, 1),
        "best": {"epoch": best["epoch"], "val_score": round(best["score"], 4),
                 "val_metrics": {k: (round(v, 5)
                                     if isinstance(v, float) else v)
                                 for k, v in best["metrics"].items()}},
        "val_score_formula": "dphi_med/2 + v0relerr/0.10 + (1-side)/0.20 "
                             "+ (1-vert)/0.20 (1.0 moi ve = bar G-27.3)",
        "augment": ("slot dropout TRAIN-only, on-the-fly tu 150k co san "
                    "(KHONG sinh dataset moi ra dia); keep-rate ~ U[min,max] "
                    "Bernoulli doc lap tren o (frame, slot khong-cue); cue "
                    "KHONG dropout; RNG khoa [seed, epoch, shot_idx]"
                    if cfg.slot_dropout else "khong"),
        "reweight": (f"loss weighting (KHONG oversample - giu nguyen so buoc "
                     f"moi epoch nhu c3): w={cfg.target_reweight.weight} cho "
                     f"cu co cham dau >= {cfg.target_reweight.t_first_s}s "
                     f"(13.62% train), trung binh co trong so nen thang "
                     f"gradient giu ~1"
                     if cfg.target_reweight else "khong"),
        "note": "checkpoint chon theo VAL - khong dung held-out (BRIEF 27)",
    }
    (out_dir / "spec.json").write_text(json.dumps(spec, indent=2),
                                       encoding="utf-8")
    print(f"XONG {elapsed_min:.1f} phut | best ep {best['epoch']} "
          f"score {best['score']:.4f} -> {out_dir}")


if __name__ == "__main__":
    main()