"""
Export a trained YOLO26 semantic checkpoint to Core ML and pick the best
quantization variant.

Path: manual trace -> coremltools (NOT `yolo export`, whose baked argmax breaks
MIL; and coremltools 9.0 x numpy 2.x needs coreml_patch). Output contract is
locked to the app: input "image" (512x512 RGB), output "logits" fp16
(1,19,64,64); computeUnits .cpuAndNeuralEngine.

Variants: fp16 baseline, int8 weight-only linear, 6-bit palettization.
Each is parity-checked against PyTorch on real val images (argmax agreement)
and timed on the ANE. Writes a summary + copies the recommended variant to
--app-dest as FaceSegModel.mlpackage (same name -> Xcode swap is automatic).

Usage:
  export_coreml.py --weights /Users/ari/FaceSegmentation/runs_semantic/celeba_large/weights/best.pt \
      --tag large [--app-dest .../facesegmentation/FaceSegModel.mlpackage]
"""
import argparse, glob, json, os, shutil, 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_IMAGES = "/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 torch_logits(wrap, img):
    x = torch.from_numpy(np.asarray(img, dtype=np.float32) / 255.0).permute(2, 0, 1)[None]
    with torch.no_grad():
        return wrap(x).float().numpy()


def evaluate(mlpath, wrap, imgs, R):
    m = ct.models.MLModel(mlpath, compute_units=ct.ComputeUnit.CPU_AND_NE)
    agree, maxdiff = [], []
    for img in imgs:
        ref = torch_logits(wrap, img)                      # (1,19,g,g)
        out = m.predict({"image": img})
        got = np.asarray(out["logits"], dtype=np.float32)
        maxdiff.append(float(np.abs(ref - got).max()))
        agree.append(float((ref.argmax(1) == got.argmax(1)).mean()))
    # latency
    m.predict({"image": imgs[0]})
    t0 = time.time()
    N = 30
    for _ in range(N):
        m.predict({"image": imgs[0]})
    ms = (time.time() - t0) / N * 1000
    size_mb = sum(os.path.getsize(p) for p in glob.glob(mlpath + "/**/*", recursive=True) if os.path.isfile(p)) / 1e6
    return {"argmax_agreement": float(np.mean(agree)), "max_abs_diff": float(np.mean(maxdiff)),
            "latency_ms": round(ms, 2), "size_mb": round(size_mb, 2)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", required=True)
    ap.add_argument("--tag", required=True, help="output name suffix, e.g. large / nano")
    ap.add_argument("--imgsz", type=int, default=512)
    ap.add_argument("--out", default="/Users/ari/FaceSegmentation/exports_semantic")
    ap.add_argument("--n-val", type=int, default=8)
    ap.add_argument("--app-dest", default="", help="if set, copy recommended variant here")
    ap.add_argument("--min-agree", type=float, default=0.995)
    ap.add_argument("--quant-min-mb", type=float, default=10.0,
                    help="only consider quantized variants if fp16 is at least this large")
    args = ap.parse_args()

    os.makedirs(args.out, exist_ok=True)
    R = args.imgsz
    y = YOLO(args.weights)
    wrap = LogitsOnly(y.model).eval().float().cpu()

    val = sorted(glob.glob(os.path.join(VAL_IMAGES, "*.jpg")))[: args.n_val]
    assert val, f"no val images at {VAL_IMAGES}"
    imgs = [Image.open(p).convert("RGB").resize((R, R), Image.BILINEAR) for p in val]

    ts = torch.jit.trace(wrap, torch.rand(1, 3, R, R), strict=False)
    base = 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,
    )
    paths, results = {}, {}
    p_fp16 = f"{args.out}/FaceSeg_{args.tag}_fp16.mlpackage"
    base.save(p_fp16)
    paths["fp16"] = p_fp16

    from coremltools.optimize.coreml import (
        OpLinearQuantizerConfig, OpPalettizerConfig, OptimizationConfig,
        linear_quantize_weights, palettize_weights,
    )
    try:
        q = linear_quantize_weights(
            base, OptimizationConfig(global_config=OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int8")))
        p = f"{args.out}/FaceSeg_{args.tag}_int8.mlpackage"
        q.save(p)
        paths["int8"] = p
    except Exception as e:
        print("int8 quant failed:", repr(e)[:200])
    try:
        q = palettize_weights(
            base, OptimizationConfig(global_config=OpPalettizerConfig(mode="kmeans", nbits=6)))
        p = f"{args.out}/FaceSeg_{args.tag}_pal6.mlpackage"
        q.save(p)
        paths["pal6"] = p
    except Exception as e:
        print("palettization failed:", repr(e)[:200])

    for name, p in paths.items():
        results[name] = evaluate(p, wrap, imgs, R)
        print(f"[{name:5s}] {results[name]}")

    # Prefer fidelity. Quantizing only pays if it saves real space: for a ~3 MB
    # nano, trading measurable accuracy for 1.5 MB is a bad deal, while for a
    # 33 MB model halving the size is worth ~0.1% argmax disagreement.
    # (Measured: convert-then-quantize beats quantize-then-convert; see
    # ml/compare_quant_order.py.)
    rec = "fp16"
    base_mb = results["fp16"]["size_mb"]
    if base_mb >= args.quant_min_mb:
        ok = [n for n in results
              if n != "fp16" and results[n]["argmax_agreement"] >= args.min_agree]
        if ok:
            rec = min(ok, key=lambda n: results[n]["size_mb"])
    summary = {"weights": args.weights, "imgsz": R, "results": results, "recommended": rec,
               "recommended_path": paths[rec]}
    with open(f"{args.out}/summary_{args.tag}.json", "w") as f:
        json.dump(summary, f, indent=2)
    print("RECOMMENDED:", rec, "->", paths[rec])

    if args.app_dest:
        if os.path.exists(args.app_dest):
            shutil.rmtree(args.app_dest)
        shutil.copytree(paths[rec], args.app_dest)
        print("copied to app:", args.app_dest)
        # Keep SegModelContract in sync. A stale inputSize/gridSize makes
        # InferenceEngine.upload() reject every frame (shape guard) and the
        # overlay silently never appears.
        # NB: compute_units is REQUIRED here -- the default (.all) aborts in
        # MPSGraph ("MLIR pass manager failed") on this machine at >=384px,
        # which killed the export stage with SIGABRT after the model was saved.
        grid = int(np.asarray(
            ct.models.MLModel(paths[rec], compute_units=ct.ComputeUnit.CPU_AND_NE).predict(
                {"image": Image.new("RGB", (R, R))})["logits"]).shape[-1])
        sync_swift_contract(args.app_dest, R, grid)


def sync_swift_contract(app_dest, input_size, grid_size):
    """Rewrite inputSize/gridSize in SegmentationShared.swift to match the model."""
    shared = os.path.join(os.path.dirname(app_dest), "SegmentationShared.swift")
    if not os.path.exists(shared):
        print(f"WARNING: {shared} not found; update SegModelContract manually "
              f"(inputSize={input_size}, gridSize={grid_size})")
        return
    import re
    src = open(shared).read()
    new = re.sub(r"(inputSize:\s*Int\s*=\s*)\d+", rf"\g<1>{input_size}", src)
    new = re.sub(r"(gridSize:\s*Int\s*=\s*)\d+", rf"\g<1>{grid_size}", new)
    if new != src:
        open(shared, "w").write(new)
        print(f"synced SegModelContract -> inputSize={input_size}, gridSize={grid_size}")
    else:
        print(f"SegModelContract already correct (inputSize={input_size}, gridSize={grid_size})")


if __name__ == "__main__":
    main()