"""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.shape_inference import quant_pre_process
from PIL import Image, ImageOps

PHOTOS = 96  # 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):
    """The photographs, fed a stride at a time.

    `__len__` and `set_range` are what `CalibStridedMinMax` asks of a
    reader: the calibrator folds each stride's activations into the running
    range before asking for the next, so memory is one stride's worth and
    not the whole set's.
    """

    STRIDE = 4

    def __init__(self, model_path, photos):
        import onnxruntime as ort

        s = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
        i = s.get_inputs()[0]
        shape = [d if isinstance(d, int) else 1 for d in i.shape]
        name = os.path.basename(model_path)
        self.edge, self.pad, self.norm = preprocessing(name, shape)
        self.name = i.name
        self.photos = photos[: len(photos) - len(photos) % self.STRIDE]
        self.set_range(0, len(self.photos))

    def __len__(self):
        return len(self.photos)

    def set_range(self, start_index, end_index):
        self.items = iter(
            letterbox(Image.open(p), self.edge, self.pad, self.norm)
            for p in self.photos[start_index:end_index]
        )

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


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)
        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"    {os.path.basename(src)}: 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)
        quantize_static(
            pre,
            out,
            Photos(pre, photos),
            quant_format=QuantFormat.QDQ,
            per_channel=True,
            activation_type=QuantType.QUInt8,
            weight_type=QuantType.QInt8,
            # Min/max with a moving average across photographs, so one
            # saturated highlight in one image does not set the range for
            # every activation. Every calibrator keeps each image's whole
            # set of activations until it folds them into a range, which
            # for the 10g detector at 640² is a gigabyte an image and, left
            # to fold once at the end, an OOM kill with no message. The
            # stride folds every four (`CalibMaxIntermediateOutputs` looks
            # like the same thing and is not: in this version it clears
            # without folding). The percentile method has no such bound and
            # is not usable on these graphs.
            calibrate_method=CalibrationMethod.MinMax,
            extra_options={"CalibMovingAverage": True, "CalibStridedMinMax": Photos.STRIDE},
        )
        os.remove(work)
        os.remove(pre)
        print(f"    {out}: {os.path.getsize(out) // 1024} KB")


if __name__ == "__main__":
    main()