"""Isolated-process CPU/MPS parity and synchronized warm latency study.

Run after the reference evaluation finishes for uncontended timing. Every schema/input
gets one untimed warm call on each device before repeated measurements.
"""

import argparse
import gc
import json
from pathlib import Path

import numpy as np
import torch

from sealcell.config import config
from sealcell.model import LocalModel, schema_spec
from sealcell.policies import policy
from sealcell.serialization import render
from sealcell.simulator import Simulator
from sealcell.supervisor import Supervisor


def episode_parity(model):
    reference = {}
    source = Path("reports/runs/final-primary/decisions.jsonl")
    if not source.exists():
        return {"status": "not run; CPU reference traces missing"}
    with source.open() as f:
        for line in f:
            row = json.loads(line)
            if row["episode_id"] > "final-00012":
                break
            if row["policy"] in {"hybrid", "direct"}:
                reference.setdefault((row["episode_id"], row["policy"]), []).append(row)
    settings = config("evaluation")
    failures, agreements = [], 0
    with Path("reports/mps-episode-parity.jsonl").open("w") as output:
        for (episode_id, name), cpu_trace in reference.items():
            index = int(episode_id.split("-")[1])
            s = Simulator(
                settings["scenarios"][index % len(settings["scenarios"])],
                settings["splits"]["final"]["seed_start"] + index,
                settings["splits"]["final"]["recipe"],
                settings["splits"]["final"]["text_family"],
            )
            sup = Supervisor(
                s,
                policy(
                    name,
                    model=model,
                    issue_threshold=settings["issue_threshold"],
                    evidence_threshold=settings["evidence_threshold"],
                    filter_actions=settings.get("filter_actions", False),
                    **settings["conditions"]["primary"],
                ),
            ).run()

            def signature(trace):
                return [
                    (
                        r["proposal"]["action"],
                        r["executed"],
                        r["gate_rejection"],
                        r["stage_after"],
                        r["time_after"],
                        (r["proposal"]["model"] or {}).get("input_hash"),
                        {
                            h: v["selected_label"]
                            for h, v in (r["proposal"]["model"] or {}).get("heads", {}).items()
                        },
                    )
                    for r in trace
                ]

            equal = signature(cpu_trace) == signature(sup.trace)
            agreements += equal
            if not equal:
                failures.append({"episode_id": episode_id, "policy": name})
            output.write(
                json.dumps(
                    {
                        "episode_id": episode_id,
                        "policy": name,
                        "matches_cpu": equal,
                        "trace": sup.trace,
                    }
                )
                + "\n"
            )
    return {"episodes": len(reference), "matching_trajectories": agreements, "failures": failures}


parser = argparse.ArgumentParser()
parser.add_argument("--repeats", type=int, default=3)
parser.add_argument("--output", default="reports/device-benchmark.json")
args = parser.parse_args()
if not torch.backends.mps.is_available():
    raise SystemExit("MPS unavailable")
source = json.loads(Path("data/fixtures/playground-probes.json").read_text())
cases = [{"id": c["id"], "text": c["text"], "schema": source["schema"]} for c in source["cases"]]
for line in Path("data/fixtures/development.jsonl").read_text().splitlines():
    case = json.loads(line)
    cases.append({"id": case["id"], "text": case["text"], "schema": schema_spec()})
for name in config("scenarios"):
    observation = Simulator(name, 42).observe()
    for direct in [False, True]:
        cases.append(
            {
                "id": f"initial-{name}-{direct}",
                "text": render(observation),
                "schema": schema_spec(direct),
            }
        )
result = {
    "precision": "float32",
    "fallback_enabled": False,
    "repeats": args.repeats,
    "timing": "synchronized, one untimed warm call for each exact input/schema shape",
    "records": {},
    "metadata": {},
}
for device in ["cpu", "mps"]:
    model = LocalModel(device=device)
    result["metadata"][device] = model.metadata
    records = []
    for i, case in enumerate(cases):
        warm = model.classify(case["text"], case["schema"])
        if warm["error"]:
            raise RuntimeError(f"{device}/{case['id']}: {warm['error']}")
        for repeat in range(args.repeats):
            record = model.classify(case["text"], case["schema"])
            records.append({"case_id": case["id"], "repeat": repeat, **record})
        if i % 10 == 0:
            print(device, i + 1, "/", len(cases), flush=True)
    result["records"][device] = records
    if device == "mps":
        result["episode_parity"] = episode_parity(model)
        result["mps_memory_bytes"] = {
            "allocated": torch.mps.current_allocated_memory(),
            "driver": torch.mps.driver_allocated_memory(),
        }
    del model
    gc.collect()
    if device == "mps":
        torch.mps.empty_cache()
cpu, mps = result["records"]["cpu"], result["records"]["mps"]
max_difference = 0.0
max_logit_difference = 0.0
agree = head_count = errors = 0
for a, b in zip(cpu, mps):
    assert (a["case_id"], a["repeat"], a["schema_hash"], a["input_hash"]) == (
        b["case_id"],
        b["repeat"],
        b["schema_hash"],
        b["input_hash"],
    )
    if a["error"] or b["error"]:
        errors += 1
        continue
    for name, head in a["heads"].items():
        other = b["heads"][name]
        head_count += 1
        agree += head["selected_label"] == other["selected_label"]
        for label, value in head["probabilities"].items():
            max_difference = max(max_difference, abs(value - other["probabilities"][label]))
            max_logit_difference = max(
                max_logit_difference, abs(head["logits"][label] - other["logits"][label])
            )
summary = {
    "cases": len(cases),
    "head_comparisons": head_count,
    "selected_label_agreements": agree,
    "errors": errors,
    "max_absolute_probability_difference": max_difference,
    "max_absolute_logit_difference": max_logit_difference,
}
for device, records in result["records"].items():
    summary[device] = {
        key: dict(
            zip(
                ["p50", "p95"],
                [float(v) for v in np.quantile([r[key] for r in records], [0.5, 0.95])],
            )
        )
        for key in ["total_ms", "inference_ms"]
    }
summary["median_total_speedup"] = (
    summary["cpu"]["total_ms"]["p50"] / summary["mps"]["total_ms"]["p50"]
)
summary["parity_passed"] = errors == 0 and agree == head_count and max_difference < 1e-3
result["summary"] = summary
Path(args.output).write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(summary, indent=2))