"""cap_test 5-shot evaluation harness (Annulus base-v4, year-aware).

Dataset: lfqian/cap_test  — 6 tasks (splits), each row = {input, output, task}:
    copying, uppercase, basic_arithmetic, plural_upper, gerund_upper, math

We evaluate OUR year-aware model under BOTH prompt conventions it was trained on:

  --mode task1  (year-conditioned generation)
      Sequence starts with a year token [Yt], then a 5-shot demo block, then the
      query. The model just continues (plain LM). This is exactly how Task1
      samples look:  [Yt] <text...>.

  --mode task2  (self-emit-year-then-continue)
      Mirrors Task2 layout  ctx <answer> [Yt] cont : the demo block is the ctx,
      then we append <answer>, force-feed the correct [Yt], and let the model
      produce the answer as `cont`. Tests that Task2 wiring didn't break the
      model's ability to still do the surface task after emitting the year token.

Both share the same 5-shot demonstrations and exact-match scoring; the only
difference is where the [Yt]/<answer> special tokens sit — so a gap between the
two modes localises to the year-token machinery, not the task itself.

5-shot: for each query we sample 5 OTHER (input,output) pairs from the same task
as in-context demonstrations (fixed rng per task so runs are comparable).

Output: per-task exact-match accuracy + overall, optional --csv for review.

Model loading: our ckpt is a NeMo/Megatron distcp MoE with a custom year router,
NOT loadable by AutoModelForCausalLM. Plug the real generate() in `load_model()`
(NeMo inference or the offline-inference path). A HF AutoModel branch is provided
for a converted ckpt / for running a Qwen baseline for comparison.
"""
import argparse, csv, json, os, random, sys

# *** CRITICAL PARITY (固化, 2026-09-18): v5 was trained with ANNULUS_YEAR_OUTPUT=0.
# Eval MUST run YEAR_OUTPUT=0 too, else apply_year_output_logits injects untrained ct
# into the year-vocab columns → model spams [Yxxxx] tokens → cap generation garbage.
# code_v7 default is "1"; force 0 here so the harness never drifts (was the cap-crash bug). ***
os.environ.setdefault("ANNULUS_YEAR_OUTPUT", "0")

# ---- year-token config (v5f tokenizer) ---------------------------------------
YEAR_TOKEN_DEFAULT = "[Y2000]"   # conditioning year for the eval prompts
ANSWER_TOKEN = "<answer>"        # id 151875 in v5f
DATE_TOKEN   = "<date?>"         # id 151874 in v5f -- leads every Task2 sample

TASKS = ["copying", "uppercase", "basic_arithmetic",
         "plural_upper", "gerund_upper", "math"]


# ---- data --------------------------------------------------------------------
def load_cap_test():
    """Return {task: [ {input, output}, ... ]} for all 6 splits.

    CAP_TEST_DIR lets us read the parquet from disk: compute nodes have no
    internet. Verified md5-identical to lfqian/cap_test on the Hub.
    """
    local = os.environ.get("CAP_TEST_DIR")
    if local:
        import pyarrow.parquet as pq, glob
        out = {}
        for t in TASKS:
            hits = glob.glob(os.path.join(local, f"{t}*.parquet"))
            if not hits:
                raise SystemExit(f"no parquet for task {t} in {local}")
            rows = pq.read_table(hits[0]).to_pylist()
            out[t] = [{"input": r["input"], "output": r["output"]} for r in rows]
        return out
    from datasets import load_dataset
    out = {}
    for t in TASKS:
        ds = load_dataset("lfqian/cap_test", split=t)
        out[t] = [{"input": r["input"], "output": r["output"]} for r in ds]
    return out


# ---- prompt construction -----------------------------------------------------
def build_demos(rows, query_idx, k, rng):
    """Pick k demonstration rows != query_idx."""
    pool = [i for i in range(len(rows)) if i != query_idx]
    idx = rng.sample(pool, min(k, len(pool)))
    return [rows[i] for i in idx]


def render_shots(demos, query_in):
    # Arrow format (no Input:/Output: labels). The arrow lives in the query too,
    # so the model only has to emit the RESULT after "-> ".
    body = "".join(f"{d['input']} -> {d['output']}\n" for d in demos)
    return f"{body}{query_in} ->"


def make_prompt(mode, year_token, demos, query_in):
    """Return the string fed to the tokenizer. Arrow format.

    task1:  [Yt] {demos} q ->        model emits only the result
    task2:  <date?> {demos} q -> <answer>    the arrow is IN the input; after
            <answer> the model self-emits [Yt] then the result.
    """
    shots = render_shots(demos, query_in)               # ends with "q ->"
    if mode == "plain":
        return shots                                     # no special token at all
    if mode == "task1":
        return f"{year_token}{shots}"
    elif mode == "task2":
        return f"{DATE_TOKEN}{shots}{ANSWER_TOKEN}"
    raise ValueError(mode)


