import asyncio
import json
import shutil
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor

from coreai_embeddinggemma2 import EmbeddingGemma2CoreAI


def cases():
    wave = (0.2 * np.sin(2 * np.pi * 440 * np.arange(16000) / 16000)).astype(np.float32)
    image = Image.new("RGB", (320, 192), (20, 80, 210))
    gradient = Image.fromarray(
        np.tile(np.arange(256, dtype=np.uint8)[None, :, None], (128, 1, 3))
    )
    frames = np.stack(
        [
            np.asarray(image.resize((128, 96))),
            np.asarray(Image.new("RGB", (128, 96), (210, 50, 20))),
        ]
    )
    return [
        (
            "query",
            {"text": "What causes the northern lights?", "prompt_name": "SearchQuery"},
        ),
        (
            "document",
            {
                "text": "The northern lights are caused by charged particles from the sun.",
                "prompt_name": "Document",
            },
        ),
        ("unprefixed", {"text": "An apple a day keeps the doctor away."}),
        (
            "code",
            {"text": "def add(a, b): return a + b", "prompt_name": "CodeRetrieval"},
        ),
        ("spanish", {"text": "La biblioteca está abierta hasta las ocho."}),
        ("japanese", {"text": "桜の花が春に咲きます。"}),
        ("short", {"text": "cat"}),
        ("long", {"text": "The oceans contain diverse ecosystems. " * 90}),
        ("near-capacity", {"text": "The oceans contain diverse ecosystems. " * 160}),
        ("image", {"images": image}),
        ("gradient", {"images": gradient}),
        ("image-text", {"text": "<|image|> A blue rectangle.", "images": image}),
        ("two-images", {"images": [image, gradient]}),
        ("audio", {"audio": wave}),
        ("audio-text", {"text": "<|audio|> A musical tone.", "audio": wave}),
        ("audio-short", {"audio": wave[:3200]}),
        ("audio-long", {"audio": np.tile(wave, 5)}),
        ("audio-limit", {"audio": np.pad(np.tile(wave, 11), (0, 3200))}),
        ("video", {"videos": frames}),
        ("video-text", {"text": "<|video|> Colors changing.", "videos": frames}),
        (
            "mixed",
            {
                "text": "<|image|> Blue. <|audio|> Tone. <|video|> Colors.",
                "images": image,
                "audio": wave,
                "videos": frames,
            },
        ),
    ]


async def main():
    torch.set_num_threads(4)
    root = Path("bundles/embeddinggemma-2-CoreAI")
    source = Path("upstream/embeddinggemma-2")
    for f in source.iterdir():
        if f.is_file() and f.suffix in [".json", ".jinja", ".model"]:
            shutil.copy2(f, root / f.name)
    m = AutoModel.from_pretrained(
        source,
        dtype=torch.float32,
        attn_implementation={
            "language_model": "sdpa",
            "vision_tower": "sdpa",
            "audio_tower": "eager",
        },
    ).eval()
    p = AutoProcessor.from_pretrained(source)
    print("Loading Core AI", flush=True)
    runtime = await EmbeddingGemma2CoreAI.load(root)
    print("Loaded Core AI", flush=True)
    records = []
    for name, args in cases():
        data = dict(args)
        prompt = data.pop("prompt_name", None)
        if prompt:
            data["text"] = runtime.prompts[prompt] + data["text"]
        if "audio" in data:
            data["sampling_rate"] = 16000
        if "videos" in data:
            data["videos_kwargs"] = {"do_sample_frames": False}
        e = p(**data, return_tensors="pt")
        with torch.no_grad():
            h = m(**e).last_hidden_state
            mask = e["attention_mask"][:, :, None].float()
            reference = torch.nn.functional.normalize(
                (h * mask).sum(1) / mask.sum(1), dim=-1
            )[0].numpy()
        actual = await runtime.encode_processed({k: v.numpy() for k, v in e.items()})
        record = {
            "name": name,
            "tokens": e["input_ids"].shape[1],
            "max_abs_error": float(np.max(np.abs(actual - reference))),
            "cosine": float(actual @ reference),
            "finite": bool(np.isfinite(actual).all()),
            "unit_norm": float(np.linalg.norm(actual)),
        }
        print(record, flush=True)
        records.append(record)
        assert (
            record["finite"]
            and record["cosine"] > 0.99999
            and record["max_abs_error"] < 0.00005
        ), record
    for dim in [128, 256, 512]:
        actual = await runtime.encode("The sky is blue.", truncate_dim=dim)
        assert actual.shape == (dim,) and abs(np.linalg.norm(actual) - 1) < 1e-5
    try:
        await runtime.encode("word " * 3000)
    except ValueError:
        pass
    else:
        raise AssertionError("Over-capacity input accepted")
    (root / "validation.json").write_text(
        json.dumps(
            {
                "thresholds": {"min_cosine": 0.99999, "max_abs_error": 0.00005},
                "cases": records,
                "truncation_dims": [128, 256, 512],
                "oversized_input_rejected": True,
            },
            indent=2,
        )
        + "\n"
    )


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