#!/usr/bin/env python3
"""Guided-filter occluder refinement using depth + RGB guidance."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import cv2
import numpy as np


def guided_filter_gray(I: np.ndarray, p: np.ndarray, radius: int, eps: float) -> np.ndarray:
    """Fast guided filter (grayscale guidance) fallback when ximgproc is unavailable."""
    r = max(1, int(radius))
    ksize = (r * 2 + 1, r * 2 + 1)
    mean_I = cv2.boxFilter(I, -1, ksize, normalize=True)
    mean_p = cv2.boxFilter(p, -1, ksize, normalize=True)
    mean_Ip = cv2.boxFilter(I * p, -1, ksize, normalize=True)
    cov_Ip = mean_Ip - mean_I * mean_p
    mean_II = cv2.boxFilter(I * I, -1, ksize, normalize=True)
    var_I = mean_II - mean_I * mean_I
    a = cov_Ip / (var_I + eps)
    b = mean_p - a * mean_I
    mean_a = cv2.boxFilter(a, -1, ksize, normalize=True)
    mean_b = cv2.boxFilter(b, -1, ksize, normalize=True)
    q = mean_a * I + mean_b
    return q


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Refine depth occluder mask using guided filtering.")
    parser.add_argument("--depth", required=True, help="Depth map (grayscale PNG)")
    parser.add_argument("--image", required=True, help="Original RGB image")
    parser.add_argument("--output", required=True, help="Output mask path (PNG)")
    parser.add_argument("--json", help="Optional JSON metadata output path")
    parser.add_argument("--radius", type=int, default=4, help="Guided filter radius")
    parser.add_argument("--eps", type=float, default=50.0, help="Guided filter epsilon")
    parser.add_argument("--erode", type=int, default=0, help="Erode radius after threshold")
    parser.add_argument("--min-coverage", type=float, default=0.30, help="Minimum coverage target")
    parser.add_argument("--retry-bias", type=float, default=18.0, help="Lower Otsu threshold by this value when coverage is low")
    parser.add_argument("--retry-steps", type=int, default=3, help="Retry steps when coverage is low")
    parser.add_argument("--dilate-low", type=int, default=1, help="Dilate radius if coverage is still low after retries")
    parser.add_argument("--invert", action="store_true", help="Invert depth before threshold")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    depth = cv2.imread(args.depth, cv2.IMREAD_GRAYSCALE)
    if depth is None:
        raise RuntimeError("failed to read depth image")
    image = cv2.imread(args.image, cv2.IMREAD_COLOR)
    if image is None:
        raise RuntimeError("failed to read guidance image")

    h, w = image.shape[:2]
    depth_resized = cv2.resize(depth, (w, h), interpolation=cv2.INTER_CUBIC)
    if args.invert:
        depth_resized = 255 - depth_resized

    radius = max(1, int(args.radius))
    eps = max(1.0, float(args.eps))
    if hasattr(cv2, "ximgproc") and hasattr(cv2.ximgproc, "guidedFilter"):
        guided = cv2.ximgproc.guidedFilter(image, depth_resized, radius, eps)
    else:
        guide_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY).astype(np.float32)
        depth_f = depth_resized.astype(np.float32)
        guided = guided_filter_gray(guide_gray, depth_f, radius, eps)
    guided_u8 = cv2.normalize(guided, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)

    otsu_thresh, mask = cv2.threshold(guided_u8, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    coverage = float((mask > 0).sum()) / float(mask.size)
    min_cov = max(0.0, min(1.0, float(args.min_coverage)))
    if coverage < min_cov:
        bias = max(0.0, float(args.retry_bias))
        steps = max(1, int(args.retry_steps))
        if bias > 0:
            for idx in range(1, steps + 1):
                retry_thresh = max(0.0, float(otsu_thresh) - bias * idx)
                _, mask = cv2.threshold(guided_u8, retry_thresh, 255, cv2.THRESH_BINARY)
                coverage = float((mask > 0).sum()) / float(mask.size)
                if coverage >= min_cov:
                    break
        if coverage < min_cov and int(args.dilate_low) > 0:
            k = max(1, int(args.dilate_low)) * 2 + 1
            kernel = np.ones((k, k), np.uint8)
            mask = cv2.dilate(mask, kernel, iterations=1)
            coverage = float((mask > 0).sum()) / float(mask.size)
    if args.erode > 0:
        k = max(1, int(args.erode)) * 2 + 1
        kernel = np.ones((k, k), np.uint8)
        mask = cv2.erode(mask, kernel, iterations=1)

    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    cv2.imwrite(str(out_path), mask)

    if args.json:
        meta = {
            "ok": True,
            "width": int(w),
            "height": int(h),
            "coverage": coverage,
            "otsu": float(otsu_thresh),
            "radius": radius,
            "eps": eps,
            "erode": int(args.erode),
            "minCoverage": float(args.min_coverage),
            "retryBias": float(args.retry_bias),
            "retrySteps": int(args.retry_steps),
            "dilateLow": int(args.dilate_low),
            "invert": bool(args.invert),
        }
        Path(args.json).write_text(json.dumps(meta, indent=2), encoding="utf8")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())