#!/usr/bin/env python3
"""
Validate background removal strategies.

Usage:
    python test_transparency.py [image.png ...]

Writes comparison files next to each input:
    *_rembg.png          — pure rembg (bg_removal=rembg path)
    *_blackbg.png        — simulated black-bg composite (what Qwen renders in sam2 mode)
    *_thresh.png         — threshold mask only (non-black pixels → person)
    *_thresh_sam2.png    — threshold bbox → SAM2 edge refinement (new sam2 mode path)
"""
import io, sys, os
import numpy as np
from PIL import Image, ImageFilter

OUTPUT_DIR = "/mnt/zim/tour-comfy/output"
VENV_SITE  = "/home/mike/comfyui/venv/lib/python3.13/site-packages"
SAM2_CKPT  = os.path.expanduser("~/.sam/sam2.1_hiera_base_plus.pt")
SAM2_CFG   = "configs/sam2.1/sam2.1_hiera_b+.yaml"


# ── rembg ──────────────────────────────────────────────────────────────────────
def apply_rembg(png_bytes: bytes) -> bytes:
    from rembg import remove
    return remove(png_bytes)


# ── SAM2 loader ────────────────────────────────────────────────────────────────
_predictor = None
def load_sam2():
    global _predictor
    if _predictor is not None:
        return _predictor
    try:
        import torch
        from sam2.build_sam import build_sam2
        from sam2.sam2_image_predictor import SAM2ImagePredictor
        model = build_sam2(SAM2_CFG, SAM2_CKPT, device="cuda")
        _predictor = SAM2ImagePredictor(model)
        print("[sam2] loaded")
    except Exception as e:
        print(f"[sam2] FAILED: {e}")
        _predictor = False
    return _predictor


# ── Simulate black-bg Qwen output ─────────────────────────────────────────────
def make_black_bg(png_bytes: bytes) -> bytes:
    """Composite a rembg cutout onto pure black — simulates Qwen 'black background' output."""
    rgba = Image.open(io.BytesIO(apply_rembg(png_bytes))).convert("RGBA")
    bg   = Image.new("RGBA", rgba.size, (0, 0, 0, 255))
    bg.paste(rgba, mask=rgba.split()[3])
    out  = bg.convert("RGB")
    buf  = io.BytesIO(); out.save(buf, "PNG"); return buf.getvalue()


# ── Threshold-only mask ────────────────────────────────────────────────────────
def apply_threshold_mask(png_bytes: bytes, threshold: int = 25) -> bytes:
    """Find non-black pixels → person mask.  No SAM2 needed."""
    img  = Image.open(io.BytesIO(png_bytes)).convert("RGB")
    arr  = np.array(img)
    h, w = arr.shape[:2]

    is_person = np.max(arr, axis=2) > threshold
    coverage  = is_person.sum() / (h * w)
    print(f"  [threshold] person coverage: {coverage:.1%}")

    if not is_person.any():
        print("  [threshold] all-black image — no person found")
        return png_bytes

    mask_np   = is_person.astype(np.uint8) * 255
    alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=2))
    rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
    out  = Image.merge("RGBA", (r, g, b, alpha_img))
    buf  = io.BytesIO(); out.save(buf, "PNG"); return buf.getvalue()


