#!/usr/bin/env python
"""Routing comparison on a split (teacher-forced loss + extra ops) for: frozen backbone (budget 0), random /
uniform / fixed-pair schedules at budgets {1,2,4,8}, entropy and norm routers (threshold grid), the learned
controller (locked λ, κ; and a λ grid for the Pareto curve), the greedy one-step oracle and (on a small subset)
the two-step oracle. Oracles are DIAGNOSTIC (they read the target).

Usage: python -m evaluation.routing --config configs/experiments/routing_p4.yaml --split dev_iid [--device mps]
Writes results/raw/routing_<run>.json (one record per policy with loss, edits/token, extra ops ratio).
"""
from __future__ import annotations

import argparse, json, sys, time
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from prisme.flops import action_cost_proxy  # noqa: E402
from prisme.losses import token_nll  # noqa: E402
from prisme.oracle import greedy_oracle_rollout, two_step_oracle  # noqa: E402
from baselines.routers import rollout as router_rollout  # noqa: E402
from training.common import TokenBin, load_checkpoint, load_yaml, pick_device, seed_all, git_sha  # noqa: E402
from training.train_editor import build_model, sample_forced_actions  # noqa: E402
from data.synthetic_dataset import AnyTokenizer, SyntheticBatches, load_items  # noqa: E402


def ops_of(actions, cost, cfg):
    a = actions
    c = torch.cat([cost, torch.zeros(1, device=cost.device)])
    return c[a.clamp(max=cfg.n_actions)].sum(-1)


