#!/usr/bin/env python
"""Reproduce / update the analytic ops-per-token table of the WP31 brief (Sec. 3.9) from the implemented model.
Usage: python scripts/estimate_flops.py --config configs/prisme_459m.yaml --T 4096 [--json results/tables/ops_459m.json]
"""
from __future__ import annotations

import argparse, json, sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from prisme.config import load_config  # noqa: E402
from prisme.flops import ops_per_token, action_cost_proxy  # noqa: E402

BRIEF_459M_T4096 = {"backbone_only": 1.296e9, "4_actions_k1": 1.645e9, "4_actions_k2": 1.767e9,
                    "8_actions_k2": 2.195e9, "8_actions_dense16": 5.608e9}


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--config", default="configs/prisme_459m.yaml")
    ap.add_argument("--T", type=int, default=4096)
    ap.add_argument("--json", default=None)
    args = ap.parse_args()
    c = load_config(args.config)
    rows = {
        "backbone_only": ops_per_token(c, args.T, 0, 1),
        "4_actions_k1": ops_per_token(c, args.T, 4, 1),
        "4_actions_k2": ops_per_token(c, args.T, 4, 2),
        "8_actions_k2": ops_per_token(c, args.T, 8, 2),
        "8_actions_dense16": ops_per_token(c, args.T, 8, c.n_slots),
    }
    out = {"config": c.name, "T": args.T, "assumptions": __import__("prisme.flops", fromlist=["x"]).__doc__,
           "rows": {k: v.as_dict() for k, v in rows.items()}, "brief_reference_459m_T4096": BRIEF_459M_T4096,
           "action_cost_proxy": action_cost_proxy(c, args.T)}
    print(f"{c.name} @ T={args.T}: ops/token (x1e9); controller evaluated n_actions+1 times (capped at r_max)")
    print(f"{'policy':20s} {'total':>8s} {'backbone':>9s} {'workspace':>10s} {'editor':>8s} {'ctrl':>7s} {'readout':>8s} {'brief':>8s}")
    for k, v in rows.items():
        ed = (v.editor_source_attn + v.editor_slot_attn + v.editor_ffn + v.editor_gate) / 1e9
        ref = BRIEF_459M_T4096[k] / 1e9 if (c.name == "prisme_459m" and args.T == 4096) else float("nan")
        print(f"{k:20s} {v.total/1e9:8.3f} {v.backbone/1e9:9.3f} {v.workspace/1e9:10.3f} {ed:8.3f} {v.controller/1e9:7.3f} {v.readout/1e9:8.3f} {ref:8.3f}")
    print("cost proxy c(a) (single, pair):", round(out["action_cost_proxy"][0], 4), round(out["action_cost_proxy"][-1], 4))
    if args.json:
        Path(args.json).parent.mkdir(parents=True, exist_ok=True)
        Path(args.json).write_text(json.dumps(out, indent=2))
    return 0


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