"""Small MPS compatibility probe. Concurrent CPU evaluation means timings are provisional."""

import json
import os
import time
from pathlib import Path

os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "0"
import torch
from gliner2.classification import ClassificationSchema

from sealcell.model import LocalModel, validate_head

capabilities = {
    "torch": torch.__version__,
    "mps_built": torch.backends.mps.is_built(),
    "mps_available": torch.backends.mps.is_available(),
    "fallback_enabled": False,
}
print(json.dumps(capabilities), flush=True)
if not capabilities["mps_available"]:
    raise SystemExit("MPS unavailable")
model = LocalModel()
source = json.loads(Path("data/fixtures/playground-probes.json").read_text())
schema = ClassificationSchema()
for name, labels in source["schema"].items():
    schema.single(name, labels)
model.classifier.to(device="mps", dtype=torch.float32).eval()
records = []
for case in source["cases"]:
    torch.mps.synchronize()
    start = time.perf_counter()
    try:
        result = model.classifier.classify(case["text"], schema)
        torch.mps.synchronize()
        elapsed = (time.perf_counter() - start) * 1000
        heads = {
            name: {
                "selected_label": result.value(name),
                "probabilities": dict(result.probabilities(name)),
            }
            for name in source["schema"]
        }
        for name, head in heads.items():
            validate_head(head["selected_label"], head["probabilities"], source["schema"][name])
        record = {
            "fixture_id": case["id"],
            "device": "mps",
            "dtype": "float32",
            "heads": heads,
            "total_ms": elapsed,
            "error": None,
        }
    except Exception as exc:
        record = {
            "fixture_id": case["id"],
            "device": "mps",
            "error": f"{type(exc).__name__}: {exc}",
        }
    records.append(record)
    print(case["id"], record.get("total_ms"), record["error"], flush=True)
Path("reports/mps-compatibility.json").write_text(
    json.dumps(
        {
            "capabilities": capabilities,
            "model_revision": model.metadata["revision"],
            "timing_context": "provisional; CPU evaluation running concurrently; first shape includes compilation",
            "records": records,
        },
        indent=2,
    )
    + "\n"
)