"""
Guarded background masking with U^2-Net (rembg).

For every image: strip near-black padding, run salient-object segmentation,
then ACCEPT the mask only if it clearly removes a plain photographic backdrop
around the artwork — never content. Five checks, all must pass:

  1. kept fraction in [0.20, 0.90]      — mask keeps a plausible artwork share
  2. no interior holes (> 2%)           — artwork regions are never punched out
  3. convex solidity >= 0.97            — one solid blob, not scattered figures
  4. mask must not touch-fill the border — something around it was removed
  5. removed pixels are uniform (std <= 28) — what's removed looks like backdrop

Accepted masks ("applied") are saved as PNGs + a bbox; everything else is
"rejected" and downstream features use the full image. On our gold set this
applies to ~10% of images (museum photos of framed/mounted works).

Output: data/masks/<stem>.png + data/masks/verdicts.csv
        (filename, verdict, y0, y1, x0, x1)
Usage:  python preprocessing/generate_masks.py
"""

import csv
from pathlib import Path

import cv2
import numpy as np
import pandas as pd
from PIL import Image
from rembg import new_session, remove
from tqdm import tqdm

IMAGES = Path("data/images")
SELECTED = Path("data/artwork_metadata.csv")
MASK_DIR = Path("data/masks")
VERDICTS = MASK_DIR / "verdicts.csv"

Image.MAX_IMAGE_PIXELS = None


def crop_padding(img_rgb, threshold=5):
    gray = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    rows = np.where(gray.max(axis=1) > threshold)[0]
    cols = np.where(gray.max(axis=0) > threshold)[0]
    if len(rows) == 0 or len(cols) == 0:
        return img_rgb
    return img_rgb[rows[0]:rows[-1] + 1, cols[0]:cols[-1] + 1]


def mask_verdict(raw, img, lo=0.20, hi=0.90, max_bg_std=28):
    """True (apply) only if the mask removes a solid, uniform border region."""
    if raw is None:
        return False
    m = (raw > 0).astype(np.uint8)
    kept = m.mean()
    if not (lo <= kept <= hi):
        return False
    cnts, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not cnts:
        return False
    filled = m.copy()
    cv2.drawContours(filled, cnts, -1, 1, -1)
    if (filled - m).sum() / max(m.sum(), 1) > 0.02:
        return False
    hull = cv2.convexHull(np.vstack([c.reshape(-1, 2) for c in cnts]))
    if cv2.contourArea(hull) == 0 or m.sum() / cv2.contourArea(hull) < 0.97:
        return False
    border = np.zeros_like(m)
    border[0, :] = border[-1, :] = border[:, 0] = border[:, -1] = 1
    if (border & (1 - m)).sum() == 0:
        return False
    removed = img[m == 0]
    if removed.std(axis=0).mean() > max_bg_std:
        return False
    return True


def main():
    MASK_DIR.mkdir(parents=True, exist_ok=True)
    sel = pd.read_csv(SELECTED, dtype=str).drop_duplicates("filename")
    done = set()
    if VERDICTS.exists():
        done = set(pd.read_csv(VERDICTS, dtype=str)["filename"])
    todo = [f for f in sel["filename"] if f not in done]
    print(f"total={len(sel)}  done={len(done)}  todo={len(todo)}")

    session = new_session("u2net")
    mode = "a" if VERDICTS.exists() else "w"
    with open(VERDICTS, mode, newline="") as fh:
        writer = csv.writer(fh)
        if mode == "w":
            writer.writerow(["filename", "verdict", "y0", "y1", "x0", "x1"])
        applied = rejected = 0
        for fn in tqdm(todo):
            try:
                img = crop_padding(np.array(Image.open(IMAGES / fn).convert("RGB")))
                raw = np.array(remove(Image.fromarray(img), session=session,
                                      only_mask=True))
                if mask_verdict(raw, img):
                    m = (raw > 0).astype(np.uint8)
                    ys, xs = np.where(m > 0)
                    y0, y1, x0, x1 = ys.min(), ys.max(), xs.min(), xs.max()
                    cv2.imwrite(str(MASK_DIR / (Path(fn).stem + ".png")), m * 255)
                    writer.writerow([fn, "applied", y0, y1, x0, x1])
                    applied += 1
                else:
                    writer.writerow([fn, "rejected", "", "", "", ""])
                    rejected += 1
            except Exception as e:
                print(f"FAIL {fn}: {e}")
                writer.writerow([fn, "rejected", "", "", "", ""])
                rejected += 1
            fh.flush()
    print(f"applied={applied}  rejected={rejected}")


if __name__ == "__main__":
    main()