#!/usr/bin/env python3
"""Shared ECSeg preprocessing + output post-processing, byte-faithful to the app.

The preprocessing here reproduces exactly what the browser app feeds the ONNX graph, so that
a Python correctness comparison feeds baseline and candidate the SAME `images` tensor the
product would:

  * OpenCV path mirrors `packages/smart-tools/src/segment-anything/pre-processing.ts`:
      - decode → RGB (drop alpha)
      - stretch resize to 640×640 with cv2.INTER_LINEAR (bilinear), aspect NOT preserved
      - scale by 1/255 (convertTo CV_32F)
      - ImageNet normalize: (x - mean) / std, mean=[0.485,0.456,0.406] std=[0.229,0.224,0.225]
      - NCHW float32, shape [1,3,640,640]

The post-processing mirrors the `edgecrafter-seg` parser + mask decode:
  * score filter at confidenceThreshold (0.4)
  * per-instance 160×160 logit map → bilinear upscale to image size → threshold at maskThreshold
    (raw LOGIT cut, default 0.0) → binary mask
"""

from __future__ import annotations

import glob
import os
from typing import List, Tuple

import cv2
import numpy as np

SIZE = 640
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

CONF_THRESHOLD = 0.4
MASK_THRESHOLD = 0.0  # raw logit cut (logit > 0 ⇔ sigmoid > 0.5)


def preprocess_bgr(bgr: np.ndarray) -> np.ndarray:
    """cv2-decoded BGR image (H,W,3 uint8) -> float32 NCHW [1,3,640,640] tensor."""
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    resized = cv2.resize(rgb, (SIZE, SIZE), interpolation=cv2.INTER_LINEAR)
    x = resized.astype(np.float32) / 255.0
    x = (x - MEAN) / STD  # broadcast over channels
    # HWC -> CHW -> NCHW
    x = np.transpose(x, (2, 0, 1))[None, ...]
    return np.ascontiguousarray(x, dtype=np.float32)


def preprocess_file(path: str) -> np.ndarray:
    bgr = cv2.imread(path, cv2.IMREAD_COLOR)
    if bgr is None:
        raise ValueError(f"cv2 failed to read {path}")
    return preprocess_bgr(bgr)


def list_images(dirs: List[str], limit: int | None = None) -> List[str]:
    """Deterministically enumerate images across directories (sorted, dedup by basename order)."""
    exts = ("*.jpg", "*.jpeg", "*.png", "*.JPG", "*.JPEG", "*.PNG")
    files: List[str] = []
    for d in dirs:
        for ext in exts:
            files.extend(glob.glob(os.path.join(d, "**", ext), recursive=True))
    files = sorted(set(files))
    if limit is not None:
        files = files[:limit]
    return files


# ------------------------------------------------------------------ output post-processing ----

def parse_instances(
    labels: np.ndarray,
    boxes: np.ndarray,
    scores: np.ndarray,
    conf_threshold: float = CONF_THRESHOLD,
    num_classes: int = 80,
) -> List[dict]:
    """Mirror parseEdgecrafterSeg: score-filter, validate class id, clamp+order box corners.

    Returns list of {q, classId, score, box(xyxy in [0,1])}, in raw query order (no NMS).
    """
    labels = np.asarray(labels).reshape(-1)
    scores = np.asarray(scores).reshape(-1)
    boxes = np.asarray(boxes).reshape(-1, 4)
    out = []
    for q in range(scores.shape[0]):
        s = float(scores[q])
        if not (s >= conf_threshold):  # NaN-safe
            continue
        cls = int(labels[q])
        if cls < 0 or cls >= num_classes:
            continue
        x0, y0, x1, y1 = [float(v) for v in boxes[q]]
        x0, x1 = sorted((min(max(x0, 0.0), 1.0), min(max(x1, 0.0), 1.0)))
        y0, y1 = sorted((min(max(y0, 0.0), 1.0), min(max(y1, 0.0), 1.0)))
        out.append({"q": q, "classId": cls, "score": s, "box": (x0, y0, x1, y1)})
    return out


def decode_mask(
    mask_logits_160: np.ndarray,
    out_w: int,
    out_h: int,
    mask_threshold: float = MASK_THRESHOLD,
) -> np.ndarray:
    """Mirror decodeEdgeSegmentation: bilinear upscale 160×160 logits to (out_h,out_w), threshold.

    Returns a boolean mask (out_h, out_w).
    """
    m = np.asarray(mask_logits_160, dtype=np.float32)
    up = cv2.resize(m, (out_w, out_h), interpolation=cv2.INTER_LINEAR)
    return up > mask_threshold


def mask_iou(a: np.ndarray, b: np.ndarray) -> float:
    """IoU of two boolean masks of identical shape."""
    a = a.astype(bool)
    b = b.astype(bool)
    inter = np.logical_and(a, b).sum(dtype=np.int64)
    union = np.logical_or(a, b).sum(dtype=np.int64)
    if union == 0:
        return 1.0  # both empty -> identical
    return float(inter) / float(union)


def box_iou(a: Tuple[float, float, float, float], b: Tuple[float, float, float, float]) -> float:
    ax0, ay0, ax1, ay1 = a
    bx0, by0, bx1, by1 = b
    ix0, iy0 = max(ax0, bx0), max(ay0, by0)
    ix1, iy1 = min(ax1, bx1), min(ay1, by1)
    iw, ih = max(0.0, ix1 - ix0), max(0.0, iy1 - iy0)
    inter = iw * ih
    area_a = max(0.0, ax1 - ax0) * max(0.0, ay1 - ay0)
    area_b = max(0.0, bx1 - bx0) * max(0.0, by1 - by0)
    union = area_a + area_b - inter
    if union <= 0:
        return 1.0 if inter == 0 else 0.0
    return inter / union