#!/usr/bin/env python
"""Exact parameter accounting for PRISME configurations (meta-device instantiation).

Reports unique parameters per component and reconciles against the reference counts of the WP31 brief.
Usage: python scripts/count_params.py --all-configs --device meta [--json results/tables/param_counts.json]
"""
from __future__ import annotations

import argparse, json, sys
from pathlib import Path

import torch

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

REFERENCE = {  # from the WP31 brief (reference design, pre-implementation)
    "prisme_103m": {"backbone": 99_502_848, "workspace": 3_088_260, "total": 102_591_108},
    "prisme_459m": {"backbone": 446_744_064, "workspace": 12_336_900, "total": 459_080_964},
    "prisme_1180m": {"backbone": 1_149_339_648, "workspace": 25_779_844, "total": 1_175_119_492},
}


def closed_form_backbone(c: PrismeConfig) -> int:
    d, f, V, L = c.d_model, c.ffn_width, c.vocab_size, c.n_layers
    kv = c.kv_dim
    per_layer = 2 * d * d + 2 * d * kv + 3 * d * f + 2 * d
    return V * d + L * per_layer + d


def closed_form_workspace(c: PrismeConfig) -> dict[str, int]:
    d, w, fw, M = c.d_model, c.w, c.ffn_w, c.n_slots
    hidden = c.controller_hidden or fw
    parts = {
        "bank_norm": d,
        "bank_projections (W_K, W_V, P)": 3 * d * w,
        "slot_embed e": M * w,
        "editor.block (shared x Lw)": 6 * w * w + 3 * w * fw + 3 * w,
        "editor.ctx_proj": 3 * w + w,
        "editor.gate": 2 * w + 1,
        "readout (q,k,U)": 2 * w * w + w * d,
        "controller": 2 * w * w + (2 * w + w) + w + 3 * w * hidden + 2 * w + (4 * w + 4),
    }
    return parts


def reference_closed_form(c: PrismeConfig) -> int:
    """Algebraic reconstruction of the brief's workspace/controller count from its three published values:
    4*d*w + M*w + 17*w^2 + 6*w*f_w + 15.5*w + 4 (fits all three configurations exactly)."""
    d, w, fw, M = c.d_model, c.w, c.ffn_w, c.n_slots
    return int(4 * d * w + M * w + 17 * w * w + 6 * w * fw + 15.5 * w + 4)


def count(model: torch.nn.Module) -> dict[str, int]:
    seen, out = set(), {}
    for n, p in model.named_parameters():
        if id(p) in seen:
            continue
        seen.add(id(p))
        top = n.split(".")[0]
        out[top] = out.get(top, 0) + p.numel()
    return out


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--all-configs", action="store_true")
    ap.add_argument("--config", nargs="*", default=[])
    ap.add_argument("--device", default="meta")
    ap.add_argument("--json", default="results/tables/param_counts.json")
    args = ap.parse_args()
    paths = sorted(Path("configs").glob("prisme_*.yaml")) if args.all_configs else [Path(p) for p in args.config]
    report = {}
    for p in paths:
        cfg = load_config(p)
        with torch.device(args.device):
            model = PrismeModel(cfg)
        by_top = count(model)
        backbone = by_top.get("backbone", 0)
        workspace = sum(v for k, v in by_top.items() if k != "backbone")
        total = backbone + workspace
        cf_b = closed_form_backbone(cfg)
        cf_w = closed_form_workspace(cfg)
        ref = REFERENCE.get(cfg.name)
        rec = {
            "config": str(p), "name": cfg.name, "config_sha256": cfg.sha256(),
            "by_module": by_top, "backbone": backbone, "workspace_controller": workspace, "total": total,
            "closed_form_backbone": cf_b, "closed_form_workspace_parts": cf_w, "closed_form_workspace": sum(cf_w.values()),
            "reference_brief": ref,
            "reference_workspace_closed_form_reconstruction": reference_closed_form(cfg),
            "backbone_matches_reference": (ref is not None and backbone == ref["backbone"]),
            "workspace_delta_vs_reference": (workspace - ref["workspace"]) if ref else None,
        }
        assert backbone == cf_b, (backbone, cf_b)
        assert workspace == sum(cf_w.values()), (workspace, sum(cf_w.values()))
        report[cfg.name] = rec
        print(f"{cfg.name}: backbone={backbone:,} workspace+controller={workspace:,} total={total:,} "
              f"| ref backbone={ref['backbone']:,} ({'OK' if rec['backbone_matches_reference'] else 'MISMATCH'}), "
              f"ref workspace={ref['workspace']:,} (delta {rec['workspace_delta_vs_reference']:+,})")
    Path(args.json).parent.mkdir(parents=True, exist_ok=True)
    Path(args.json).write_text(json.dumps(report, indent=2))
    return 0


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