#!/usr/bin/env python3
"""Compare synthetic PyTorch and Core ML CPU/GPU classifier probabilities."""

import json
from pathlib import Path

import coremltools as ct
import numpy as np
import torch
from transformers import CLIPVisionModelWithProjection

from convert import Classifier, LABELS


def main():
    torch.set_num_threads(4)
    rng = np.random.default_rng(20260927)
    # RGB-like solid, gradient, and noise, after the CLIP image normalization.
    mean = np.array([0.48145466, 0.4578275, 0.40821073], np.float32)[None, :, None, None]
    std = np.array([0.26862954, 0.26130258, 0.27577711], np.float32)[None, :, None, None]
    rgb = [
        np.full((1, 3, 224, 224), 0.5, np.float32),
        np.broadcast_to(np.linspace(0, 1, 224, dtype=np.float32)[None, None, None, :], (1, 3, 224, 224)).copy(),
        rng.random((1, 3, 224, 224), dtype=np.float32),
    ]
    samples = [(x - mean) / std for x in rgb]
    vision = CLIPVisionModelWithProjection.from_pretrained(".source", local_files_only=True).eval()
    proto = torch.from_numpy(np.load("text_prototypes.npy"))
    reference = Classifier(vision, proto).eval()
    with torch.inference_mode():
        expected = [reference(torch.from_numpy(x)).numpy() for x in samples]
    del vision, reference, proto

    results = {"source_revision": "3d74acf9a28c67741b2f4f2ea7635f0aaf6f0268", "labels": [x[0] for x in LABELS], "cases": []}
    for unit_name, unit in [("cpu", ct.ComputeUnit.CPU_ONLY), ("cpu_gpu", ct.ComputeUnit.CPU_AND_GPU)]:
        core = ct.models.MLModel("RegionClassifier.mlpackage", compute_units=unit)
        for index, (sample, want) in enumerate(zip(samples, expected)):
            got = core.predict({"image": sample})["probabilities"]
            if got.shape != (1, 12) or got.dtype != np.float32:
                raise AssertionError((got.shape, got.dtype))
            max_error = float(np.max(np.abs(got - want)))
            sum_error = float(abs(float(got.sum()) - 1.0))
            same_label = int(got.argmax()) == int(want.argmax())
            case = {"compute_unit": unit_name, "sample": index, "max_absolute_probability_error": max_error,
                    "probability_sum_error": sum_error, "same_top_label": same_label,
                    "pytorch_top_label": LABELS[int(want.argmax())][0], "coreml_top_label": LABELS[int(got.argmax())][0]}
            results["cases"].append(case)
            print(case, flush=True)
            if max_error > 0.03 or sum_error > 0.001 or not same_label:
                raise AssertionError(case)
    Path("validation.json").write_text(json.dumps(results, indent=2) + "\n")
    print("VALIDATED")


if __name__ == "__main__":
    main()