"""슬롯 판정 — 등록 레이아웃 + 검사 씬 → 슬롯별 present / missing / wrong_tool + 슬롯밖 foreign(이물).

3단계 (클래스 무관, Track A 탐지기 클래스에 의존하지 않음):
  1) 점유(occupancy): 슬롯 모서리 배경색 대비 전경 픽셀 비율 → 낮으면 missing (고전 CV)
  2) 정체(identity): 전경 FFA 임베딩과 등록 지문의 코사인 유사도 → 낮으면 wrong_tool (DINOv2).
     tier3(기본): 등록 크롭을 여러 각도로 회전한 임베딩 뱅크와 max-sim → 회전으로 흔들린 정상 공구
     회수(present→wrong_tool 오판 0.07→0.023, 씬 오경보 0.806→0.514, wrong_tool recall 유지)
  3) 이물(foreign): 등록 슬롯 박스 '바깥' 자유공간에서 보드 배경 대비 전경 덩어리 → 이물 후보
     (슬롯 점유와 동일한 상대임계 median+MAD_K·MAD·denoise 엔진 재사용)

임계값(occ_thr, id_thr)은 dev에서 슬롯 상태 balanced accuracy를, foreign_min_area는 dev에서
이물 F1을 (씬 오경보 게이트 하에) 최대화하도록 탐색. 씬 pass/fail은 슬롯만/슬롯+이물 2종을
같이 보고해 이물탐지 on/off 트레이드오프(침입탐지↑ ↔ 정상 씬 오경보↑)를 정량화한다.

실행: python -m src.layout --tune-eval
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
from PIL import Image
from scipy import ndimage

from src.enroll import DinoEmbedder, cos

ROOT = Path(__file__).resolve().parents[1]
BENCH = ROOT / "data/synthbench"
POOL = BENCH / "pool"
INSET = 16          # 슬롯 테두리(그려진 윤곽선) 회피 여백
COLOR_THR = 38.0    # 배경색 대비 전경 판정 색거리


def _slot_region(scene: np.ndarray, bbox):
    x, y, w, h = bbox
    H, W = scene.shape[:2]
    x1, y1 = max(0, x + INSET), max(0, y + INSET)
    x2, y2 = min(W, x + w - INSET), min(H, y + h - INSET)
    return scene[y1:y2, x1:x2], (x1, y1, x2, y2)


MAD_K = 3.0     # tier1a 상대임계 민감도 (median + k·MAD)
TAU_CD = 3.0    # tier1b Horprasert 색성분(CD) 임계 (bg MAD 배수)
TAU_A = 3.0     # tier1b Horprasert 밝기성분(BD) 임계 (어두워지는 방향)

EDGE_K = 5.0               # 점유 엣지 보강: 얇은/저조도 금속(색은 비슷해도 윤곽 강함)을 색거리 OR 엣지에너지로 회수
FOREIGN_MIN_AREA = 0.0005  # 이물 최소 전경면적(씬 대비) — 0.0005가 마지노선(더 낮추면 잡음덩어리 통과 P붕괴, dev 튜닝)
FOREIGN_DILATE = 10        # 슬롯 박스 제외 시 팽창(테두리 그림자·정렬오차 여유)
FOREIGN_MED = 5            # 이물 스캔 median 커널: 배경 speckle 억제(이물 ≥40px는 생존)
FOREIGN_GAUSS = 3.5        # 이물 스캔 gaussian σ: 가산잡음까지 평균 → 배경 MAD 급감, 오탐 P≈1.0 달성
FOREIGN_K = 3.5            # 이물 상대임계 배수(정밀 우선; 낮추면 recall↑ precision↓ — 나쁜 거래)


def _mad(x):
    med = float(np.median(x))
    return med, float(np.median(np.abs(x - med))) * 1.4826 + 1e-6


def _fg_mask(reg: np.ndarray, bg: np.ndarray, mode: str) -> np.ndarray:
    """모드별 전경 픽셀 마스크(연결성분 필터 전).
    raw=고정 색거리 임계(밝기보정 없음) / tier0=밝기비례 색거리+고정임계(빈칸오점유↑, 폐기) /
    mad=색거리+median+k·MAD 상대임계 /
    horprasert=배경색축 분해(BD 밝기·CD 색)로 '색 다름 OR 더 어두움' 판정(저조도 회수)."""
    regf = reg.astype(np.float32)
    if mode == "raw":
        return np.linalg.norm(regf - bg, axis=2) > COLOR_THR
    if mode == "mad":
        dist = np.linalg.norm(regf - bg, axis=2)
        med, mad = _mad(dist.reshape(-1))
        return dist > med + MAD_K * mad
    if mode == "horprasert":
        E = bg.astype(np.float32)
        ee = float(np.dot(E, E)) + 1e-6
        alpha = (regf @ E) / ee                              # BD: 배경색 축 밝기 성분 (1=배경과 동일 밝기)
        cd = np.linalg.norm(regf - alpha[..., None] * E, axis=2)  # CD: 축에 수직인 색 성분
        a_med, a_mad = _mad(alpha.reshape(-1))
        c_med, c_mad = _mad(cd.reshape(-1))
        fg_color = cd > c_med + TAU_CD * c_mad               # 배경과 색이 다름
        fg_dark = alpha < a_med - TAU_A * a_mad               # 배경과 색은 비슷하나 더 어두움(저조도 공구)
        return fg_color | fg_dark
    # tier0 (기본): 밝기 비례 스케일링 + 고정 색거리 임계
    scale = min(4.0, 128.0 / max(float(bg.mean()), 8.0))
    dist = np.linalg.norm(regf - bg, axis=2) * scale
    return dist > COLOR_THR


def occupancy(scene: np.ndarray, bbox, mode: str = "raw", denoise: bool = False, edge: bool = False):
    """슬롯 4모서리에서 배경색을 추정하고, 내부 전경 픽셀 비율과 마스크를 반환.
    denoise=True면 median filter로 노이즈 억제 후 판정(노이즈 조건 회복).
    edge=True면 색거리 마스크에 엣지에너지(Sobel) 마스크를 OR — 얇은/저조도 금속 회수."""
    reg, _ = _slot_region(scene, bbox)
    if reg.size == 0:
        return 0.0, np.zeros((1, 1), bool)
    if denoise:
        reg = ndimage.median_filter(reg, size=(3, 3, 1))  # edge 보존 경량 denoise
    h, w = reg.shape[:2]
    cs = max(4, min(h, w) // 5)
    corners = np.concatenate([
        reg[:cs, :cs].reshape(-1, 3), reg[:cs, -cs:].reshape(-1, 3),
        reg[-cs:, :cs].reshape(-1, 3), reg[-cs:, -cs:].reshape(-1, 3)], 0)
    bg = np.median(corners, 0)
    mask = _fg_mask(reg, bg, mode)
    if edge:  # 색이 배경과 비슷해도 윤곽이 강한 얇은 금속(가위·나사)을 엣지로 회수
        gray = reg.astype(np.float32).mean(2)
        gm = np.hypot(ndimage.sobel(gray, axis=1), ndimage.sobel(gray, axis=0))
        gmed, gmad = _mad(gm.reshape(-1))
        mask = mask | (gm > gmed + EDGE_K * gmad)
    # 연결 성분 필터: 배경 무늬(pegboard 구멍·나뭇결)는 작은 조각 → 제거, 공구=큰 덩어리만 유지
    lbl, n = ndimage.label(mask)
    if n:
        areas = ndimage.sum(mask, lbl, range(1, n + 1))
        min_area = 0.004 * mask.size
        keep_ids = {i for i, a in enumerate(areas, 1) if a >= min_area}
        mask = np.isin(lbl, list(keep_ids)) if keep_ids else np.zeros_like(mask)
    return float(mask.mean()), mask


def _iou(a, b) -> float:
    ax, ay, aw, ah = a
    bx, by, bw, bh = b
    ix = max(0, min(ax + aw, bx + bw) - max(ax, bx))
    iy = max(0, min(ay + ah, by + bh) - max(ay, by))
    inter = ix * iy
    return inter / max(1, aw * ah + bw * bh - inter)


def _foreign_scan(scene: np.ndarray, slot_boxes, denoise: bool,
                  min_area_frac: float = FOREIGN_MIN_AREA) -> list:
    """등록 슬롯 박스 '바깥'(자유공간)에서 보드 배경 대비 전경 덩어리를 찾아 이물 후보 bbox 목록 반환.
    슬롯 점유와 동일한 상대임계(median+MAD_K·MAD)·denoise 엔진 재사용, 클래스 무관."""
    H, W = scene.shape[:2]
    img = scene.astype(np.float32)
    if denoise:  # median(speckle 억제)+gaussian(가산잡음 평균) → 거친/저조도 배경에서도 이물이 드러남
        img = ndimage.median_filter(img, size=(FOREIGN_MED, FOREIGN_MED, 1))
        img = ndimage.gaussian_filter(img, sigma=(FOREIGN_GAUSS, FOREIGN_GAUSS, 0))
    free = np.ones((H, W), bool)                              # 슬롯(팽창) 밖 = 자유공간
    for x, y, w, h in slot_boxes:
        free[max(0, y - FOREIGN_DILATE):y + h + FOREIGN_DILATE,
             max(0, x - FOREIGN_DILATE):x + w + FOREIGN_DILATE] = False
    if free.sum() < 100:
        return []
    bg = np.median(img[free], axis=0)                        # 보드 배경색(이물은 소수 → median≈보드)
    dist = np.linalg.norm(img - bg, axis=2)
    med, mad = _mad(dist[free])
    fg = ndimage.binary_opening((dist > med + FOREIGN_K * mad) & free, iterations=1)  # 흩뿌린 잡음 제거
    lbl, n = ndimage.label(fg)
    if not n:
        return []
    areas = ndimage.sum(fg, lbl, range(1, n + 1))
    slices = ndimage.find_objects(lbl)
    min_area = min_area_frac * H * W
    out = []
    for i, sl in enumerate(slices, 1):
        if sl is None or areas[i - 1] < min_area:
            continue
        ys, xs = sl
        out.append([xs.start, ys.start, xs.stop - xs.start, ys.stop - ys.start])
    return out


TTA_ANGLES = (-20.0, -10.0, 0.0, 10.0, 20.0)  # tier3 회전 증강 각도(합성 paste의 ±20° 회전 범위를 덮음)


class Judge:
    def __init__(self, embedder: DinoEmbedder, mode: str = "mad", denoise: bool = True,
                 tier3: bool = True, edge: bool = False):
        self.emb = embedder
        self.mode = mode         # 프로덕션 확정: mad + denoise (전 조건 실패모드 없음)
        self.denoise = denoise
        # tier3 확정: 회전증강 뱅크 max-sim → present→wrong_tool 0.07→0.023, 씬오경보 0.806→0.514 (recall 유지)
        self.tier3 = tier3       # True=등록 크롭 회전증강 뱅크로 max-sim(identity 재검증)
        self.edge = edge         # True=점유에 엣지에너지 보강(얇은/저조도 금속 회수, A/B 중)
        self._reg_cache: dict[str, list] = {}

    def _registered_bank(self, rel_path: str) -> list:
        """등록 크롭 임베딩 뱅크. tier3면 여러 각도 회전본, 아니면 원본 1개.
        회전은 RGBA째 돌려 알파(전경 마스크)도 함께 회전 → FFA 전경 정합 유지."""
        if rel_path not in self._reg_cache:
            rgba = Image.open(POOL / rel_path).convert("RGBA")
            angles = TTA_ANGLES if self.tier3 else (0.0,)
            bank = []
            for ang in angles:
                r = rgba.rotate(ang, expand=True, resample=Image.BICUBIC) if ang else rgba
                bank.append(self.emb.embed_ffa(r.convert("RGB"), np.asarray(r.split()[-1])))
            self._reg_cache[rel_path] = bank
        return self._reg_cache[rel_path]

    def judge_scene(self, scene_path: Path, layout: dict, occ_thr: float, id_thr: float,
                    foreign_min_area: float = FOREIGN_MIN_AREA, do_foreign: bool = True):
        scene = np.asarray(Image.open(scene_path).convert("RGB"))
        out = []
        for s in layout["slots"]:
            occ, mask = occupancy(scene, s["bbox"], self.mode, self.denoise, self.edge)
            if occ < occ_thr:
                out.append({"slot_id": s["slot_id"], "state": "missing", "occ": round(occ, 3), "sim": None})
                continue
            reg_region, _ = _slot_region(scene, s["bbox"])
            e = self.emb.embed_ffa(Image.fromarray(reg_region), mask)
            sim = max(cos(e, r) for r in self._registered_bank(s["registered"]))  # tier3=회전뱅크 max-sim
            state = "present" if sim >= id_thr else "wrong_tool"
            out.append({"slot_id": s["slot_id"], "state": state, "occ": round(occ, 3), "sim": round(sim, 3)})
        # do_foreign=False → 슬롯 임계 튜닝 중엔 무거운 이물 스캔 생략(속도)
        foreign = _foreign_scan(scene, [s["bbox"] for s in layout["slots"]], self.denoise, foreign_min_area) \
            if do_foreign else []
        return {"slots": out, "foreign": foreign}


def _load(split: str):
    man = json.loads((BENCH / "manifest.json").read_text())["scenes"]
    layouts = {p.stem.replace("layout_", ""): json.loads(p.read_text())
               for p in (BENCH / "layouts").glob("layout_*.json")}
    rows = []
    for m in man:
        if m["split"] != split:
            continue
        gt = json.loads((BENCH / "gt" / f"{m['scene_id']}.json").read_text())
        rows.append((m["scene_id"], layouts[m["layout_id"]],
                     {s["slot_id"]: s["state"] for s in gt["slots"]},
                     [f["obj_bbox"] for f in gt.get("foreign_objects", [])]))
    return rows


def _scene_metrics(sc: dict) -> dict:
    """씬 단위 {tp,fp,fn,tn}(fail=positive) → pass/fail 정확도·불량탐지율·오경보율."""
    n = sum(sc.values())
    return {
        "pass_fail_acc": round((sc["tp"] + sc["tn"]) / max(1, n), 3),
        "fail_recall": round(sc["tp"] / max(1, sc["tp"] + sc["fn"]), 3),   # 불량 공구함을 불량으로 잡는 비율
        "false_alarm": round(sc["fp"] / max(1, sc["fp"] + sc["tn"]), 3),   # 정상을 불량이라 오경보하는 비율
        "counts": sc,
    }


IOU_MATCH = 0.15   # 이물 예측 bbox ↔ GT bbox 매칭 임계 (작은 이물이라 관대하게)


def _score(judge: Judge, rows, occ_thr, id_thr, foreign_min_area: float = FOREIGN_MIN_AREA,
           do_foreign: bool = True) -> dict:
    """슬롯 상태 + 이물 예측을 GT와 대조해 지표 산출.
    씬 pass/fail은 슬롯만(scene_slotonly) / 슬롯+이물(scene) 2종을 함께 내 트레이드오프를 드러낸다."""
    states = ["present", "missing", "wrong_tool"]
    conf = {a: {b: 0 for b in states} for a in states}  # conf[gt][pred]
    sc_slot = {"tp": 0, "fp": 0, "fn": 0, "tn": 0}      # 씬 fail 판정 — 슬롯만(이물 제외)
    sc_full = {"tp": 0, "fp": 0, "fn": 0, "tn": 0}      # 씬 fail 판정 — 슬롯+이물
    cond_pm = {"normal": [0, 0], "dark": [0, 0], "noisy": [0, 0]}  # 조건별 [present총, present→missing 오판]
    f_tp = f_fp = f_fn = 0                               # 이물 객체 단위(bbox 매칭)
    fs = {"tp": 0, "fp": 0, "fn": 0, "tn": 0}           # 이물 씬 단위(이물 존재 유무)
    for sid, layout, gt, gt_foreign in rows:
        cond = sid.rsplit("_", 1)[-1]
        res = judge.judge_scene(BENCH / "scenes" / f"{sid}.jpg", layout, occ_thr, id_thr,
                                foreign_min_area, do_foreign)
        pred = {p["slot_id"]: p["state"] for p in res["slots"]}
        pred_foreign = res["foreign"]
        for slot_id, g in gt.items():
            conf[g][pred[slot_id]] += 1
            if g == "present" and cond in cond_pm:
                cond_pm[cond][0] += 1
                cond_pm[cond][1] += pred[slot_id] == "missing"
        # 이물 객체 매칭 (greedy IoU)
        matched = set()
        for pb in pred_foreign:
            hit = next((j for j, gb in enumerate(gt_foreign) if j not in matched and _iou(pb, gb) >= IOU_MATCH), None)
            if hit is None:
                f_fp += 1
            else:
                matched.add(hit)
                f_tp += 1
        f_fn += len(gt_foreign) - len(matched)
        g_has_f, p_has_f = len(gt_foreign) > 0, len(pred_foreign) > 0
        fs["tp" if g_has_f and p_has_f else "fn" if g_has_f else "fp" if p_has_f else "tn"] += 1
        # 씬 pass/fail — 슬롯만 / 슬롯+이물
        gt_slot_fail = any(v != "present" for v in gt.values())
        pred_slot_fail = any(v != "present" for v in pred.values())
        sc_slot["tp" if gt_slot_fail and pred_slot_fail else "fn" if gt_slot_fail
                else "fp" if pred_slot_fail else "tn"] += 1
        gt_full_fail, pred_full_fail = gt_slot_fail or g_has_f, pred_slot_fail or p_has_f
        sc_full["tp" if gt_full_fail and pred_full_fail else "fn" if gt_full_fail
                else "fp" if pred_full_fail else "tn"] += 1
    total = sum(conf[a][b] for a in states for b in states)
    acc = sum(conf[s][s] for s in states) / max(1, total)
    per = {}
    for s in states:
        tp = conf[s][s]
        fn = sum(conf[s][b] for b in states) - tp
        fp = sum(conf[a][s] for a in states) - tp
        p = tp / max(1, tp + fp)
        r = tp / max(1, tp + fn)
        per[s] = {"P": round(p, 3), "R": round(r, 3), "F1": round(2 * p * r / max(1e-9, p + r), 3), "n": tp + fn}
    bal = float(np.mean([per[s]["R"] for s in states]))
    fp_p = f_tp / max(1, f_tp + f_fp)
    fp_r = f_tp / max(1, f_tp + f_fn)
    foreign = {
        "obj_P": round(fp_p, 3), "obj_R": round(fp_r, 3),
        "obj_F1": round(2 * fp_p * fp_r / max(1e-9, fp_p + fp_r), 3), "n_obj": f_tp + f_fn,
        "scene_R": round(fs["tp"] / max(1, fs["tp"] + fs["fn"]), 3),     # 이물 있는 씬을 잡는 비율
        "scene_FA": round(fs["fp"] / max(1, fs["fp"] + fs["tn"]), 3),    # 이물 없는데 있다 오경보
        "counts": {"obj_tp": f_tp, "obj_fp": f_fp, "obj_fn": f_fn, "scene": fs},
    }
    cond_err = {c: round(v[1] / max(1, v[0]), 3) for c, v in cond_pm.items()}  # 조건별 present→missing 오판율
    return {"slot_acc": round(acc, 3), "balanced_acc": round(bal, 3),
            "per_state": per, "confusion": conf,
            "scene": _scene_metrics(sc_full), "scene_slotonly": _scene_metrics(sc_slot),
            "foreign": foreign, "cond_pm_err": cond_err}


def _tune(judge, dev):
    best = None
    for occ_thr in [0.015, 0.03, 0.05, 0.08, 0.12]:
        for id_thr in [0.35, 0.42, 0.50, 0.58, 0.66]:
            r = _score(judge, dev, occ_thr, id_thr, do_foreign=False)  # 슬롯 임계 탐색엔 이물 불필요
            if best is None or r["balanced_acc"] > best[2]:
                best = (occ_thr, id_thr, r["balanced_acc"])
    return best


def _tune_foreign(judge, dev, occ_thr, id_thr, fa_gate=0.15):
    """이물 min_area를 dev에서 탐색: 이물 씬 오경보(scene_FA) ≤ fa_gate 게이트 하에 이물 F1 최대화.
    게이트 통과 후보가 없으면 오경보 최소를 택한다(트레이드오프 게이트 방식)."""
    cands = []
    for ma in [0.0005, 0.0008, 0.0012, 0.002]:  # 0.0003↓는 잡음덩어리 통과로 P붕괴 → 제외
        f = _score(judge, dev, occ_thr, id_thr, ma)["foreign"]
        cands.append((ma, f["obj_F1"], f["scene_FA"]))
    gated = [c for c in cands if c[2] <= fa_gate]
    return (max(gated, key=lambda c: c[1]) if gated else min(cands, key=lambda c: c[2]))[0]


def _report(label, occ_thr, id_thr, fma, res):
    ce, sc, ss, fo = res["cond_pm_err"], res["scene"], res["scene_slotonly"], res["foreign"]
    pw = res["confusion"]["present"]
    p2w = pw["wrong_tool"] / max(1, sum(pw.values()))   # present→wrong_tool 오판율(씬 오경보 주범)
    print(f"[{label:11}] thr=({occ_thr},{id_thr}) fmin_area={fma}  slot_acc={res['slot_acc']}  balanced={res['balanced_acc']}")
    print(f"              present→missing: normal={ce['normal']} dark={ce['dark']} noisy={ce['noisy']} | present→wrong_tool={round(p2w, 3)}")
    print(f"              missing F1={res['per_state']['missing']['F1']} wrong_tool P={res['per_state']['wrong_tool']['P']} R={res['per_state']['wrong_tool']['R']} F1={res['per_state']['wrong_tool']['F1']}")
    print(f"              이물 객체: P={fo['obj_P']} R={fo['obj_R']} F1={fo['obj_F1']} (n={fo['n_obj']}) | 씬탐지R={fo['scene_R']} 씬오경보={fo['scene_FA']}")
    print(f"              씬 불량탐지 슬롯만={ss['fail_recall']}→+이물={sc['fail_recall']} | 씬 오경보 슬롯만={ss['false_alarm']}→+이물={sc['false_alarm']}\n")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tune-eval", action="store_true")
    ap.add_argument("--ab", action="store_true", help="occupancy 모드 A/B 비교 (tier0/mad/horprasert)")
    ap.add_argument("--ab-tier3", action="store_true", help="identity 재검증 A/B (tier2 단일 vs tier3 회전뱅크)")
    ap.add_argument("--ab-edge", action="store_true", help="점유 엣지보강 A/B (색거리 vs 색거리+엣지)")
    ap.add_argument("--modes", default="tier0,mad,horprasert")
    ap.add_argument("--model", default="facebook/dinov2-small")
    a = ap.parse_args()

    emb = DinoEmbedder(a.model)
    dev, test = _load("dev"), _load("test")
    print(f"dev {len(dev)} scenes / test {len(test)} scenes | model {a.model}\n")

    if a.ab:                                                # (label, judge kwargs)
        configs = [(m, {"mode": m}) for m in a.modes.split(",")]
    elif a.ab_tier3:
        configs = [("tier2", {"tier3": False}), ("tier3", {"tier3": True})]
    elif a.ab_edge:
        configs = [("occ_color", {"edge": False}), ("occ_edge", {"edge": True})]  # 점유 엣지보강 A/B
    else:
        configs = [("prod", {})]                            # 프로덕션 = mad + tier3 (Judge 기본값)

    results = {}
    for label, kw in configs:
        judge = Judge(emb, **kw)
        occ_thr, id_thr, _ = _tune(judge, dev)
        fma = _tune_foreign(judge, dev, occ_thr, id_thr)
        res = _score(judge, test, occ_thr, id_thr, fma)
        results[label] = {"occ_thr": occ_thr, "id_thr": id_thr, "foreign_min_area": fma, "test": res}
        _report(label, occ_thr, id_thr, fma, res)

    (ROOT / "outputs").mkdir(exist_ok=True)
    out = "synthbench_ab_tier3.json" if a.ab_tier3 else "synthbench_ab.json"
    (ROOT / "outputs" / out).write_text(json.dumps(
        {"model": a.model, "results": results}, ensure_ascii=False, indent=1))
    print(f"저장: outputs/{out}")


if __name__ == "__main__":
    main()