"""Validate Core AI against full upstream BF16 text and image inference."""

import argparse
import asyncio
import json
import sys
import time
from pathlib import Path

import numpy as np
import torch

from coreai_d3 import D3CoreAI
from validation_cases import CASES


async def main():
    p = argparse.ArgumentParser()
    p.add_argument("name")
    a = p.parse_args()
    root = Path("upstream") / a.name
    bundle = Path("bundles") / (a.name + "-CoreAI")
    sys.path.insert(0, str(root.resolve()))
    from d3_runtime import D3

    torch.set_num_threads(4)
    reference = D3.from_pretrained(root, device="cpu", verify="full")
    engine = await D3CoreAI.load(bundle)
    cases = [(s, q, None, None) for s, q in CASES]
    cases.append(
        (
            "refund " * 350,
            {
                "type": "choice",
                "instructions": "What topic?",
                "criteria": {"refund": None, "shipping": None},
            },
            None,
            None,
        )
    )
    cases.append(
        (
            "A receipt.",
            {
                "type": "choice",
                "instructions": "Which item appears on the receipt?",
                "criteria": {"blender": None, "laptop": None, "bicycle": None},
            },
            [root / "assets/example-receipt.png"],
            None,
        )
    )
    cases.append(
        (
            "A receipt.",
            {"type": "noul", "instructions": "Does this receipt list a blender?"},
            [root / "assets/example-receipt.png"],
            None,
        )
    )
    cases.append(
        (
            "Two receipts.",
            {"type": "noul", "instructions": "Do the receipts list a blender?"},
            [root / "assets/example-receipt.png"] * 2,
            None,
        )
    )
    cases.append(
        (
            "A camera clip.",
            {
                "type": "choice",
                "instructions": "Which way does the square move?",
                "criteria": {
                    "right": "From left to right",
                    "left": "From right to left",
                    "still": "It does not move",
                },
            },
            None,
            [root / "assets/example-video.mp4"],
        )
    )
    records = []
    for state, q, images, videos in cases:
        start = time.monotonic()
        logits, keys, encoded = await engine.logits(state, q, images, videos)
        actual = torch.tensor(logits).softmax(-1).numpy()
        if images or videos:
            # Compare identical resized pixels, prompts, positions and the source BF16 model.
            inputs = {
                k: torch.from_numpy(np.asarray(v))
                for k, v in encoded.items()
                if k
                in [
                    "input_ids",
                    "attention_mask",
                    "pixel_values",
                    "pixel_values_videos",
                    "image_grid_thw",
                    "video_grid_thw",
                    "mm_token_type_ids",
                ]
            }
            with torch.inference_mode():
                h = reference.backbone(**inputs, use_cache=False).last_hidden_state[
                    :, -1
                ]
                scores = (h.float() @ reference.readout.T)[0, : len(keys)]
                expected = scores.softmax(-1).numpy()
        else:
            seq = encoded["input_ids"][0].tolist()
            expected = np.array(reference.probabilities([seq], [len(keys)])[0])
        r = {
            "type": q["type"],
            "modality": "video"
            if videos
            else "images"
            if images and len(images) > 1
            else "image"
            if images
            else "text",
            "tokens": len(encoded["input_ids"][0]),
            "options": len(keys),
            "coreai_choice": int(actual.argmax()),
            "upstream_choice": int(expected.argmax()),
            "max_probability_error": float(abs(expected - actual).max()),
            "total_seconds": time.monotonic() - start,
        }
        print(json.dumps(r), flush=True)
        records.append(r)
        if not np.isfinite(actual).all():
            raise AssertionError("Nonfinite Core AI probabilities")
        if r["coreai_choice"] != r["upstream_choice"]:
            raise AssertionError(r)
    (bundle / "validation.json").write_text(json.dumps(records, indent=2) + "\n")
    try:
        await engine.logits("refund " * 1000, CASES[0][1])
    except ValueError:
        pass
    else:
        raise AssertionError("Oversized request was accepted")
    print(
        "PASS: selected answers agree; probability drift is recorded, not treated as strict parity",
        flush=True,
    )


if __name__ == "__main__":
    asyncio.run(main())