@torch.no_grad()
def evaluate_policy(model, cfg, batches, device, cost, kind, **kw):
    tot, n, ed, ops = 0.0, 0, 0.0, 0.0
    extra = {}
    for x, y in batches:
        v = y != -100
        if kind == "budget0":
            logits = model(x, None)["logits"]; acts = None
        elif kind == "forced":
            rng = kw["rng"]; acts, _ = sample_forced_actions(cfg, x.shape[0], x.shape[1], [kw["budget"]], [1.0], rng, device)
            out = model(x, acts); logits = out["logits"]
        elif kind in ("uniform", "fixed_pair", "entropy", "norm", "random"):
            out = router_rollout(model, x, kind, kw["budget"], threshold=kw.get("threshold", 0.0), seed=kw.get("seed", 0))
            logits, acts = out["logits"], out["actions"]
        elif kind == "controller":
            out = model.forward_policy(x, kw["budget"], cost, lambda_cost=kw.get("lam"), kappa_risk=kw.get("kap"))
            logits, acts = out["logits"], out["actions"]
        elif kind == "oracle1":
            out = greedy_oracle_rollout(model, x, y, kw["budget"])
            nll = out["loss_per_step"][-1]; acts = out["actions"]; logits = None
        else:
            raise ValueError(kind)
        if logits is not None:
            nll = token_nll(logits, y)
        tot += float(nll[v].sum()); n += int(v.sum())
        if acts is not None:
            ed += float((acts != cfg.stop_action).sum(-1)[v].float().sum()); ops += float(ops_of(acts, cost, cfg)[v].sum())
    return {"loss": tot / max(n, 1), "edits_per_token": ed / max(n, 1), "extra_ops_ratio": ops / max(n, 1), "n_tokens": n}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--config", required=True)
    ap.add_argument("--split", default=None)
    ap.add_argument("--device", default="auto")
    ap.add_argument("--out", default=None)
    args = ap.parse_args(argv)
    exp = load_yaml(args.config); seed_all(exp.get("seed", 1)); device = pick_device(args.device)
    cfg, model = build_model(exp); model = model.to(device).eval()
    meta = load_checkpoint(exp["checkpoint"], model, strict=False)
    data = exp["data"]; split = args.split or data.get("split", "dev_iid")
    bs, T = exp["batch_size"], exp["seq_len"]
    def batches(n=None):
        if data["kind"] == "tokens":
            return TokenBin(data["bin"], T).sequential_batches(bs, device, n or exp.get("n_batches"))
        tok = AnyTokenizer(data.get("tokenizer", "char"))
        sb = SyntheticBatches(load_items(Path(data["splits_dir"]) / f"{split}.jsonl"), tok, T)
        return sb.sequential_batches(bs, device, data.get("loss_span", "answer"), (n or exp.get("n_batches", 10**9)) * bs)
    cost = torch.tensor(action_cost_proxy(cfg, T), device=device)
    budgets = exp.get("budgets", [1, 2, 4, 8])
    recs = []
    def add(name, res, **tags):
        r = {"policy": name, **tags, **res}; recs.append(r); print(json.dumps(r), flush=True)
    t0 = time.time()
    add("budget0", evaluate_policy(model, cfg, batches(), device, cost, "budget0"))
    for b in budgets:
        add("random_forced", evaluate_policy(model, cfg, batches(), device, cost, "forced", budget=b, rng=np.random.default_rng(0)), budget=b)
        add("uniform", evaluate_policy(model, cfg, batches(), device, cost, "uniform", budget=b), budget=b)
        add("fixed_pair", evaluate_policy(model, cfg, batches(), device, cost, "fixed_pair", budget=b), budget=b)
    for thr in exp.get("entropy_thresholds", [0.1, 0.5, 1.0, 2.0]):
        add("entropy", evaluate_policy(model, cfg, batches(), device, cost, "entropy", budget=max(budgets), threshold=thr), threshold=thr, budget=max(budgets))
    for thr in exp.get("norm_thresholds", [0.02, 0.05, 0.1, 0.2]):
        add("norm", evaluate_policy(model, cfg, batches(), device, cost, "norm", budget=max(budgets), threshold=thr), threshold=thr, budget=max(budgets))
    if exp.get("controller", True):
        add("controller_locked", evaluate_policy(model, cfg, batches(), device, cost, "controller", budget=max(budgets), lam=cfg.lambda_cost, kap=cfg.kappa_risk), lam=cfg.lambda_cost, kap=cfg.kappa_risk, budget=max(budgets))
        for lam in exp.get("lambda_grid", [0.0, 0.01, 0.03, 0.1, 0.3]):
            add("controller", evaluate_policy(model, cfg, batches(), device, cost, "controller", budget=max(budgets), lam=lam, kap=cfg.kappa_risk), lam=lam, kap=cfg.kappa_risk, budget=max(budgets))
    nb_or = exp.get("oracle_batches", 4)
    for b in budgets:
        add("oracle1_greedy", evaluate_policy(model, cfg, batches(nb_or), device, cost, "oracle1", budget=b), budget=b, diagnostic=True)
    # two-step oracle on a small subset
    nb2 = exp.get("two_step_batches", 1)
    acc = {"one": 0.0, "two": 0.0, "greedy_then_best": 0.0, "n": 0}
    for i, (x, y) in enumerate(batches(nb2)):
        o = two_step_oracle(model, x, y)
        v = y != -100
        acc["one"] += float(o["one_step_best"][v].sum()); acc["two"] += float(o["two_step_best"][v].sum())
        acc["greedy_then_best"] += float(o["greedy_then_best"][v].sum()); acc["n"] += int(v.sum())
    add("oracle_two_step", {"one_step_gain": acc["one"] / max(acc["n"], 1), "two_step_gain": acc["two"] / max(acc["n"], 1),
                            "greedy_then_best_gain": acc["greedy_then_best"] / max(acc["n"], 1), "n_tokens": acc["n"]}, diagnostic=True)
    out = {"config": args.config, "checkpoint": exp["checkpoint"], "checkpoint_sha256": meta.get("weights_sha256"), "split": split,
           "code_commit": git_sha(), "device": str(device), "seconds": round(time.time() - t0, 1), "records": recs}
    path = Path(args.out or f"results/raw/routing_{Path(exp['checkpoint']).parent.name}_{split}.json")
    path.parent.mkdir(parents=True, exist_ok=True); path.write_text(json.dumps(out, indent=1))
    print("wrote", path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())