"""
Answer the open question directly: is it better to QUANTIZE-THEN-CONVERT or
CONVERT-THEN-QUANTIZE for this model?

Variants compared, all ending in a runnable .mlpackage:
  A fp16                     convert only (baseline)
  B convert -> int8 weights  coremltools.optimize.coreml.linear_quantize_weights
  C convert -> 6-bit palette coremltools.optimize.coreml.palettize_weights (kmeans)
  D int8 weights -> convert  coremltools.optimize.torch PostTrainingQuantizer, then trace+convert

Each is scored on real val faces by agreement with the *unquantized PyTorch*
argmax (the thing users actually see), plus mean IoU against the fp16 baseline's
labels, ANE latency, and on-disk size.

Usage: compare_quant_order.py --weights <best.pt> [--imgsz 512] [--n-val 12]
"""
import argparse, glob, json, os, time

import numpy as np
import torch
from PIL import Image

import coremltools as ct
import coreml_patch  # noqa: F401
from ultralytics import YOLO

VAL = "/Users/ari/FaceSegmentation/dataset_celebamaskhq_semantic/images/val"


class LogitsOnly(torch.nn.Module):
    def __init__(self, m):
        super().__init__()
        self.m = m

    def forward(self, x):
        z = self.m(x)
        return z[0] if isinstance(z, (list, tuple)) else z


def convert(mod, R):
    ts = torch.jit.trace(mod, torch.rand(1, 3, R, R), strict=False)
    return ct.convert(
        ts,
        inputs=[ct.ImageType(name="image", shape=(1, 3, R, R), scale=1 / 255.0,
                             bias=[0, 0, 0], color_layout=ct.colorlayout.RGB)],
        outputs=[ct.TensorType(name="logits")],
        convert_to="mlprogram",
        compute_precision=ct.precision.FLOAT16,
        compute_units=ct.ComputeUnit.CPU_AND_NE,
        minimum_deployment_target=ct.target.iOS17,
    )


def dirsize_mb(p):
    return round(sum(os.path.getsize(f) for f in glob.glob(p + "/**/*", recursive=True)
                     if os.path.isfile(f)) / 1e6, 2)


def score(path, ref_labels, imgs):
    m = ct.models.MLModel(path, compute_units=ct.ComputeUnit.CPU_AND_NE)
    agree, ious = [], []
    for img, ref in zip(imgs, ref_labels):
        got = np.asarray(m.predict({"image": img})["logits"], dtype=np.float32).argmax(1)[0]
        agree.append(float((got == ref).mean()))
        per = []
        for c in np.union1d(np.unique(ref), np.unique(got)):
            inter = np.logical_and(ref == c, got == c).sum()
            union = np.logical_or(ref == c, got == c).sum()
            if union:
                per.append(inter / union)
        ious.append(float(np.mean(per)) if per else 1.0)
    m.predict({"image": imgs[0]})
    t0 = time.time()
    for _ in range(30):
        m.predict({"image": imgs[0]})
    return {"argmax_agreement": round(float(np.mean(agree)), 5),
            "mean_iou_vs_torch": round(float(np.mean(ious)), 5),
            "latency_ms": round((time.time() - t0) / 30 * 1000, 2),
            "size_mb": dirsize_mb(path)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", required=True)
    ap.add_argument("--imgsz", type=int, default=512)
    ap.add_argument("--n-val", type=int, default=12)
    ap.add_argument("--out", default="/Users/ari/FaceSegmentation/exports_semantic/quant_order")
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)
    R = args.imgsz

    files = sorted(glob.glob(os.path.join(VAL, "*.jpg")))[: args.n_val]
    imgs = [Image.open(f).convert("RGB").resize((R, R), Image.BILINEAR) for f in files]

    base_model = LogitsOnly(YOLO(args.weights).model).eval().float().cpu()
    # reference labels from unquantized PyTorch
    ref = []
    for img in imgs:
        x = torch.from_numpy(np.asarray(img, np.float32) / 255.0).permute(2, 0, 1)[None]
        with torch.no_grad():
            ref.append(base_model(x).float().numpy().argmax(1)[0])

    results = {}
    mlbase = convert(base_model, R)
    pA = f"{args.out}/A_fp16.mlpackage"; mlbase.save(pA)
    results["A_fp16_convert_only"] = score(pA, ref, imgs)

    from coremltools.optimize.coreml import (
        OpLinearQuantizerConfig, OpPalettizerConfig, OptimizationConfig,
        linear_quantize_weights, palettize_weights)
    try:
        q = linear_quantize_weights(mlbase, OptimizationConfig(
            global_config=OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int8")))
        p = f"{args.out}/B_convert_then_int8.mlpackage"; q.save(p)
        results["B_convert_then_int8"] = score(p, ref, imgs)
    except Exception as e:
        results["B_convert_then_int8"] = {"error": repr(e)[:200]}
    try:
        q = palettize_weights(mlbase, OptimizationConfig(
            global_config=OpPalettizerConfig(mode="kmeans", nbits=6)))
        p = f"{args.out}/C_convert_then_palette6.mlpackage"; q.save(p)
        results["C_convert_then_palette6"] = score(p, ref, imgs)
    except Exception as e:
        results["C_convert_then_palette6"] = {"error": repr(e)[:200]}

    # D: quantize the TORCH model first, then convert.
    try:
        from coremltools.optimize.torch.quantization import (
            PostTrainingQuantizer, PostTrainingQuantizerConfig)
        tmod = LogitsOnly(YOLO(args.weights).model).eval().float().cpu()
        cfg = PostTrainingQuantizerConfig.from_dict(
            {"global_config": {"weight_dtype": "int8", "granularity": "per_channel"}})
        tq = PostTrainingQuantizer(tmod, cfg).compress()
        p = f"{args.out}/D_int8_then_convert.mlpackage"
        convert(tq.eval(), R).save(p)
        results["D_int8_then_convert"] = score(p, ref, imgs)
    except Exception as e:
        results["D_int8_then_convert"] = {"error": repr(e)[:300]}

    print(json.dumps(results, indent=2))
    with open(f"{args.out}/comparison.json", "w") as f:
        json.dump({"weights": args.weights, "imgsz": R, "results": results}, f, indent=2)

    ok = {k: v for k, v in results.items() if "error" not in v}
    if ok:
        best = max(ok.items(), key=lambda kv: (kv[1]["argmax_agreement"], -kv[1]["size_mb"]))
        print("\nBEST by fidelity:", best[0], best[1])
        small = min(ok.items(), key=lambda kv: kv[1]["size_mb"])
        print("SMALLEST        :", small[0], small[1])


if __name__ == "__main__":
    main()