# ---- model interface ---------------------------------------------------------
def load_model(model_path, backend, trust_remote_code=False):
    """Return a callable generate(prompt:str, max_new_tokens:int)->str."""
    if backend == "hf":
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        tok = AutoTokenizer.from_pretrained(model_path, trust_remote_code=trust_remote_code)
        model = AutoModelForCausalLM.from_pretrained(
            model_path, torch_dtype=torch.bfloat16, device_map="auto",
            trust_remote_code=trust_remote_code)
        model.eval()

        def generate(prompt, max_new_tokens=24):
            ids = tok(prompt, return_tensors="pt").to(model.device)
            with torch.no_grad():
                out = model.generate(**ids, max_new_tokens=max_new_tokens,
                                     do_sample=False,
                                     pad_token_id=tok.eos_token_id)
            new = out[0][ids["input_ids"].shape[1]:]
            # keep special tokens so task2's emitted [Yt] is visible to strip_special
            return tok.decode(new, skip_special_tokens=False)
        return generate, tok

    if backend == "nemo":
        # annulus-base-v4: Megatron distcp MoE + custom year router.
        # model_path = the ckpt dir containing weights/ and context/.
        import types, glob, torch
        for _v in ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN"):
            os.environ.pop(_v, None)

        _m = "megatron.core.dist_checkpointing.strategies.tensorstore"
        try:
            __import__(_m)
        except Exception:
            st = types.ModuleType(_m)

            class _T:
                def __init__(s, *a, **k): pass
            st.TensorStoreLoadShardedStrategy = _T
            st.TensorStoreSaveShardedStrategy = _T
            sys.modules[_m] = st
            import megatron.core.dist_checkpointing.strategies as _s
            _s.tensorstore = st

        import torch.distributed as dist
        for k, v in dict(MASTER_ADDR="localhost",
                         MASTER_PORT=os.environ.get("MASTER_PORT", "29761"),
                         RANK="0", WORLD_SIZE="1", LOCAL_RANK="0").items():
            os.environ.setdefault(k, v)
        torch.cuda.set_device(0)
        if not dist.is_initialized():
            dist.init_process_group("nccl", rank=0, world_size=1)
        from megatron.core import parallel_state as ps
        if not ps.model_parallel_is_initialized():
            ps.initialize_model_parallel(1, 1)
        from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed
        model_parallel_cuda_manual_seed(1234)

        WEIGHTS = os.path.join(model_path, "weights")
        CTX = os.path.join(model_path, "context")
        from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
        _td = [d for d in sorted(glob.glob(os.path.join(CTX, "*"))) if os.path.isdir(d)
               and os.path.exists(os.path.join(d, "tokenizer_config.json"))]
        if not _td:
            raise SystemExit(f"no tokenizer dir under {CTX}")
        tok = AutoTokenizer(pretrained_model_name=_td[0]); ht = tok.tokenizer
        print(f"[nemo] tokenizer {_td[0]}", flush=True)

        # --- year-token ids ---------------------------------------------------
        # On the tf4-launch-fix baseline YEAR_TOKEN_BASE_ID is None by design and
        # resolve_year_token_base_id(tokenizer) is the single source of truth
        # (fail-loud, plus a [Y2000..Y2025] contiguity check). For v5f it returns
        # [Y2000]=151830. NOTE: 151830 is the anchor of the 26-row year_context.e
        # table (YEAR_MIN..YEAR_MAX = 2000..2025) -- NOT to be confused with 151646,
        # which is the base of the year-token VOCAB block ([Ypre1700], 210 bins per
        # year_bin_spec.json). Two different things.
        import annulus.model.year_routing as _yr
        import annulus.model.year_embedding as _ye
        if _yr.YEAR_TOKEN_BASE_ID is not None:
            raise SystemExit(
                f"[baseline] annulus.year_routing.YEAR_TOKEN_BASE_ID is hardcoded to "
                f"{_yr.YEAR_TOKEN_BASE_ID} -- that is the PRE-v4 package. Point "
                f"PYTHONPATH at the tf4-launch-fix code (code_tf4/).")
        _y2000 = _yr.resolve_year_token_base_id(ht)      # fail-loud, verifies contiguity
        _ye.verify_year_vocab_layout(ht)
        print(f"[nemo] year ids resolved from tokenizer: [Y2000]={_y2000}", flush=True)

        # --- config -----------------------------------------------------------
        # tf4-launch-fix ships AnnulusGPTConfig (what context/io.json references)
        # and its env defaults already are the v4 cut (4 routed / top-3 / 1 shared
        # / 27L), so the clean io.load_context path works. model.yaml is still
        # applied afterwards as the checkpoint-authoritative cross-check.
        cfg = None
        try:
            from nemo.lightning import io as _nlio
            cfg = _nlio.load_context(CTX, subpath="model.config")
            print("[nemo] CONFIG: from context/io.json (AnnulusGPTConfig)", flush=True)
        except Exception as _e:
            print(f"[nemo] io.load_context failed: {repr(_e)[:200]}", flush=True)
            from annulus.model.base_config import make_base_config
            cfg = make_base_config()
            print("[nemo] CONFIG: make_base_config()", flush=True)

        import yaml as _yaml
        _mp = os.path.join(CTX, "model.yaml")
        if not os.path.exists(_mp):
            raise SystemExit(f"no model.yaml in {CTX}; arch cannot be trusted")
        _y = (_yaml.safe_load(open(_mp)) or {}).get("config", {}) or {}
        _chg = []
        for _k, _v in _y.items():
            if _k.startswith("_") or isinstance(_v, dict) or _v is None: continue
            if isinstance(_v, (list, tuple)) and any(isinstance(_x, (dict, list)) for _x in _v):
                continue
            if hasattr(cfg, _k) and getattr(cfg, _k) != _v:
                try:
                    _chg.append((_k, getattr(cfg, _k), _v)); setattr(cfg, _k, _v)
                except Exception: pass
        print(f"[nemo] {len(_chg)} config overrides from model.yaml "
              f"(0 expected if the baseline matches the ckpt)", flush=True)
        for _k, _o, _n in _chg:
            print(f"        {_k}: {_o} -> {_n}", flush=True)
        for _a in ("num_layers", "num_moe_experts", "moe_router_topk",
                   "moe_shared_expert_intermediate_size", "hidden_size"):
            print(f"        cfg.{_a} = {getattr(cfg, _a, 'n/a')}", flush=True)
        try:
            from megatron.core.transformer.enums import AttnBackend
            cfg.attention_backend = AttnBackend.unfused
        except Exception: pass

        # ---- CRITICAL: wire the year hooks, exactly as training does -------------
        # train_ckpt.py wraps cfg.configure_model so wire_year_hooks() runs on the freshly
        # built mcore model: it swaps every MoE router for AnnulusYearRouter and registers
        # year_context (e_t). WITHOUT this the model is a plain TopKRouter GPT with NO year
        # conditioning, the ckpt's year_context.e is never even requested by
        # sharded_state_dict(), and train/inference silently diverge.
        from annulus.model.year_aware_model import wire_year_hooks
        mcore = cfg.configure_model(tok)
        _yc = wire_year_hooks(mcore, tok, register_on_model=True, verbose=True)
        _routers = [m for _, m in mcore.named_modules()
                    if m.__class__.__name__ == "AnnulusYearRouter"]
        if not _routers:
            raise SystemExit("[wiring] no AnnulusYearRouter after wire_year_hooks")
        _e = [n for n, _ in mcore.named_parameters() if n.endswith("year_context.e")]
        if len(_e) != 1:
            raise SystemExit(f"[wiring] expected 1 year_context.e, found {len(_e)}")
        print(f"[wiring] AnnulusYearRouter x{len(_routers)}, year_context.e ✓", flush=True)
        mcore = mcore.cuda().eval()
        from megatron.core import dist_checkpointing
        sd = mcore.sharded_state_dict(prefix="module.")
        loaded = dist_checkpointing.load(sd, WEIGHTS)
        strip = {(k[7:] if k.startswith("module.") else k): v for k, v in loaded.items()}
        miss = mcore.load_state_dict(strip, strict=False)
        real = [m for m in miss.missing_keys if "optim" not in m]   # do NOT hide year_context
        print(f"[nemo] loaded {WEIGHTS}  missing_nonopt={len(real)}", flush=True)
        if real[:5]:
            print(f"        e.g. {real[:5]}", flush=True)

        EOD = ht.convert_tokens_to_ids("<|endoftext|>")

        @torch.no_grad()
        def generate(prompt, max_new_tokens=24):
            # greedy, no early newline stop -> mirrors the HF branch, so task2's
            # `gold in first-N-token window` sees the same window.
            ids = ht(prompt, add_special_tokens=False)["input_ids"]
            out = []
            for _ in range(max_new_tokens):
                s = len(ids)
                inp = torch.tensor([ids], device="cuda")
                pos = torch.arange(s, device="cuda")[None]
                m = torch.triu(torch.ones(s, s, dtype=torch.bool, device="cuda"), 1)[None, None]
                o = mcore(input_ids=inp, position_ids=pos, attention_mask=m)
                v = (o[0, -1] if o.shape[0] == 1 else o[-1, 0])
                nx = int(v.float().argmax())
                ids.append(nx); out.append(nx)
                if nx == EOD: break
            # keep special tokens so split_year() can see an emitted [Yt]
            return ht.decode(out, skip_special_tokens=False)

        return generate, ht
    raise ValueError(backend)


