#!/usr/bin/env python3
"""Side-by-side FP32-vs-candidate mask visualization for the report.

For one image, renders three panels: FP32 baseline instances, candidate instances, and a per-pixel
mask disagreement map (baseline-only / candidate-only / agreement), plus the mean mask IoU. Uses the
same app-faithful preprocessing + edgecrafter-seg decode as the rest of the pipeline.

Usage:
  python visualize.py --baseline FP32.onnx --candidate CAND.onnx --image IMG.jpg --out OUT.png
"""

from __future__ import annotations

import argparse
import os
import sys

import cv2
import numpy as np
import onnxruntime as ort

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import ecseg_common as ec  # noqa: E402

# Distinct BGR colors for instance overlays.
PALETTE = [
    (0, 0, 255), (0, 255, 0), (255, 0, 0), (0, 255, 255), (255, 0, 255), (255, 255, 0),
    (0, 128, 255), (128, 0, 255), (0, 255, 128), (255, 128, 0), (128, 255, 0), (255, 0, 128),
]


def run(sess, x):
    names = [o.name for o in sess.get_outputs()]
    return dict(zip(names, sess.run(names, {"images": x})))


def instances_with_masks(out, w, h):
    inst = ec.parse_instances(out["labels"], out["boxes"], out["scores"], num_classes=80)
    masks = out["masks"].astype(np.float32)
    for it in inst:
        it["mask"] = ec.decode_mask(masks[0, it["q"]], w, h)
    return inst


def overlay(base_img, instances):
    canvas = base_img.copy()
    for i, it in enumerate(instances):
        color = PALETTE[i % len(PALETTE)]
        m = it["mask"]
        canvas[m] = (0.5 * canvas[m] + 0.5 * np.array(color)).astype(np.uint8)
        ys, xs = np.where(m)
        if len(xs):
            cv2.rectangle(canvas, (xs.min(), ys.min()), (xs.max(), ys.max()), color, 2)
    return canvas


def disagreement(base_inst, cand_inst, w, h):
    """Union-of-masks disagreement: red=baseline-only, blue=candidate-only, gray=agree."""
    b = np.zeros((h, w), bool)
    c = np.zeros((h, w), bool)
    for it in base_inst:
        b |= it["mask"]
    for it in cand_inst:
        c |= it["mask"]
    img = np.zeros((h, w, 3), np.uint8)
    img[np.logical_and(b, c)] = (90, 90, 90)
    img[np.logical_and(b, ~c)] = (0, 0, 255)   # baseline-only (lost)
    img[np.logical_and(~b, c)] = (255, 0, 0)   # candidate-only (spurious)
    iou = ec.mask_iou(b, c)
    return img, iou


def label(img, text):
    out = img.copy()
    cv2.rectangle(out, (0, 0), (img.shape[1], 26), (0, 0, 0), -1)
    cv2.putText(out, text, (6, 18), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--baseline", required=True)
    ap.add_argument("--candidate", required=True)
    ap.add_argument("--image", required=True)
    ap.add_argument("--out", required=True)
    args = ap.parse_args()

    bgr = cv2.imread(args.image, cv2.IMREAD_COLOR)
    h, w = bgr.shape[:2]
    x = ec.preprocess_bgr(bgr)

    base = ort.InferenceSession(args.baseline, providers=["CPUExecutionProvider"])
    cand = ort.InferenceSession(args.candidate, providers=["CPUExecutionProvider"])
    base_inst = instances_with_masks(run(base, x), w, h)
    cand_inst = instances_with_masks(run(cand, x), w, h)

    p1 = label(overlay(bgr, base_inst), f"FP32 baseline  ({len(base_inst)} inst)")
    p2 = label(overlay(bgr, cand_inst), f"{os.path.basename(args.candidate)}  ({len(cand_inst)} inst)")
    dis, iou = disagreement(base_inst, cand_inst, w, h)
    p3 = label(dis, f"disagreement  union maskIoU={iou:.3f}  (red=lost blue=spurious)")

    strip = np.concatenate([p1, p2, p3], axis=1)
    cv2.imwrite(args.out, strip)
    print(f"wrote {args.out}  base={len(base_inst)} cand={len(cand_inst)} unionMaskIoU={iou:.3f}")


if __name__ == "__main__":
    main()