"""The body of quantise-models.sh; see there. Run through it, not directly."""

import glob
import os
import sys

import numpy as np
import onnx
from onnx import version_converter
from onnxruntime.quantization import (
    CalibrationDataReader,
    CalibrationMethod,
    QuantFormat,
    QuantType,
    quantize_static,
)
from onnxruntime.quantization.calibrate import create_calibrator
from onnxruntime.quantization.calibrate import save_tensors_data
from onnxruntime.quantization.shape_inference import quant_pre_process
from pathlib import Path
from PIL import Image, ImageOps

PHOTOS = 64  # enough for a stable range; more only costs time


def letterbox(img, edge, pad, norm):
    """The app's Letterbox::sample: fit the long side to `edge`, centre, pad."""
    img = ImageOps.exif_transpose(img).convert("RGB")
    w, h = img.size
    scale = edge / max(w, h)
    nw, nh = max(1, round(w * scale)), max(1, round(h * scale))
    img = img.resize((nw, nh), Image.BILINEAR)
    canvas = Image.new("RGB", (edge, edge), (pad, pad, pad))
    canvas.paste(img, ((edge - nw) // 2, (edge - nh) // 2))
    x = np.asarray(canvas, dtype=np.float32)  # HWC, 0..255
    x = norm(x)
    return np.ascontiguousarray(x.transpose(2, 0, 1))[None]  # NCHW


def preprocessing(name, shape):
    """Which normalisation this model is fed in the app.

    SCRFD (`dr-face::detect`): `(v - 127.5) / 128`, padded with 114.
    ArcFace (`dr-face::embed`): the same, on an aligned 112 crop — a
    letterboxed photograph is the wrong distribution, but the embedder is
    never quantised (§7), so this is only ever a fallback.
    YOLO (`dr-segment`): `v / 255`, padded with 0.5.
    """
    edge = shape[-1]
    if name.startswith("scrfd") or name.startswith("arcface"):
        return edge, 114, lambda x: (x - 127.5) / 128.0
    return edge, 128, lambda x: x / 255.0


class Photos(CalibrationDataReader):
    """One chunk of photographs, fed as the app would feed them."""

    def __init__(self, input_name, paths, edge, pad, norm):
        self.name = input_name
        self.items = iter(letterbox(Image.open(p), edge, pad, norm) for p in paths)

    def get_next(self):
        x = next(self.items, None)
        return None if x is None else {self.name: x}


# Photographs whose activations are held in memory at once. Every ONNX
# Runtime calibrator keeps each image's whole set of activations until it
# folds them into a range — a gigabyte an image on the 10g detector at 640²,
# and folded once at the end, an OOM kill with no message. Folding every
# `CHUNK` images gives ranges identical to folding once (checked on
# scrfd_500m, 129 tensors, no difference) at a bounded cost.
CHUNK = 4


def calibrate(pre, name, photos, cache):
    """Min/max ranges over `photos`, written to `cache` for quantize_static.

    Plain min/max: the moving average and the strided option of
    `quantize_static` both measured worse than this on held-out proxies, and
    the percentile method has no memory bound at all.
    """
    import onnxruntime as ort

    s = ort.InferenceSession(pre, providers=["CPUExecutionProvider"])
    i = s.get_inputs()[0]
    shape = [d if isinstance(d, int) else 1 for d in i.shape]
    edge, pad, norm = preprocessing(name, shape)
    calibrator = create_calibrator(
        Path(pre),
        None,
        augmented_model_path=f"{pre}.augmented.onnx",
        calibrate_method=CalibrationMethod.MinMax,
    )
    for start in range(0, len(photos), CHUNK):
        calibrator.collect_data(Photos(i.name, photos[start : start + CHUNK], edge, pad, norm))
    ranges = calibrator.compute_data()
    save_tensors_data(ranges, cache)
    os.remove(f"{pre}.augmented.onnx")


def main():
    photo_dir, models = sys.argv[1], sys.argv[2:]
    photos = sorted(
        p
        for ext in ("jpg", "jpeg", "JPG", "JPEG", "png")
        for p in glob.glob(os.path.join(photo_dir, "**", f"*.{ext}"), recursive=True)
    )[:PHOTOS]
    if len(photos) < 20:
        sys.exit(f"only {len(photos)} photographs under {photo_dir}; calibration wants dozens")
    print(f"==> calibrating on {len(photos)} photographs")

    for src in models:
        stem, _ = os.path.splitext(src)
        name = os.path.basename(src)
        out = f"{stem}.int8.onnx"
        m = onnx.load(src)
        opset = next((o.version for o in m.opset_import if o.domain in ("", "ai.onnx")), 0)
        work = f"{stem}.quant-work.onnx"
        if opset < 13:
            print(f"    {name}: opset {opset} -> 17")
            m = version_converter.convert_version(m, 17)
            m.ir_version = 8
        onnx.save(m, work)
        pre = f"{stem}.quant-pre.onnx"
        quant_pre_process(work, pre)
        cache = f"{stem}.quant-ranges.json"
        calibrate(pre, name, photos, cache)
        quantize_static(
            pre,
            out,
            None,
            quant_format=QuantFormat.QDQ,
            per_channel=True,
            activation_type=QuantType.QUInt8,
            weight_type=QuantType.QInt8,
            calibrate_method=CalibrationMethod.MinMax,
            calibration_cache_path=cache,
        )
        for f in (work, pre, cache):
            os.remove(f)
        print(f"    {out}: {os.path.getsize(out) // 1024} KB")


if __name__ == "__main__":
    main()