# ---- eval --------------------------------------------------------------------
import re
_YEAR_RE = re.compile(r"^\s*(\[Y[^\]]*\]|<beyond_horizon>|<far_future>)\s*")
_SPECIAL_RE = re.compile(r"\[Y[^\]]*\]|<answer>|<future>|<far_future>|<beyond_horizon>")


def split_year(text):
    """For task2: pull the model-emitted leading year token off the front.
    Returns (year_token_or_None, remaining_text)."""
    m = _YEAR_RE.match(text)
    if m:
        return m.group(1), text[m.end():]
    return None, text


def strip_special(text):
    """Remove year/special tokens so a bare gold substring can't collide with
    the digits inside an emitted year token (e.g. '8' inside '[Y1998]')."""
    return _SPECIAL_RE.sub(" ", text)


def first_line(text):
    return text.strip().split("\n")[0].strip()


def run(args):
    rng = random.Random(args.seed)
    data = load_cap_test()
    generate, _ = load_model(args.model, args.backend, args.trust_remote_code)

    rows_out, totals = [], {}
    # historical cap metric (cap_icl_v4:64,219): first-line + strip + lower, substring containment.
    def _norm(s): return str(s).split("\n")[0].strip().lower()
    for task in TASKS:
        rows = data[task]
        n_len = 0; n_str = 0
        n = min(args.max_per_task, len(rows)) if args.max_per_task else len(rows)
        task_rng = random.Random(args.seed + hash(task) % 1000)
        for qi in range(n):
            q = rows[qi]
            demos = build_demos(rows, qi, args.shots, task_rng)
            prompt = make_prompt(args.mode, args.year_token, demos, q["input"])
            gold = q["output"].strip()
            emitted_year = None
            if args.mode == "task2":
                raw = generate(prompt, max_new_tokens=args.task2_window)
                emitted_year, _ = split_year(raw)
                window = strip_special(raw)
                pred = window.strip()
                ok_str = (gold in window); ok_len = (_norm(gold) in window.lower())
            else:
                raw = generate(prompt)
                pred = first_line(raw)
                ok_str = (pred == gold)                          # strict, case-sensitive (reference col)
                ok_len = (_norm(gold) in raw.strip().lower())    # lenient = cap_icl_v4:219 (historical acc)
            n_str += ok_str; n_len += ok_len
            if qi < 10:   # real pred-vs-gold samples (Manager gate #2)
                print(f"[SAMPLE] {task:12s} input={q['input']!r} gold={gold!r} pred={pred!r} "
                      f"raw={raw.strip()[:50]!r} strict={int(ok_str)} lenient={int(ok_len)}", flush=True)
            rows_out.append({"task": task, "mode": args.mode, "input": q["input"],
                             "gold": gold, "pred": pred, "raw": raw.strip()[:80],
                             "lenient": int(ok_len), "strict": int(ok_str),
                             "emitted_year": emitted_year or ""})
        totals[task] = (n_len, n_str, n)
        print(f"[{args.mode}] {task:16s} lenient={n_len/n:.3f} strict={n_str/n:.3f}  (n={n})", flush=True)

    tl = sum(l for l, _, _ in totals.values()); ts = sum(s for _, s, _ in totals.values())
    tn = sum(nn for _, _, nn in totals.values())
    print(f"[{args.mode}] {'OVERALL':16s} lenient={tl/tn:.3f} strict={ts/tn:.3f}  (n={tn})", flush=True)

    if args.csv:
        with open(args.csv, "w", newline="") as f:
            w = csv.DictWriter(f, fieldnames=list(rows_out[0].keys()))
            w.writeheader(); w.writerows(rows_out)
        print("wrote", args.csv)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("model", help="ckpt path / HF repo (or converted annulus-base-v4)")
    ap.add_argument("--backend", choices=["hf", "nemo"], default="nemo")
    ap.add_argument("--mode", choices=["plain", "task1", "task2"], default="task1")
    ap.add_argument("--shots", type=int, default=5)
    ap.add_argument("--task2-window", dest="task2_window", type=int, default=8,
                    help="task2: match gold within the first N generated tokens")
    ap.add_argument("--year-token", dest="year_token", default=YEAR_TOKEN_DEFAULT)
    ap.add_argument("--max-per-task", dest="max_per_task", type=int, default=0)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--trust-remote-code", action="store_true")
    ap.add_argument("--csv")
    args = ap.parse_args()
    run(args)


if __name__ == "__main__":
    main()