"""Test domykania maski: (a) GrabCut z inicjalizacji maska, (b) unia SAM 'sofa'.
Uzycie: python tmp/dbg_grow.py <master.png> <matte.png| -> recompute>
"""
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

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

rgb = np.asarray(Image.open(src).convert("RGB"), dtype=np.uint8)
matte = np.asarray(Image.open(sys.argv[3]).convert("L"), dtype=np.uint8).astype(np.float32) / 255.0
h, w = matte.shape
if rgb.shape[:2] != (h, w):
    rgb = cv2.resize(rgb, (w, h), interpolation=cv2.INTER_LINEAR)

mask = (matte >= 0.5).astype(np.uint8)

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}")

# --- (a) GrabCut: sure-FG = erodowana maska; probable-FG = dylatacja; reszta sure-BG poza pasmem
gc = np.full((h, w), cv2.GC_BGD, np.uint8)
k_fg = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
sure_fg = cv2.erode(mask, k_fg)
band = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (121, 121)))
gc[sure_fg > 0] = cv2.GC_FGD
gc[(band > 0) & (sure_fg == 0)] = cv2.GC_PR_FGD
# cienki pas tuz za maska = probable BG zamiast BGD, zeby mogl sie rozszerzyc
gc[(band > 0) & (mask == 0)] = cv2.GC_PR_BGD
bgm = np.zeros((1, 65), np.float64); fgm = np.zeros((1, 65), np.float64)
bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
cv2.grabCut(bgr, gc, None, bgm, fgm, 4, cv2.GC_INIT_WITH_MASK)
gc_mask = np.where((gc == cv2.GC_FGD) | (gc == cv2.GC_PR_FGD), 1, 0).astype(np.uint8)
dump("gc_mask.png", gc_mask)

# roznica: co dostawil grabcut
added = gc_mask & (1 - mask)
dump("gc_added.png", added)

ov = rgb.copy().astype(np.float32)
a3 = gc_mask[..., None].astype(np.float32)
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, "gc_overlay.png"))

# --- (b) unia z SAM z boxa DINO 'sofa'
from fabric_recolor.segmentation.grounded_sam import GroundedSamSegmenter
gs = GroundedSamSegmenter(device="cuda", max_side=1600)
gs._load()
det = gs._detect(Image.fromarray(rgb), ["sofa", "couch", "settee"])
labels = det.get("text_labels") or det.get("labels") or []
print("DINO:", labels, det["scores"].cpu().numpy())
boxes = det["boxes"].cpu().numpy()
masks = gs._boxes_to_masks(Image.fromarray(rgb), boxes)
if len(masks):
    sam = masks.any(axis=0).astype(np.uint8)
    dump("sam_sofa.png", sam)
    uni = np.maximum(mask, sam)
    dump("union.png", uni)
    ov2 = rgb.copy().astype(np.float32)
    a4 = uni[..., None].astype(np.float32)
    ov2 = ov2 * (1 - 0.5 * a4) + np.array([0, 128, 255]) * 0.5 * a4
    Image.fromarray(np.clip(ov2, 0, 255).astype(np.uint8)).save(os.path.join(out_dir, "union_overlay.png"))
print("done")