#!/usr/bin/env python3
"""Trimap-based image matting utilities.

This module implements classic trimap matting using the pymatting library.
It allows preserving entire objects from sparse interior marks by
morphological reconstruction to expand seeds before running matting.

Dependencies: Pillow, numpy, opencv-python, pymatting
"""

from pathlib import Path
from typing import Optional

import numpy as np
from PIL import Image

try:
    import cv2  # type: ignore
except Exception as e:  # pragma: no cover
    cv2 = None

try:
    from pymatting.alpha.estimate_alpha_knn import estimate_alpha_knn  # KNN matting
except Exception as e:  # pragma: no cover
    estimate_alpha_knn = None


def _load_gray(path: str) -> np.ndarray:
    """Load an image as 8-bit grayscale numpy array."""
    arr = np.array(Image.open(path).convert("L"))
    return arr


def build_trimap(
    image_path: str,
    *,
    keep_mask_path: Optional[str] = None,
    bg_mask_path: Optional[str] = None,
    keep_dilate: int = 2,
    bg_dilate: int = 2,
    auto_bg_from_edges: bool = True,
    white_bg_threshold: int = 240,
    auto_fg_from_nonwhite: bool = True,
    nonwhite_threshold: int = 25,
    bg_mode: str = "dilate",  # 'dilate' or 'graphcut'
    grabcut_iters: int = 5,
) -> np.ndarray:
    """Construct a trimap from optional foreground/background seeds.

    Returns a uint8 array with values: 0 (background), 128 (unknown), 255 (foreground).
    """
    if cv2 is None:
        raise RuntimeError("OpenCV (cv2) is required for trimap construction. Install opencv-python.")

    rgb = np.array(Image.open(image_path).convert("RGB"))
    h, w, _ = rgb.shape
    trimap = np.full((h, w), 128, dtype=np.uint8)  # start as unknown

    k = np.ones((3, 3), np.uint8)

    # Foreground seeds
    seeds_fg = None
    if keep_mask_path:
        seeds_fg = (_load_gray(keep_mask_path) > 0).astype(np.uint8) * 255
        if keep_dilate > 0:
            seeds_fg = cv2.dilate(seeds_fg, k, iterations=int(keep_dilate))
    elif auto_fg_from_nonwhite:
        diff = np.abs(rgb.astype(np.int16) - 255)
        maxdiff = diff.max(axis=2).astype(np.uint8)
        _, seeds_fg = cv2.threshold(maxdiff, int(nonwhite_threshold), 255, cv2.THRESH_BINARY)
        if keep_dilate > 0:
            seeds_fg = cv2.dilate(seeds_fg, k, iterations=int(keep_dilate))
    if seeds_fg is not None:
        trimap[seeds_fg > 0] = 255

    # Background seeds
    seeds_bg = None
    if bg_mask_path:
        seeds_bg = (_load_gray(bg_mask_path) > 0).astype(np.uint8) * 255
    elif auto_bg_from_edges:
        # Auto background from edges: assume near-white background connected to borders
        gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
        bg_candidate = (gray >= white_bg_threshold).astype(np.uint8) * 255
        mask = np.zeros((h + 2, w + 2), np.uint8)
        seeds_bg = np.zeros_like(bg_candidate)
        for sx, sy in [(0, 0), (w - 1, 0), (0, h - 1), (w - 1, h - 1)]:
            if bg_candidate[sy, sx] > 0:
                tmp = bg_candidate.copy()
                cv2.floodFill(tmp, mask.copy(), (sx, sy), 127)
                seeds_bg[tmp == 127] = 255
    # Two strategies to convert seeds to trimap
    if bg_mode == "graphcut" and (seeds_bg is not None or seeds_fg is not None):
        # Initialize GrabCut mask
        GC_BGD, GC_FGD, GC_PR_BGD, GC_PR_FGD = 0, 1, 2, 3
        gmask = np.full((h, w), GC_PR_BGD, dtype=np.uint8)
        if seeds_bg is not None:
            gmask[seeds_bg > 0] = GC_BGD
        if seeds_fg is not None:
            gmask[seeds_fg > 0] = GC_FGD
        bgdModel = np.zeros((1, 65), np.float64)
        fgdModel = np.zeros((1, 65), np.float64)
        # GrabCut expects BGR
        bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
        cv2.grabCut(bgr, gmask, None, bgdModel, fgdModel, max(1, int(grabcut_iters)), cv2.GC_INIT_WITH_MASK)
        # Convert to trimap: treat probable background as background to ensure known BG
        trimap[:, :] = 128
        trimap[(gmask == GC_BGD) | (gmask == GC_PR_BGD)] = 0
        trimap[gmask == GC_FGD] = 255
        # Always enforce original seeds as known
        if seeds_bg is not None:
            trimap[seeds_bg > 0] = 0
        if seeds_fg is not None:
            trimap[seeds_fg > 0] = 255
    else:
        # Simple dilation expansion
        if seeds_bg is not None:
            sbg = seeds_bg.copy()
            if bg_dilate > 0:
                sbg = cv2.dilate(sbg, k, iterations=int(bg_dilate))
            trimap[sbg > 0] = 0

    return trimap


def apply_trimap_matting(
    image_path: str,
    trimap: np.ndarray,
    output_path: str,
) -> Path:
    """Run closed-form matting using the provided trimap and save RGBA PNG."""
    if estimate_alpha_knn is None:
        raise RuntimeError("pymatting is required. Install with 'pip install pymatting'.")

    img = np.array(Image.open(image_path).convert("RGB"), dtype=np.float64) / 255.0
    if trimap.ndim == 3:
        trimap = cv2.cvtColor(trimap, cv2.COLOR_BGR2GRAY)

    # Convert trimap to float in [0,1]: 0 (BG), 0.5 (Unknown), 1 (FG)
    t = trimap.astype(np.float32) / 255.0
    t[t < 0.25] = 0.0
    t[(t >= 0.25) & (t <= 0.75)] = 0.5
    t[t > 0.75] = 1.0

    # Estimate alpha (0..1) using KNN matting (robust and no heavy tuning)
    alpha = estimate_alpha_knn(img, t)

    # Compose RGBA
    rgba = np.concatenate([img, alpha[..., None]], axis=2)
    rgba = (np.clip(rgba, 0.0, 1.0) * 255.0).astype(np.uint8)
    out = Image.fromarray(rgba, mode="RGBA")
    out.save(output_path, "PNG")
    return Path(output_path)
