"""Classify an image on CPU with a packaged FP32, A16W8 or A8W8 TFLite model."""
from pathlib import Path
import argparse
import json

import numpy as np
from PIL import Image
from transformers import AutoImageProcessor
from ai_edge_litert.interpreter import Interpreter


def quantization(detail):
    params = detail["quantization_parameters"]
    if len(params["scales"]) != 1 or len(params["zero_points"]) != 1:
        raise ValueError("This example requires per-tensor input/output quantization")
    scale = float(params["scales"][0])
    if not np.isfinite(scale) or scale <= 0:
        raise ValueError("Invalid input/output quantization scale")
    return scale, int(params["zero_points"][0])


def main():
    directory = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", type=Path, required=True)
    parser.add_argument("--model", type=Path, default=directory / "model_fp32.tflite",
                        help="TFLite model path; defaults to the packaged FP32 model")
    args = parser.parse_args()
    processor = AutoImageProcessor.from_pretrained(directory, backend="pil", local_files_only=True)
    with Image.open(args.image) as image:
        pixels = processor(images=image.convert("RGB"), return_tensors="np")["pixel_values"]
    runtime = Interpreter(model_path=str(args.model), num_threads=4)
    runtime.allocate_tensors()
    input_detail, = runtime.get_input_details()
    output_detail, = runtime.get_output_details()
    pixels = np.ascontiguousarray(pixels, dtype=np.float32)
    if tuple(input_detail["shape"]) != pixels.shape:
        raise ValueError(f"Processor/model shape mismatch: {pixels.shape} vs {input_detail['shape']}")
    if input_detail["dtype"] == np.int8:
        scale, zero = quantization(input_detail)
        pixels = np.clip(np.rint(pixels / scale + zero), -128, 127).astype(np.int8)
    elif input_detail["dtype"] != np.float32:
        raise ValueError("Expected FP32 or INT8 input")
    runtime.set_tensor(input_detail["index"], pixels)
    runtime.invoke()
    logits = runtime.get_tensor(output_detail["index"])[0]
    if output_detail["dtype"] == np.int8:
        scale, zero = quantization(output_detail)
        logits = (logits.astype(np.float32) - zero) * scale
    elif output_detail["dtype"] != np.float32:
        raise ValueError("Expected FP32 or INT8 output")
    if logits.shape != (1000,) or not np.isfinite(logits).all():
        raise ValueError("Expected 1000 finite classifier logits")
    labels = json.loads((directory / "labels.json").read_text())
    for index in np.argsort(-logits, kind="stable")[:5]:
        print(f"{int(index)}\t{labels[str(int(index))]}\tlogit={float(logits[index]):.6f}")


if __name__ == "__main__":
    main()