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

import glob
import os
import sys
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
from onnx import version_converter
from onnxruntime.quantization import CalibrationDataReader, CalibrationMethod, QuantType, quantize_static
from onnxruntime.quantization.calibrate import create_calibrator, save_tensors_data
from onnxruntime.quantization.execution_providers.qnn import get_qnn_qdq_config
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import htp_graph  # noqa: E402

MODELS = Path(__file__).resolve().parents[1] / "models"
PHOTOS = 300  # calibration photographs; face crops and keypoints come from fewer
CHUNK = 4  # inputs whose activations are held at once (scrfd_10g: ~1 GB each)

Q = QuantType
FORMS = {"int8": (Q.QUInt8, Q.QInt8), "a16w8": (Q.QUInt16, Q.QInt8), "a16w16": (Q.QUInt16, Q.QInt16)}

# Per model: where it lives, the form `Rung::form` gives its role, the exact
# rewrites its graph needs, and nodes that stay float on the CPU because one
# scale cannot serve the tensor (the segmenter's rows: boxes in pixels beside
# scores in 0..1) or because the HTP's 16-bit arithmetic drifts there (the
# scene model's attention). Measured, inference.md §1.5.
TABLE = {
    "scrfd_500m_640": dict(dir="face", form="a16w8", feed="scrfd"),
    "scrfd_2.5g_640": dict(dir="face", form="a16w8", feed="scrfd"),
    "scrfd_10g_640": dict(dir="face", form="a16w8", feed="scrfd"),
    "2d106det_b1": dict(dir="face", form="a16w8", feed="landmarks"),
    "yolo26n-seg": dict(dir="segment", form="a16w16", feed="yolo", float_from="/model.23/Concat_4"),
    "yolo26s-sem-ade20k": dict(dir="scene", form="a16w16", feed="yolo",
                               float_nodes=["/model.10/m/m.0/attn/MatMul", "/model.10/m/m.0/attn/Softmax",
                                            "/model.10/m/m.0/attn/MatMul_1"]),
    "migan-512": dict(dir="inpaint", form="a16w16", feed="migan"),
    "xfeat-1024": dict(dir="keypoints", form="int8", feed="xfeat", rewrites=["unfold", "resize"]),
    "xfeat-768": dict(dir="keypoints", form="int8", feed="xfeat", rewrites=["unfold", "resize"]),
    "mosaic-1408": dict(dir="denoise", form="a16w16", feed=None, rewrites=["bayer"]),
}


# ---- the app's samplers (dr-face Letterbox, dr-segment Letterbox::sample, align.rs) ----

def load(p):
    return np.asarray(Image.open(p).convert("RGB"), np.float32) / 255.0


def bilinear(img, sx, sy):
    h, w = img.shape[:2]
    x0, y0 = np.floor(sx).astype(np.int64), np.floor(sy).astype(np.int64)
    fx, fy = (sx - x0)[..., None], (sy - y0)[..., None]
    c = lambda a, n: np.clip(a, 0, n - 1)  # noqa: E731 - neighbours clamp at the edge
    top = img[c(y0, h), c(x0, w)] * (1 - fx) + img[c(y0, h), c(x0 + 1, w)] * fx
    bot = img[c(y0 + 1, h), c(x0, w)] * (1 - fx) + img[c(y0 + 1, h), c(x0 + 1, w)] * fx
    return top * (1 - fy) + bot * fy


def chw(x):
    return np.ascontiguousarray(x.transpose(2, 0, 1))[None].astype(np.float32)


def letterbox(img, edge, yolo):
    h, w = img.shape[:2]
    s = min(edge / w, edge / h)
    px, py = (edge - w * s) / 2, (edge - h * s) / 2
    ix, iy = np.meshgrid(np.arange(edge) + 0.5, np.arange(edge) + 0.5)
    if yolo:  # semantic.rs: no -0.5, pad 0.5, 0..1
        sx, sy = (ix - px) / s, (iy - py) / s
        out = bilinear(img, sx, sy)
        out[(sx < 0) | (sx >= w) | (sy < 0) | (sy >= h)] = 0.5
    else:  # detect.rs: -0.5, pad 114, (v·255 − 127.5)/128
        sx, sy = (ix - px) / s - 0.5, (iy - py) / s - 0.5
        out = (bilinear(img, sx, sy) * 255 - 127.5) / 128
        out[(sx < -0.5) | (sx > w - 0.5) | (sy < -0.5) | (sy > h - 0.5)] = (114 - 127.5) / 128
    return chw(out), (s, px, py)


