"""Check native scorer versus fixed K16 Core ML over real upstream demo requests."""

from __future__ import annotations

import argparse
import json
import statistics
import time
from pathlib import Path

import numpy as np

from assets import ROOT, check_base_access
from native_reference import choose, encode_options, pad_batch
from scorer import score_coreml


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--package", type=Path, default=ROOT / "build" / "system_one_gemma_fp16_L256_K16.mlpackage")
    parser.add_argument("--repeats", type=int, default=10)
    args = parser.parse_args()
    check_base_access()

    import coremltools as ct
    import torch

    from export_model import load_trained_scorer

    tokenizer, native = load_trained_scorer()
    coreml = ct.models.MLModel(str(args.package), compute_units=ct.ComputeUnit.ALL)
    demos = json.loads((ROOT / "build" / "source" / "demos.json").read_text())
    rows = []
    for demo in demos:
        sequences = encode_options(tokenizer, demo["state"], demo["question"], demo["options"])
        ids, masks = pad_batch(sequences, tokenizer.pad_token_id)
        with torch.inference_mode():
            native_logits = (
                native(
                    input_ids=torch.tensor(ids, dtype=torch.long),
                    attention_mask=torch.tensor(masks, dtype=torch.long),
                )
                .logits.reshape(-1)
                .float()
                .tolist()
            )
        native_choice = choose(demo["options"], native_logits)
        times = []
        converted = None
        for _ in range(args.repeats):
            start = time.perf_counter()
            converted = score_coreml(coreml, tokenizer, demo["state"], demo["question"], demo["options"])
            times.append((time.perf_counter() - start) * 1000)
        assert converted is not None
        error = max(abs(a - b) for a, b in zip(native_choice["probabilities"], converted["probabilities"]))
        rows.append(
            {
                "demo": demo["name"],
                "candidates": len(demo["options"]),
                "native_selected_index": native_choice["selected_index"],
                "coreml_selected_index": converted["selected_index"],
                "max_calibrated_probability_error": error,
                "median_end_to_end_ms": statistics.median(times),
            }
        )
    report = {
        "package": args.package.name,
        "source": "pinned upstream demos.json",
        "questions": len(rows),
        "selected_id_agreements": sum(r["native_selected_index"] == r["coreml_selected_index"] for r in rows),
        "max_calibrated_probability_error": float(np.max([r["max_calibrated_probability_error"] for r in rows])),
        "rows": rows,
    }
    report["passed"] = (
        report["selected_id_agreements"] == len(rows) and report["max_calibrated_probability_error"] <= 0.02
    )
    target = ROOT / "reports" / "verification-L256-K16.json"
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))
    if not report["passed"]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()