"""Run local text, image and video inference checks for one quantization."""
import argparse
import json
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
from PIL import Image

from metaencoder_mlx import MetaEncoderMLX


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", type=Path, required=True)
    parser.add_argument("--report", type=Path, required=True)
    args = parser.parse_args()
    start = time.monotonic()
    model = MetaEncoderMLX(args.model)
    report = {"model": args.model.name, "load_seconds": time.monotonic()-start,
              "embedding_dim": model.embedding_dim, "checks": []}
    red = Image.new("RGB", (56, 56), (255, 0, 0))
    cases = [
        ("text_animal", "Which animal says meow?", ["cat", "dog", "horse"]),
        ("text_action", "Select the action to cancel a pending order.",
         ["Click Cancel order", "Open profile settings", "Search products"]),
        ("image_color", {"image": red, "text": "What is the main color of this image?"},
         ["red", "blue", "green"]),
        ("video_color", {"video": [red, red], "text": "What is the main color of this video?"},
         ["red", "blue", "green"]),
    ]
    vectors = []
    for name, task, candidates in cases:
        started = time.monotonic()
        embeddings = model.encode([task, *candidates])
        array = np.asarray(embeddings)
        assert array.shape == (4, model.embedding_dim), array.shape
        assert np.isfinite(array).all(), name
        norms = np.linalg.norm(array, axis=-1)
        assert np.allclose(norms, 1.0, atol=1e-5), norms
        scores = array[1:] @ array[0]
        predicted = int(scores.argmax())
        report["checks"].append({"name": name, "seconds": time.monotonic()-started,
                                 "scores": scores.tolist(), "candidates": candidates,
                                 "predicted": candidates[predicted], "expected": candidates[0],
                                 "correct": predicted == 0, "norms": norms.tolist()})
        vectors.append(array)
        print(json.dumps(report["checks"][-1]), flush=True)
    report["total_seconds"] = time.monotonic()-start
    report["peak_mlx_memory_gb"] = mx.get_peak_memory()/1e9
    args.report.write_text(json.dumps(report, indent=2) + "\n")
    np.savez(args.report.with_suffix(".npz"), **{case[0]: vec for case, vec in zip(cases, vectors)})


if __name__ == "__main__":
    main()