#!/usr/bin/env python3
"""Minimal zero-shot inference for ReLU-CLIP. Depends only on numpy, Pillow and a
TFLite runtime -- no torch, no CLIP, no code from this project.

    python example_infer.py cat.jpg
    python example_infer.py cat.jpg --text text/demo_prompt_emb_vitl14.npz

The image tower is a plain int8 TFLite graph; classification is a cosine between its
L2-normalised output and a stored text-embedding matrix. To classify your own classes,
build a new matrix once with sources/make_prompt_emb.py (needs torch + open_clip, on any
machine) and pass it with --text; the device side never needs those.
"""
import argparse
import json
import numpy as np
from PIL import Image

try:                                                          # the lightweight device runtime
    from tflite_runtime.interpreter import Interpreter
except ImportError:
    try:                                                      # its maintained successor
        from ai_edge_litert.interpreter import Interpreter
    except ImportError:                                       # full tensorflow, if that is what you have
        import tensorflow as tf
        Interpreter = tf.lite.Interpreter


def preprocess(path, spec):
    """Resize shortest side to 224 (bicubic), centre-crop 224, scale to [0,1], CLIP-normalise."""
    s = spec["input"]["size"]
    img = Image.open(path).convert("RGB")
    w, h = img.size
    # torchvision Resize(int) truncates the long side -- int(), not round(). One pixel of
    # difference here shifts the centre crop and silently costs accuracy.
    nw, nh = (s, int(s * h / w)) if w <= h else (int(s * w / h), s)
    img = img.resize((nw, nh), Image.BICUBIC)
    # ...and CenterCrop rounds the offset rather than flooring it, so on an odd difference
    # the two disagree by another pixel.
    left, top = int(round((nw - s) / 2.0)), int(round((nh - s) / 2.0))
    img = img.crop((left, top, left + s, top + s))
    x = np.asarray(img, np.float32) / 255.0
    x = (x - spec["input"]["mean"]) / spec["input"]["std"]
    return x[None]                                            # [1, 224, 224, 3] NHWC


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("image")
    ap.add_argument("--model", default="efflite4-vitl14/student_int8.tflite")
    ap.add_argument("--text", default="text/imagenet1k_text_emb_vitl14.npz")
    ap.add_argument("--spec", default="preprocessing.json")
    ap.add_argument("--topk", type=int, default=5)
    a = ap.parse_args()

    spec = json.load(open(a.spec))

    z = np.load(a.text, allow_pickle=True)
    key = "embs" if "embs" in z.files else "text_emb"          # the two shipped files differ
    T, labels = z[key].astype(np.float32), z["labels"]
    T = T / np.linalg.norm(T, axis=1, keepdims=True)           # idempotent; imagenet1k needs it

    itp = Interpreter(a.model)
    itp.allocate_tensors()
    inp, out = itp.get_input_details()[0], itp.get_output_details()[0]

    x = preprocess(a.image, spec)
    s_in, z_in = inp["quantization"]
    itp.set_tensor(inp["index"], np.clip(np.round(x / s_in + z_in), -128, 127).astype(inp["dtype"]))
    itp.invoke()

    s_out, z_out = out["quantization"]
    emb = (itp.get_tensor(out["index"]).astype(np.float32) - z_out) * s_out
    emb /= np.linalg.norm(emb, axis=-1, keepdims=True)         # the graph emits it unnormalised

    sims = (emb @ T.T)[0]
    for i in sims.argsort()[::-1][:a.topk]:
        print(f"{sims[i]:+.4f}  {labels[i]}")


if __name__ == "__main__":
    main()