def crop_box(img, x0, y0, bw, bh, ow, oh):
    """align.rs crop_box: output (u+.5) → source, −0.5, bilinear, outside black."""
    u, v = np.meshgrid(np.arange(ow) + 0.5, np.arange(oh) + 0.5)
    sx, sy = x0 + u * bw / ow - 0.5, y0 + v * bh / oh - 0.5
    out = bilinear(img, sx, sy)
    h, w = img.shape[:2]
    out[(sx < -1) | (sx > w) | (sy < -1) | (sy > h)] = 0
    return out


def scrfd_boxes(outs, s, px, py):
    """detect.rs decode: score ≥ 0.5, greedy NMS at 0.4, min side 24 px."""
    fmc = len(outs) // 3
    boxes, scores = [], []
    for i, st in enumerate([8, 16, 32, 64][:fmc]):
        sc, bx = outs[i].reshape(-1), outs[fmc + i].reshape(-1, 4)
        idx = np.nonzero(sc >= 0.5)[0]
        cell = idx // 2
        cx, cy = (cell % (640 // st)) * st, (cell // (640 // st)) * st
        boxes.append(np.stack([cx - bx[idx, 0] * st, cy - bx[idx, 1] * st, cx + bx[idx, 2] * st, cy + bx[idx, 3] * st], 1))
        scores.append(sc[idx])
    b, sc = np.concatenate(boxes), np.concatenate(scores)
    keep = []
    for i in np.argsort(-sc):
        x0 = np.maximum(b[i, 0], b[keep, 0]); y0 = np.maximum(b[i, 1], b[keep, 1])
        x1 = np.minimum(b[i, 2], b[keep, 2]); y1 = np.minimum(b[i, 3], b[keep, 3])
        inter = np.clip(x1 - x0, 0, None) * np.clip(y1 - y0, 0, None)
        area = lambda r: (r[..., 2] - r[..., 0]) * (r[..., 3] - r[..., 1])  # noqa: E731
        if not keep or (inter / (area(b[i]) + area(b[keep]) - inter)).max() <= 0.4:
            keep.append(i)
    b = (b[keep] - [px, py, px, py]) / s
    return b[np.minimum(b[:, 2] - b[:, 0], b[:, 3] - b[:, 1]) >= 32]


# ---- one calibration input per photograph (or per face), as the app makes it ----

def feeds(kind, photos, model):
    name = model.get_inputs()[0].name
    if kind == "scrfd":
        for p in photos:
            yield {name: letterbox(load(p), 640, False)[0]}
    elif kind == "landmarks":  # landmarks.rs: 1.5× the box, square, 0..255
        det = ort.InferenceSession(str(MODELS / "face/scrfd_10g_640.onnx"), providers=["CPUExecutionProvider"])
        for p in photos:
            img = load(p)
            x, ctx = letterbox(img, 640, False)
            for b in scrfd_boxes(det.run(None, {"input.1": x}), *ctx):
                cx, cy, side = (b[0] + b[2]) / 2, (b[1] + b[3]) / 2, 1.5 * max(b[2] - b[0], b[3] - b[1])
                yield {name: chw(crop_box(img, cx - side / 2, cy - side / 2, side, side, 192, 192) * 255)}
    elif kind == "yolo":
        for p in photos:
            yield {name: letterbox(load(p), 640, True)[0]}
    elif kind == "migan":  # migan.rs: ch0 = known − 0.5, ch1–3 = (rgb·2 − 1)·known
        rng = np.random.default_rng(7)
        for p in photos:
            img = load(p)
            h, w = img.shape[:2]
            e = min(h, w)
            sq = Image.fromarray((img[(h - e) // 2:(h + e) // 2, (w - e) // 2:(w + e) // 2] * 255).astype(np.uint8))
            img = np.asarray(sq.resize((512, 512), Image.BILINEAR), np.float32) / 255
            known = np.ones((512, 512), np.float32)
            for _ in range(rng.integers(1, 3)):  # a panorama's unknown border: a wedge along one edge
                side = rng.integers(4)
                depth = np.linspace(rng.integers(20, 110), rng.integers(20, 110), 512).astype(int)
                edge = np.arange(512)[:, None] < depth[None, :]  # [depth, along]: inside the wedge
                wedge = edge if side % 2 == 0 else edge[::-1]  # top / bottom of a column
                known[wedge if side < 2 else wedge.T] = 0  # or left / right of a row
            x = np.concatenate([(known - 0.5)[None], ((img * 2 - 1) * known[..., None]).transpose(2, 0, 1)])[None]
            yield {name: x.astype(np.float32)}
    elif kind == "xfeat":  # xfeat.rs: grey 0..1, shrink to fit, top-left, zero pad
        _, _, H, W = [d if isinstance(d, int) else 1 for d in model.get_inputs()[0].shape]
        for p in photos:
            g = load(p).mean(2)  # a display-rendered photograph is already the app's (R+G+B)/3 ^ 1/2.2
            h, w = g.shape
            s = min(W / w, H / h, 1.0)
            if s < 1:
                g = np.asarray(Image.fromarray(g).resize((round(w * s), round(h * s)), Image.BOX))
            pad = np.zeros((H, W), np.float32)
            pad[:g.shape[0], :g.shape[1]] = g
            yield {name: pad[None, None]}


class Items(CalibrationDataReader):
    def __init__(self, items):
        self.it = iter(items)

    def get_next(self):
        return next(self.it, None)


def calibrate(path, items, cache):
    """Min/max ranges in chunks — every ORT calibrator holds all activations
    until it folds them, and the others measurably degrade the result."""
    cal = create_calibrator(Path(path), None, augmented_model_path=f"{path}.aug.onnx",
                            calibrate_method=CalibrationMethod.MinMax)
    batch, n = [], 0
    for item in items:
        batch.append(item)
        n += 1
        if len(batch) == CHUNK:
            cal.collect_data(Items(batch))
            batch = []
    if batch:
        cal.collect_data(Items(batch))
    save_tensors_data(cal.compute_data(), cache)
    os.remove(f"{path}.aug.onnx")
    return n


def downstream(m, start):
    names, live = set(), set()
    for n in m.graph.node:
        if n.name == start or any(i in live for i in n.input):
            names.add(n.name)
            live.update(n.output)
    return sorted(names)


def main():
    args = sys.argv[1:]
    ranges = None
    if args[:1] == ["--ranges"]:
        ranges, args = args[1], args[2:]
        photo_dir = None
    else:
        photo_dir, args = args[0], args[1:]
    names = args or [n for n in TABLE if TABLE[n]["feed"]]
    photos = []
    if photo_dir:
        photos = sorted(p for e in ("jpg", "jpeg", "JPG", "JPEG", "png")
                        for p in glob.glob(os.path.join(photo_dir, "**", f"*.{e}"), recursive=True))[:PHOTOS]
        if len(photos) < 50:
            sys.exit(f"only {len(photos)} photographs under {photo_dir}; calibration wants hundreds")
        print(f"==> calibrating on {len(photos)} photographs")

    for stem in names:
        spec = TABLE[stem]
        src = MODELS / spec["dir"] / f"{stem}.onnx"
        out = src.with_name(f"{stem}.{spec['form']}.onnx")
        work = src.with_name(f"{stem}.quant-work.onnx")
        print(f"  {stem} -> {out.name}")
        m = onnx.load(src)
        if next(o.version for o in m.opset_import if o.domain in ("", "ai.onnx")) < 13:
            m = version_converter.convert_version(m, 17)  # per-channel QDQ needs 13
            m.ir_version = 8
        onnx.save(m, work)
        if spec.get("rewrites"):
            onnx.save(htp_graph.rewrite(str(work), spec["rewrites"]), work)
        model = ort.InferenceSession(str(work), providers=["CPUExecutionProvider"])
        cache = str(work) + ".ranges"
        if spec["feed"] is None:
            if not ranges:
                sys.exit(f"{stem} is calibrated on mosaics, not photographs: pass --ranges")
            cache = ranges
        else:
            n = calibrate(str(work), feeds(spec["feed"], photos, model), cache)
            print(f"    {n} calibration inputs")
        act, wt = FORMS[spec["form"]]
        # The config needs a reader only to exist; the ranges come from `cache`.
        zeros = {i.name: np.zeros([d if isinstance(d, int) else 1 for d in i.shape], np.float32)
                 for i in model.get_inputs()}
        cfg = get_qnn_qdq_config(str(work), Items([zeros]), activation_type=act, weight_type=wt, per_channel=True)
        exclude = list(cfg.nodes_to_exclude or []) + spec.get("float_nodes", [])
        if spec.get("float_from"):
            exclude += downstream(onnx.load(work), spec["float_from"])
        quantize_static(str(work), str(out), None, quant_format=cfg.quant_format,
                        op_types_to_quantize=cfg.op_types_to_quantize, per_channel=True,
                        activation_type=act, weight_type=wt, nodes_to_exclude=exclude,
                        calibrate_method=CalibrationMethod.MinMax, extra_options=cfg.extra_options,
                        calibration_cache_path=cache)
        os.remove(work)
        if cache != ranges:
            os.remove(cache)
        print(f"    {out.name}: {out.stat().st_size // 1024} KB")


if __name__ == "__main__":
    main()