"""Predict with an exported ONNX model. Needs only onnxruntime, numpy, and Pillow (no PyTorch).

python predict_onnx.py --model grid9.int8.onnx --image challenge.png --target bicycles
"""

import argparse
import json
from pathlib import Path

import numpy as np
import onnxruntime as ort
from PIL import Image

from targets import target_name

MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
IMAGE_NAMES = ("image.jpg", "image.jpeg", "image.png", "image.webp")


def preprocess(image, size):
    """Matches torchvision Resize (bilinear, PIL), ToTensor and Normalize."""
    image = image.convert("RGB").resize((size, size), Image.Resampling.BILINEAR)
    return ((np.asarray(image, dtype=np.float32) / 255 - MEAN) / STD).transpose(2, 0, 1)


class GridPredictor:
    def __init__(self, model, threads=None):
        if threads is not None and threads < 1:
            raise ValueError("CPU threads must be positive")
        options = ort.SessionOptions()
        if threads:
            options.intra_op_num_threads = threads
        self.session = ort.InferenceSession(str(model), options, providers=["CPUExecutionProvider"])
        self.info = json.loads(self.session.get_modelmeta().custom_metadata_map["grid_vision"])

    def inputs(self, source):
        """Model input for a challenge image, or a dataset item folder holding image.jpg."""
        source, size, tiles = Path(source), self.info["size"], self.info["tiles"]
        if source.is_dir():
            found = next((source / name for name in IMAGE_NAMES if (source / name).is_file()), None)
            if found is None:
                raise ValueError(f"No image in {source} (expected one of {IMAGE_NAMES})")
            source = found
        with Image.open(source) as opened:
            image = opened.convert("RGB")
        if tiles == 9:
            # Same rounding as training's grid_boxes.
            w, h = image.size
            xs, ys = [round(w * i / 3) for i in range(4)], [round(h * i / 3) for i in range(4)]
            return np.stack(
                [
                    preprocess(image.crop((xs[c], ys[r], xs[c + 1], ys[r + 1])), size)
                    for r in range(3)
                    for c in range(3)
                ]
            )
        return preprocess(image, size)[None]

    def probabilities(self, inputs):
        """Per-class probabilities, shaped (classes, cells)."""
        output = self.session.run(None, {self.session.get_inputs()[0].name: inputs})[0]
        return output.T if self.info["tiles"] == 9 else output[0]

    def predict(self, source, target):
        name = target_name(target)
        if name not in self.info["classes"]:
            raise ValueError(f"Unknown target {name!r}; supported: {self.info['classes']}")
        probabilities = self.probabilities(self.inputs(source))[self.info["classes"].index(name)]
        cutoff = self.info["thresholds"].get(name, self.info["threshold"])
        return {
            "target": name,
            "threshold": cutoff,
            "selected_cells": [int(i) for i in np.flatnonzero(probabilities >= cutoff)],
            "probabilities": probabilities.tolist(),
        }


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--model", required=True)
    parser.add_argument("--image", required=True, help="Challenge image, or a dataset item folder")
    parser.add_argument("--target", required=True)
    parser.add_argument("--threads", type=int, help="CPU threads (default: onnxruntime's choice)")
    args = parser.parse_args()
    print(
        json.dumps(
            GridPredictor(args.model, args.threads).predict(args.image, args.target), indent=2
        )
    )


if __name__ == "__main__":
    main()