# ── NEW: Threshold bbox → SAM2 refinement (sam2 mode path) ────────────────────
def apply_thresh_sam2(png_bytes: bytes, threshold: int = 25) -> bytes:
    """
    For black-background Qwen output:
      1. Threshold to find person bbox (non-black pixels)
      2. Run SAM2 with that tight bbox for clean edge refinement
      3. Fallback to threshold mask if SAM2 unavailable or mask looks wrong
    """
    import torch

    img  = Image.open(io.BytesIO(png_bytes)).convert("RGB")
    arr  = np.array(img)
    h, w = arr.shape[:2]

    # Step 1 — threshold
    is_person  = np.max(arr, axis=2) > threshold
    thresh_cov = is_person.sum() / (h * w)
    print(f"  [thresh_sam2] threshold person coverage: {thresh_cov:.1%}")

    if not is_person.any():
        print("  [thresh_sam2] all-black — fallback to rembg")
        return apply_rembg(png_bytes)

    rows = np.any(is_person, axis=1)
    cols = np.any(is_person, axis=0)
    rmin = int(np.where(rows)[0][0]);  rmax = int(np.where(rows)[0][-1])
    cmin = int(np.where(cols)[0][0]);  cmax = int(np.where(cols)[0][-1])

    margin = int(min(h, w) * 0.02)
    y1 = max(0, rmin - margin); y2 = min(h, rmax + margin)
    x1 = max(0, cmin - margin); x2 = min(w, cmax + margin)
    print(f"  [thresh_sam2] person bbox (+margin): ({x1},{y1})-({x2},{y2})")

    # Step 2 — SAM2 with person-specific bbox
    predictor = load_sam2()
    if predictor is not False:
        box = np.array([[x1, y1, x2, y2]], dtype=np.float32)
        try:
            with torch.inference_mode():
                predictor.set_image(arr)
                masks, scores, _ = predictor.predict(box=box, multimask_output=True)

            if masks is not None and len(masks) > 0:
                best     = masks[int(np.argmax(scores))]
                sam_cov  = float(best.sum()) / (h * w)
                print(f"  [thresh_sam2] SAM2 coverage: {sam_cov:.1%}  (threshold was {thresh_cov:.1%})")

                # Accept SAM2 result if coverage is within reasonable range of threshold
                if 0.03 < sam_cov < 0.95 and abs(sam_cov - thresh_cov) < 0.30:
                    mask_np   = best.astype(np.uint8) * 255
                    alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=1))
                    rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
                    out  = Image.merge("RGBA", (r, g, b, alpha_img))
                    buf  = io.BytesIO(); out.save(buf, "PNG")
                    print("  [thresh_sam2] SAM2 result accepted ✓")
                    return buf.getvalue()
                else:
                    print(f"  [thresh_sam2] SAM2 coverage diverged from threshold — using threshold mask")
        except Exception as e:
            print(f"  [thresh_sam2] SAM2 error: {e} — using threshold mask")
    else:
        print("  [thresh_sam2] SAM2 not available — using threshold mask")

    # Step 3 — fallback: threshold mask with soft edges
    mask_np   = is_person.astype(np.uint8) * 255
    alpha_img = Image.fromarray(mask_np, "L").filter(ImageFilter.GaussianBlur(radius=2))
    rgba = img.convert("RGBA"); r, g, b, _ = rgba.split()
    out  = Image.merge("RGBA", (r, g, b, alpha_img))
    buf  = io.BytesIO(); out.save(buf, "PNG")
    print("  [thresh_sam2] threshold mask used as fallback")
    return buf.getvalue()


# ── main ───────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    paths = sys.argv[1:] if len(sys.argv) > 1 else [
        os.path.join(OUTPUT_DIR, "20260622_181910_0_20260619_124038_image.png"),
    ]

    for path in paths:
        if not os.path.exists(path):
            print(f"SKIP (not found): {path}"); continue
        stem = os.path.splitext(path)[0]
        print(f"\n══ {os.path.basename(path)} ══")
        with open(path, "rb") as f:
            raw = f.read()

        print("1. rembg (bg_removal=rembg path)...")
        rb = apply_rembg(raw)
        with open(stem + "_rembg.png", "wb") as f: f.write(rb)
        print(f"   → {os.path.basename(stem)}_rembg.png")

        print("2. Simulate black-bg Qwen output...")
        bb = make_black_bg(raw)
        with open(stem + "_blackbg.png", "wb") as f: f.write(bb)
        print(f"   → {os.path.basename(stem)}_blackbg.png")

        print("3. Threshold-only mask on black-bg image...")
        tm = apply_threshold_mask(bb)
        with open(stem + "_thresh.png", "wb") as f: f.write(tm)
        print(f"   → {os.path.basename(stem)}_thresh.png")

        print("4. Threshold bbox → SAM2 refinement on black-bg image (NEW sam2 mode path)...")
        ts = apply_thresh_sam2(bb)
        with open(stem + "_thresh_sam2.png", "wb") as f: f.write(ts)
        print(f"   → {os.path.basename(stem)}_thresh_sam2.png")

    print("\n── Done ──")
    print("  *_rembg.png       rembg on real background (bg_removal=rembg path)")
    print("  *_thresh.png      threshold-only on black bg")
    print("  *_thresh_sam2.png threshold-bbox → SAM2 on black bg (NEW sam2 mode path)")