"""
Visual end-to-end check of the SHIPPING path: run the .mlpackage exactly as the
iOS app does (center-square crop -> 512 input -> fp16 logits -> bilinear upsample
-> 19-way argmax -> palette LUT -> alpha mix) and write a contact sheet.

This mirrors SegmentationShaders.metal's math on the CPU, so a good-looking
sheet means the model AND the app's decode contract AND the palette are right.

Usage: verify_coreml_visual.py --model <.mlpackage> --out <dir> [--source <dir>]
"""
import argparse, glob, os

import numpy as np
from PIL import Image, ImageDraw
import coremltools as ct

from palette import CLASS_NAMES, PALETTE

PAL = np.array(PALETTE, dtype=np.float32)


def center_square(im):
    w, h = im.size
    s = min(w, h)
    return im.crop(((w - s) // 2, (h - s) // 2, (w + s) // 2, (h + s) // 2))


def bilinear_upsample(logits, out):
    """(C,g,g) -> (C,out,out), matching the shader's bilinear sampler."""
    C, g, _ = logits.shape
    # sample centers, same convention as a normalized-UV bilinear sampler
    coord = (np.arange(out) + 0.5) / out * g - 0.5
    coord = np.clip(coord, 0, g - 1)
    i0 = np.floor(coord).astype(int)
    i1 = np.minimum(i0 + 1, g - 1)
    w1 = (coord - i0).astype(np.float32)
    w0 = 1.0 - w1
    tmp = logits[:, :, i0] * w0 + logits[:, :, i1] * w1      # x interp
    return tmp[:, i0, :] * w0[None, :, None] + tmp[:, i1, :] * w1[None, :, None]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--out", default="/private/tmp/claude-501/-Users-ari-Documents-XcodeProjects-facesegmentation/149f7917-f682-4847-ba5b-1e726cecc20c/scratchpad/visual")
    ap.add_argument("--source", default="/Users/ari/FaceSegmentation/dataset_celebamaskhq_semantic/images/val")
    ap.add_argument("--n", type=int, default=8)
    ap.add_argument("--alpha", type=float, default=0.55)
    ap.add_argument("--tile", type=int, default=320)
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)

    m = ct.models.MLModel(args.model, compute_units=ct.ComputeUnit.CPU_AND_NE)
    spec_in = list(m.input_description)[0]
    R = 512
    files = sorted(glob.glob(os.path.join(args.source, "**", "*.jpg"), recursive=True))[: args.n]
    if not files:
        files = sorted(glob.glob(os.path.join(args.source, "**", "*.png"), recursive=True))[: args.n]

    tiles, class_hits = [], np.zeros(len(CLASS_NAMES))
    for p in files:
        im = center_square(Image.open(p).convert("RGB")).resize((R, R), Image.BILINEAR)
        out = m.predict({spec_in: im})
        logits = np.asarray(out["logits"], dtype=np.float32)[0]      # (19,g,g)
        up = bilinear_upsample(logits, R)                            # (19,R,R)
        cls = up.argmax(0).astype(np.int32)
        class_hits += np.bincount(cls.ravel(), minlength=len(CLASS_NAMES)) / cls.size

        rgb = np.asarray(im, dtype=np.float32)
        col = PAL[cls]
        a = (cls != 0).astype(np.float32)[..., None] * args.alpha
        blend = (rgb * (1 - a) + col * a).astype(np.uint8)
        pair = Image.new("RGB", (args.tile * 2, args.tile))
        pair.paste(im.resize((args.tile, args.tile)), (0, 0))
        pair.paste(Image.fromarray(blend).resize((args.tile, args.tile)), (args.tile, 0))
        tiles.append(pair)

    cols = 2
    rows = (len(tiles) + cols - 1) // cols
    W, H = tiles[0].size
    sheet = Image.new("RGB", (cols * W, rows * H), (16, 16, 16))
    for i, t in enumerate(tiles):
        sheet.paste(t, ((i % cols) * W, (i // cols) * H))
    sheet_path = os.path.join(args.out, "verify_sheet.jpg")
    sheet.save(sheet_path, quality=92)

    # legend of what the model actually predicted
    order = np.argsort(-class_hits)
    print("mean class coverage (top 10):")
    for i in order[:10]:
        if class_hits[i] > 0:
            print(f"  {CLASS_NAMES[i]:14s} {class_hits[i]/len(tiles):6.2%}")
    print("wrote", sheet_path)


if __name__ == "__main__":
    main()