"""Diagnostyka etapow segmentacji v7 — zrzuca maske po kazdym kroku.

Uzycie: python tmp/dbg_mask_stages.py <master.png> [out_dir]
"""
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "services", "recolor"))

import numpy as np
import cv2
from PIL import Image
from fabric_recolor.masks import (
    alpha_hysteresis, cavity_anchors, clean_binary, cut_anchored_pockets,
    feather, fill_small_holes, largest_component, recessed_pocket_labels,
    subtract_protected, _silhouette_bands,
)
from fabric_recolor.labcolor import rgb_to_lab
from fabric_recolor.segmentation.prompts import NEGATIVE_PROMPTS, KEEP_PROMPTS, CAVITY_PROMPTS
from fabric_recolor.segmentation.grounded_sam import GroundedSamSegmenter

src = sys.argv[1]
out_dir = sys.argv[2] if len(sys.argv) > 2 else "tmp/dbg_mask"
os.makedirs(out_dir, exist_ok=True)

from fabric_recolor.segmentation.birefnet import BirefnetMatteSegmenter
seg = BirefnetMatteSegmenter(device="cuda", max_side=1600)
seg._load()

rgb = np.asarray(Image.open(src).convert("RGB"), dtype=np.uint8)
h0, w0 = rgb.shape[:2]
scale = min(1.0, seg.max_side / max(h0, w0))
image = Image.fromarray(rgb).resize((round(w0 * scale), round(h0 * scale)), Image.BILINEAR) if scale < 1.0 else Image.fromarray(rgb)
arr = np.asarray(image)

def dump(name, m):
    a = m if m.dtype != np.uint8 else m * 255
    if a.max() <= 1.0:
        a = (a * 255).astype(np.uint8)
    Image.fromarray(np.clip(a, 0, 255).astype(np.uint8)).save(os.path.join(out_dir, name))
    print(f"{name}: cov={(a>25).mean():.4f}")

matte = seg._matte(image)
dump("01_matte.png", matte)
matte_h = alpha_hysteresis(matte, ramp_px=max(4, round(8 * scale)))
dump("02_hysteresis.png", matte_h)
mask = clean_binary((matte_h >= seg.matte_threshold).astype(np.uint8), min_area_frac=0.002)
mask = largest_component(mask)
mask = fill_small_holes(mask, max_hole_frac=0.01)
dump("03_mask_pre_neg.png", mask)

iso = Image.fromarray(arr * mask[..., None].astype(np.uint8))
neg = seg._neg_segmenter()
raw = neg.negatives_instances(iso, list(NEGATIVE_PROMPTS) + list(KEEP_PROMPTS))
cavity_ids = {i for i, p in enumerate(NEGATIVE_PROMPTS) if p in CAVITY_PROMPTS}
n_neg = len(NEGATIVE_PROMPTS)
instances, keep_mask = [], np.zeros(mask.shape, np.uint8)
phrases = list(NEGATIVE_PROMPTS) + list(KEEP_PROMPTS)
for m, idx in raw:
    area = int((m > 0).sum())
    lab = "NEG:" + phrases[idx] if idx < n_neg else "KEEP:" + phrases[idx]
    print(f"  inst: {lab} area={area} ({100*area/mask.size:.2f}% kadru)")
    if idx < n_neg:
        instances.append((m, idx in cavity_ids))
        dump(f"neg_{idx}_{phrases[idx].replace(' ', '_')}.png", m.astype(np.uint8))
    else:
        keep_mask |= ((m > 0) & (mask > 0)).astype(np.uint8)

depth = seg._depth_map(image)
n_vol = 0
if depth is not None:
    anchors = cavity_anchors(instances, mask, keep_mask=keep_mask)
    dump("04_cavity_anchors.png", anchors)
    pockets = recessed_pocket_labels(mask, depth)
    print("pocket labels:", int(pockets.max()), "area:", [(int(((pockets == i) & (mask > 0)).sum())) for i in range(1, pockets.max() + 1)])
    mask2, n_vol = cut_anchored_pockets(mask, pockets, anchors, depth, dilate_px=max(2, round(3 * scale)), keep_mask=keep_mask)
    dump("05_mask_post_pockets.png", mask2)
    print(f"pockets cut: {n_vol}")
    mask = mask2
else:
    print("depth niedostepna")

lab = rgb_to_lab(arr)
mask3, n_cut, n_kept = subtract_protected(mask, instances, lab, delta_e=seg.neg_delta_e, dilate_px=max(2, round(3 * scale)), keep_mask=keep_mask)
dump("06_mask_post_subtract.png", mask3)
print(f"subtract: cut={n_cut} kept={n_kept}")

alpha = matte * feather(mask3, sigma=1.0)
dump("07_alpha_final.png", alpha)

# overlay na masterze
ov = arr.copy().astype(np.float32)
a3 = alpha[..., None]
ov = ov * (1 - 0.5 * a3) + np.array([255, 0, 0]) * 0.5 * a3
Image.fromarray(np.clip(ov, 0, 255).astype(np.uint8)).save(os.path.join(out_dir, "08_overlay.png"))
print("done ->", out_dir)