"""Compare a converted Kev package with the merged PyTorch checkpoint."""

from __future__ import annotations

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

import coremltools as ct
import numpy as np
import torch

from assets import load_model
from export_model import KevExport
from preprocessing import Shape, prepare_inputs

COMPUTE_UNITS = {
    "all": ct.ComputeUnit.ALL,
    "cpu": ct.ComputeUnit.CPU_ONLY,
    "cpu-gpu": ct.ComputeUnit.CPU_AND_GPU,
    "cpu-ne": ct.ComputeUnit.CPU_AND_NE,
}


def fixtures() -> list[dict]:
    return [
        {
            "state": "The piece leaves one hole beneath it and creates a small bump on top.",
            "questions": {
                "q": {
                    "type": "choice",
                    "instructions": "Classify the placement.",
                    "criteria": {
                        "clean": "No buried holes and a flat surface",
                        "risky": "Creates a cavity or awkward surface",
                    },
                    "label": "risky",
                    "src": "coreml-fixture",
                }
            },
        },
        {
            "state": "URGENT: verify your account at http://unknown.example and enter your password.",
            "questions": {
                "q": {
                    "type": "noul",
                    "instructions": "Is this message phishing?",
                    "criteria": {"false": "legitimate", "true": "phishing"},
                    "label": True,
                    "src": "coreml-fixture",
                }
            },
        },
        {
            "state": "The customer was charged twice and wants the duplicate transaction reversed.",
            "questions": {
                "q": {
                    "type": "choice",
                    "instructions": "Route this support ticket.",
                    "criteria": {
                        "billing": "Payments and charges",
                        "technical": "Product malfunction",
                        "sales": "Buying a product",
                    },
                    "label": "billing",
                    "src": "coreml-fixture",
                }
            },
        },
        {
            "state": "The order arrived two weeks late and the outer box was damaged.",
            "questions": {
                "q": {
                    "type": "score",
                    "instructions": "Rate the delivery issue severity.",
                    "criteria": ["Low impact", "Moderate impact", "High impact"],
                    "label": 2,
                    "src": "coreml-fixture",
                }
            },
        },
    ]


def suite_requests(path: Path, limit: int) -> list[dict]:
    requests = []
    per_suite: dict[str, int] = {}
    for line in path.read_text().splitlines():
        row = json.loads(line)
        if per_suite.get(row["suite"], 0) >= 2:
            continue
        if row["type"] == "choice":
            question = {
                "type": "choice",
                "instructions": row["instructions"],
                "criteria": {key: description for key, description in row["options"]},
                "label": row["options"][row["gold"]][0],
            }
        else:
            question = {
                "type": "noul",
                "instructions": row["instructions"],
                "criteria": (
                    {key: description for key, description in row["options"]} if row["options"] else None
                ),
                "label": bool(row["gold"]),
            }
        question["src"] = row["suite"]
        requests.append({"state": json.loads(row["state"]), "questions": {"q": question}})
        per_suite[row["suite"]] = per_suite.get(row["suite"], 0) + 1
        if len(requests) == limit:
            break
    return requests


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("package", type=Path)
    parser.add_argument("--length", type=int, default=128)
    parser.add_argument("--max-options", type=int, default=32)
    parser.add_argument("--units", choices=COMPUTE_UNITS, default="all")
    parser.add_argument("--max-probability-error", type=float, default=0.02)
    parser.add_argument("--suite", type=Path)
    parser.add_argument("--suite-cases", type=int, default=20)
    args = parser.parse_args()
    _, tokenizer, decision_model = load_model()
    decision_model.eval()
    shape = Shape(args.length, args.max_options)
    wrapper = KevExport(decision_model, shape.length, shape.max_options).eval()
    started = time.perf_counter()
    coreml = ct.models.MLModel(str(args.package), compute_units=COMPUTE_UNITS[args.units])
    load_seconds = time.perf_counter() - started
    errors: list[float] = []
    times: list[float] = []
    agreements = 0
    evaluated = 0
    requests = fixtures()
    if args.suite:
        requests.extend(suite_requests(args.suite, args.suite_cases))
    for index, request in enumerate(requests):
        try:
            arrays, encoded = prepare_inputs(decision_model, tokenizer, request, shape)
        except ValueError as error:
            print(f"fixture {index}: skipped ({error})")
            continue
        evaluated += 1
        tensors = tuple(torch.from_numpy(value) for value in arrays.values())
        with torch.no_grad():
            _, reference = wrapper(*tensors)
        started = time.perf_counter()
        output = coreml.predict(arrays)
        times.append((time.perf_counter() - started) * 1000)
        options = len(encoded["opt_idx"][0])
        expected = reference[0, :options].numpy()
        actual = np.asarray(output["probabilities"])[0, :options]
        error = float(np.max(np.abs(expected - actual)))
        errors.append(error)
        agreement = int(expected.argmax()) == int(actual.argmax())
        agreements += agreement
        print(f"fixture {index}: options={options} argmax={agreement} max_probability_error={error:.6f}")
    p95 = sorted(times)[max(0, int(0.95 * len(times)) - 1)]
    print(
        f"{agreements}/{evaluated} argmax; max_probability_error={max(errors):.6f}; "
        f"p50={statistics.median(times):.2f} ms; p95={p95:.2f} ms; load={load_seconds:.2f} s"
    )
    if agreements != evaluated or max(errors) > args.max_probability_error:
        raise SystemExit(1)


if __name__ == "__main__":
    main()