"""Evaluate one checkpoint across the full sweep grid and dump JSON.

Sweeps (fresh random mappings everywhere; nothing from training is reused):
  - accuracy vs hop count k in 1..64 at the model's training-time unroll count
  - accuracy vs inference-time unroll count n_iters in 1..16, for
    k in {1, 2, 4, 8, 16, 32, 64}  (both architectures: a 'fixed' model run
    at n_iters>1 composes its stack repeatedly, i.e. deeper than trained)
  - accuracy vs state-space size S' in {16, 32, 64} across k in 1..64
  - LookupSolver reference: must be 1.0 (pipeline integrity check)
  - inference FLOPs per example recorded per setting

Usage: python evaluate.py --ckpt out/rec1x4-seed0.pt --out-dir results
"""

import argparse
import json
import os

import numpy as np
import torch

from ptr.common import LookupSolver, sample_eval_batch
from ptr.flops import step_flops
from ptr.models import PtrTransformer

KS_FULL = list(range(1, 65))
KS_ITERS = [1, 2, 4, 8, 16, 32, 64]
ITERS_GRID = list(range(1, 17))
STATE_SIZES = [16, 32, 64]


def load_model(ckpt_path, device):
    ck = torch.load(ckpt_path, map_location=device, weights_only=False)
    cfg = ck["config"]
    model = PtrTransformer(cfg["S"], cfg["d_model"], 4, 4 * cfg["d_model"],
                           depth=cfg["depth"],
                           stem_depth=cfg.get("stem_depth", 0))
    model.load_state_dict(ck["state_dict"])
    model.to(device).eval()
    return model, ck


@torch.no_grad()
def acc_at_k(model, rng, S, k, n, n_iters, device):
    b = sample_eval_batch(rng, n, S, k=k)
    f = torch.as_tensor(b["f"], device=device)
    x = torch.as_tensor(b["x"], device=device)
    kt = torch.as_tensor(b["k"], device=device)
    y = torch.as_tensor(b["answer"], device=device)
    logits = model(f, x, kt, n_iters=n_iters)
    # the head is sized to the trained S; for a smaller eval state space
    # mask the logits to the first S classes (answers are always < S)
    logits = logits[:, :S]
    return float((logits.argmax(-1) == y).float().mean())


def evaluate_ckpt(ckpt_path, out_dir, n_per_cell=1000, seed=1234):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    rng = np.random.default_rng(seed)
    model, ck = load_model(ckpt_path, device)
    name = os.path.splitext(os.path.basename(ckpt_path))[0]
    cfg = ck["config"]
    spec = ck["n_iters_spec"]
    aligned = isinstance(spec, tuple) and spec[0] in (
        "aligned", "aligned_margin")
    margin = spec[3] if isinstance(spec, tuple) and len(spec) > 3 else 0
    if isinstance(spec, tuple):
        train_iters = (spec[1] + spec[2]) / 2 + margin
    else:
        train_iters = spec

    def it_for(k):
        # aligned variants read out at their trained offset: R = k + margin
        return k + margin if aligned else int(train_iters)

    is_rec = cfg["arch"] == "rec"

    res = {"run_name": name, "config": cfg, "budget": ck["budget"],
           "device": device, "n_per_cell": n_per_cell, "aligned": aligned}

    # pipeline integrity: the lookup reference must be exact
    b = sample_eval_batch(rng, 2000, cfg["S"], k=64)
    assert (LookupSolver().solve(b["f"], b["x"], 64) == b["answer"]).all()
    res["lookup_reference_acc"] = 1.0

    # accuracy vs k at training-time unrolls (R=k for aligned models)
    acc_k = {k: acc_at_k(model, rng, cfg["S"], k, n_per_cell, it_for(k),
                         device) for k in KS_FULL}
    res["acc_vs_k"] = acc_k
    res["train_iters"] = train_iters
    res["inference_flops_per_example"] = step_flops(
        model, cfg["S"], 1, train_iters, training=False)

    # accuracy vs inference-time unroll count (OOD iterations)
    iters_grid = ITERS_GRID
    acc_iters = {str(r): {k: acc_at_k(model, rng, cfg["S"], k, n_per_cell, r,
                                      device) for k in KS_ITERS}
                 for r in iters_grid}
    res["acc_vs_iters"] = acc_iters

    # accuracy vs state-space size (smaller spaces reuse trained vocab rows)
    acc_state = {}
    for Sp in STATE_SIZES:
        if Sp > cfg["S"]:
            continue
        acc_state[str(Sp)] = {k: acc_at_k(model, rng, Sp, k, n_per_cell,
                                          it_for(k), device) for k in
                              KS_FULL}
    res["acc_vs_state_size"] = acc_state

    os.makedirs(out_dir, exist_ok=True)
    out_path = os.path.join(out_dir, f"{name}_eval.json")
    with open(out_path, "w") as fh:
        json.dump(res, fh)
    hub = os.environ.get("HUB_REPO")
    if hub:
        from huggingface_hub import HfApi
        HfApi().upload_file(path_or_fileobj=out_path,
                            path_in_repo=f"results/{name}_eval.json",
                            repo_id=hub, repo_type="model")
        print(f"[hub] uploaded eval json to {hub}")
    id64 = sum(v for k, v in acc_k.items() if k <= 8) / 8
    print(f"[eval] {name}: mean acc k=1..8 {id64:.3f}, "
          f"k=64 {acc_k[64]:.3f} (train_iters={train_iters}) -> {out_path}")
    return out_path


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--ckpt", required=True)
    p.add_argument("--out-dir", default="results")
    p.add_argument("--n-per-cell", type=int, default=1000)
    a = p.parse_args()
    evaluate_ckpt(a.ckpt, a.out_dir, a.n_per_cell)