from __future__ import annotations

from pathlib import Path
import numpy as np
import cv2
from PIL import Image

# 解除大图限制
Image.MAX_IMAGE_PIXELS = None 

from model.Depth.depth_loader import build_depth_predictor, UnifiedDepthConfig, DepthBackend

# ================= 配置区 =================
INPUT_IMAGE = "/home/dwh/Documents/毕业设计/dwh/数据集/Up the River During Qingming (detail) - Court painters.jpg"
OUTPUT_DIR = "outputs/test_depth_v4"
DEPTH_BACKEND = DepthBackend.DEPTH_ANYTHING_V2

# 边缘捕捉参数
# 增大这个值会让边缘更细，减小会让边缘更粗（捕捉更多微弱信息）
EDGE_TOP_PERCENTILE = 96.0  # 选取梯度最强的前 7.0% 的像素
# 局部增强的灵敏度，建议在 2.0 - 5.0 之间
CLAHE_CLIP_LIMIT = 3.0 
# 形态学核大小
MORPH_SIZE = 5
# =========================================

def _extract_robust_edges(depth_norm: np.ndarray) -> np.ndarray:
    """
    通过局部增强和 Sobel 梯度提取闭合边缘
    """
    # 1. 转换为 8 位灰度
    depth_u8 = (depth_norm * 255).astype(np.uint8)
    
    # 2. 【核心步骤】CLAHE 局部自适应对比度增强
    # 这会强行放大古画中细微的建筑/人物深度差
    clahe = cv2.createCLAHE(clipLimit=CLAHE_CLIP_LIMIT, tileGridSize=(16, 16))
    enhanced_depth = clahe.apply(depth_u8)
    
    # 3. 高斯模糊：减少数字化噪声
    blurred = cv2.GaussianBlur(enhanced_depth, (5, 5), 0)
    
    # 4. Sobel 算子计算梯度强度
    grad_x = cv2.Sobel(blurred, cv2.CV_64F, 1, 0, ksize=3)
    grad_y = cv2.Sobel(blurred, cv2.CV_64F, 0, 1, ksize=3)
    grad_mag = np.sqrt(grad_x**2 + grad_y**2)
    
    # 5. 【统计学阈值】不再死守固定数值，而是选 Top X%
    threshold = np.percentile(grad_mag, EDGE_TOP_PERCENTILE)
    binary_edges = (grad_mag >= threshold).astype(np.uint8) * 255
    
    # 6. 形态学闭合：桥接裂缝，让线条连起来
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (MORPH_SIZE, MORPH_SIZE))
    closed = cv2.morphologyEx(binary_edges, cv2.MORPH_CLOSE, kernel)
    
    # 7. 再次轻微膨胀，为 SAM 提供更好的引导范围
    final_mask = cv2.dilate(closed, kernel, iterations=1)
    
    return final_mask

def main() -> None:
    base_dir = Path(__file__).resolve().parent
    img_path = Path(INPUT_IMAGE)
    out_dir = base_dir / OUTPUT_DIR
    out_dir.mkdir(parents=True, exist_ok=True)

    # 1. 初始化深度模型
    predictor, used_backend = build_depth_predictor(UnifiedDepthConfig(backend=DEPTH_BACKEND))
    
    # 2. 加载图像
    print(f"[Loading] 正在处理: {img_path.name}")
    img_pil = Image.open(img_path).convert("RGB")
    w_orig, h_orig = img_pil.size
    
    # 3. 深度预测
    print(f"[Depth] 正在进行深度估计 (Large Image)...")
    depth = np.asarray(predictor(img_pil), dtype=np.float32).squeeze()
    
    # 4. 归一化与保存
    dmin, dmax = depth.min(), depth.max()
    depth_norm = (depth - dmin) / (dmax - dmin + 1e-8)
    depth_u16 = (depth_norm * 65535.0).astype(np.uint16)
    Image.fromarray(depth_u16).save(out_dir / f"{img_path.stem}.depth.png")

    # 5. 提取强鲁棒性边缘
    print(f"[Edge] 正在应用 CLAHE + Sobel 增强算法提取边缘...")
    edge_mask = _extract_robust_edges(depth_norm)
    
    # 6. 导出
    mask_path = out_dir / f"{img_path.stem}.edge_mask_robust.png"
    Image.fromarray(edge_mask).save(mask_path)

    edge_ratio = float((edge_mask > 0).sum()) / float(edge_mask.size)
    print("-" * 30)
    print(f"提取完成！")
    print(f"边缘密度: {edge_ratio:.2%} (目标通常应在 1% ~ 8% 之间)")
    print(f"如果 Mask 依然太黑，请调低 EDGE_TOP_PERCENTILE (如 90.0)")
    print(f"如果 Mask 太乱，请调高 EDGE_TOP_PERCENTILE (如 96.0)")
    print("-" * 30)

if __name__ == "__main__":
    main()