import cv2
import numpy as np

def detect_stain_mask(img2, img1, diff_thresh=0.12, grad_thresh=0.03, var_thresh=0.005):
    """
    检测图2中的印染区域（相对于图1的差异 + 高梯度 + 高局部方差）
    
    Args:
        img2: 图2（清晰但有印染），uint8 RGB
        img1: 图1（模糊但纯净），uint8 RGB
    Returns:
        mask: 浮点掩码 [0,1]，1表示印染区域
    """
    # 转浮点
    f2 = img2.astype(np.float32) / 255.0
    f1 = img1.astype(np.float32) / 255.0
    
    # 1. 差异图
    diff = np.abs(f2 - f1)
    gray_diff = cv2.cvtColor(diff, cv2.COLOR_RGB2GRAY)
    
    # 2. 图2的拉普拉斯梯度（印染边缘高梯度）
    gray2 = cv2.cvtColor(f2, cv2.COLOR_RGB2GRAY)
    laplacian2 = cv2.Laplacian(gray2, cv2.CV_32F, ksize=3)
    grad_mag = np.abs(laplacian2)
    
    # 3. 图2的局部方差（印染区域局部对比度高）
    mean2 = cv2.GaussianBlur(f2, (0, 0), 5)
    sqr_mean2 = cv2.GaussianBlur(f2 * f2, (0, 0), 5)
    local_var = np.mean(sqr_mean2 - mean2 * mean2, axis=2)
    
    # 4. 组合条件
    mask_cond = (
        (gray_diff > diff_thresh) & 
        (grad_mag > grad_thresh) & 
        (local_var > var_thresh)
    ).astype(np.float32)
    
    # 5. 形态学操作：膨胀连接区域 + 高斯模糊羽化
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
    mask = cv2.dilate(mask_cond, kernel, iterations=2)
    mask = cv2.GaussianBlur(mask, (15, 15), 0)
    
    # 归一化到 [0, 1]
    mask = np.clip(mask, 0, 1)
    
    return mask


def fuse_textures(img2, img1, mask, sharpen_img1=True):
    """
    融合两张纹理图
    
    Args:
        img2: 底图（清晰但有印染）
        img1: 补丁图（模糊但纯净）
        mask: 印染掩码 [0,1]
        sharpen_img1: 是否对图1做轻微锐化以减少"糊"感
    Returns:
        fused: 融合结果
    """
    f2 = img2.astype(np.float32) / 255.0
    f1 = img1.astype(np.float32) / 255.0
    
    # 可选：对图1做轻微锐化
    if sharpen_img1:
        # 非锐化掩模
        blurred = cv2.GaussianBlur(f1, (0, 0), 1.5)
        sharpened = cv2.addWeighted(f1, 1.5, blurred, -0.5, 0)
        f1_sharp = np.clip(sharpened, 0, 1)
    else:
        f1_sharp = f1
    
    # 融合：mask区域用图1，其余保留图2
    mask_3ch = np.stack([mask, mask, mask], axis=2)
    fused = f2 * (1.0 - mask_3ch) + f1_sharp * mask_3ch
    
    # 转回 uint8
    fused = np.clip(fused * 255.0, 0, 255).astype(np.uint8)
    
    return fused


def post_process_clahe(img, clip_limit=1.5, grid_size=(8, 8), mix_ratio=0.3):
    """
    轻微 CLAHE 后处理（仅 V 通道），与原图混合
    """
    hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV).astype(np.float32)
    h, s, v = cv2.split(hsv)
    
    # CLAHE 需要 uint8
    v_u8 = (v / 255.0 * 255).astype(np.uint8)
    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=grid_size)
    v_adjusted = clahe.apply(v_u8).astype(np.float32)
    
    # 混合
    v_mixed = v * (1 - mix_ratio) + v_adjusted * mix_ratio
    v_mixed = np.clip(v_mixed, 0, 255)
    
    hsv_out = cv2.merge([h, s, v_mixed]).astype(np.uint8)
    return cv2.cvtColor(hsv_out, cv2.COLOR_HSV2RGB)


def main():
    # === 参数配置 ===
    img1_path = "/home/algo/Documents/openMVS/openMVS/cartoon_texture/data/642713/642713.o.jpg"   # 图1：模糊但纯净
    img2_path = "/home/algo/Documents/openMVS/openMVS/cartoon_texture/data/642713/642713.mesh7.png"  # 图2：清晰但有印染
    output_path = "/home/algo/Documents/openMVS/openMVS/cartoon_texture/data/642713/642713.jpg"
    
    # 印染检测参数（可按实际效果微调）
    DIFF_THRESH = 0.12      # 差异阈值
    GRAD_THRESH = 0.03      # 梯度阈值
    VAR_THRESH = 0.005      # 局部方差阈值
    
    # 后处理参数
    APPLY_CLAHE = True      # 是否做轻微CLAHE
    SHARPEN_IMG1 = True     # 是否锐化图1
    
    # === 加载图片 ===
    img2 = cv2.imread(img2_path)
    img1 = cv2.imread(img1_path)
    
    if img2 is None or img1 is None:
        print("Error: Cannot load images")
        return
    
    # 确保尺寸一致（以图2为准）
    if img1.shape != img2.shape:
        print(f"Resizing img1 from {img1.shape} to {img2.shape}")
        img1 = cv2.resize(img1, (img2.shape[1], img2.shape[0]))
    
    # 转 RGB（OpenCV 默认 BGR）
    img2_rgb = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
    img1_rgb = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
    
    print("Detecting stain mask...")
    mask = detect_stain_mask(img2_rgb, img1_rgb, 
                             diff_thresh=DIFF_THRESH, 
                             grad_thresh=GRAD_THRESH, 
                             var_thresh=VAR_THRESH)
    
    print(f"Mask coverage: {np.mean(mask > 0.5) * 100:.1f}%")
    
    print("Fusing textures...")
    fused = fuse_textures(img2_rgb, img1_rgb, mask, sharpen_img1=SHARPEN_IMG1)
    
    if APPLY_CLAHE:
        print("Applying CLAHE post-processing...")
        fused = post_process_clahe(fused)
    
    # 保存结果
    fused_bgr = cv2.cvtColor(fused, cv2.COLOR_RGB2BGR)
    cv2.imwrite(output_path, fused_bgr)
    
    # 保存掩码可视化（调试用）
    mask_viz = (mask * 255).astype(np.uint8)
    cv2.imwrite("stain_mask_debug.png", mask_viz)
    
    print(f"Done! Result saved to {output_path}")
    print(f"Debug mask saved to stain_mask_debug.png")
    
    # 显示结果
    cv2.imshow("Original (mesh7.1)", img2)
    cv2.imshow("Reference (n1)", img1)
    cv2.imshow("Stain Mask", mask_viz)
    cv2.imshow("Fused Result", fused_bgr)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()