"""
calibrate.py —— 在跑 OpenMVS 之前，对输入照片做白平衡 + 曝光统一。
这是消除"源视图色温差异"最彻底、最安全的办法：在像素进入纹理管线之前就对齐。
用法: python calibrate.py <image_dir> <out_dir>

做了两件事（都可选、都保守，不会破坏细节）：
1) Gray-World 白平衡：把每张图的通道均值拉向全局中值（去色温差）
2) 全局曝光归一化：把所有图的亮度中位数对齐到全局中位数（去曝光差）

注意：OpenMVS 的 EstimateGlobalPhotometricCorrection 强依赖稠密点云，
如果没有点云它会被跳过（你之前日志里 "Photometric: no point cloud, skip"）。
所以在这里预处理是最稳的兜底。
"""
import os, sys, glob
import numpy as np
import cv2


def gray_world(img_bgr, ref_bgr=None):
    """灰度世界白平衡。

    若给定 ref_bgr（全局参考图的 BGR 均值），则把本图通道均值拉向 ref；
    否则用本图三通道的均值作为中性灰参考（单图自校正）。
    不做任何“保守衰减”——必须完全对齐，否则等于没做。
    """
    b, g, r = [img_bgr[:, :, c].astype(np.float32) for c in range(3)]
    mb, mg, mr = b.mean(), g.mean(), r.mean()
    if ref_bgr is not None:
        rb, rg, rr = ref_bgr
    else:
        # 单图自校正：以三通道共同均值作为中性灰目标
        m = (mb + mg + mr) / 3.0
        rb, rg, rr = m, m, m
    kb = rb / mb if mb > 1e-6 else 1.0
    kg = rg / mg if mg > 1e-6 else 1.0
    kr = rr / mr if mr > 1e-6 else 1.0
    # 限制增益上限，防止极端单图把噪声放大
    def clamp_gain(k):
        return min(max(k, 0.3), 3.0)
    kb, kg, kr = clamp_gain(kb), clamp_gain(kg), clamp_gain(kr)
    out = np.stack([
        np.clip(b * kb, 0, 255),
        np.clip(g * kg, 0, 255),
        np.clip(r * kr, 0, 255),
    ], axis=2).astype(np.uint8)
    return out


def normalize_exposure(img_bgr, target_median):
    """曝光归一化：按亮度中位数线性缩放，对齐到全局目标中位数。
    只做温和限制（上限 2.0），保证不同视图曝光真正一致。"""
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY).astype(np.float32)
    med = np.median(gray)
    if med < 1.0:
        return img_bgr
    scale = target_median / med
    scale = min(scale, 2.0)   # 防止极暗图放大噪声，但允许充分提亮
    out = np.clip(img_bgr.astype(np.float32) * scale, 0, 255).astype(np.uint8)
    return out


def main():
    if len(sys.argv) < 3:
        print("usage: python calibrate.py <image_dir> <out_dir>")
        sys.exit(1)
    in_dir, out_dir = sys.argv[1], sys.argv[2]
    os.makedirs(out_dir, exist_ok=True)

    paths = sorted(glob.glob(os.path.join(in_dir, "*.*")))
    paths = [p for p in paths if p.lower().endswith((".jpg", ".jpeg", ".png", ".tif", ".bmp"))]
    if not paths:
        print("no images found in", in_dir)
        sys.exit(1)

    # ---- 第一遍：统计所有图，确定全局参考（目标白平衡 + 目标亮度）----
    means = []   # (mb, mg, mr)
    meds = []
    for p in paths:
        img = cv2.imread(p)
        if img is None:
            continue
        b, g, r = img[:, :, 0].astype(np.float32), img[:, :, 1].astype(np.float32), img[:, :, 2].astype(np.float32)
        means.append((b.mean(), g.mean(), r.mean()))
        gray = 0.299 * r + 0.587 * g + 0.114 * b
        meds.append(np.median(gray))

    means = np.array(means)
    meds = np.array(meds)

    # 全局参考 = 所有图通道均值的均值（= 数据集整体真实色彩），亮度取中位数
    ref_b = means[:, 0].mean()
    ref_g = means[:, 1].mean()
    ref_r = means[:, 2].mean()
    ref_bgr = (ref_b, ref_g, ref_r)
    target_median = float(np.median(meds))

    print(f"[calibrate] {len(paths)} images")
    print(f"[calibrate] global reference white-balance (B,G,R) = ({ref_b:.1f}, {ref_g:.1f}, {ref_r:.1f})")
    print(f"[calibrate] global reference luminance median = {target_median:.1f}")

    # ---- 第二遍：所有图统一校正到同一参考 ----
    for i, p in enumerate(paths):
        img = cv2.imread(p)
        if img is None:
            continue

        # 1) 白平衡：把本图拉向全局参考 ref_bgr
        out = gray_world(img, ref_bgr=ref_bgr)

        # 2) 曝光归一化：亮度中位数对齐到全局目标
        out = normalize_exposure(out, target_median)

        name = os.path.basename(p)
        cv2.imwrite(os.path.join(out_dir, name), out)

    print(f"[calibrate] done -> {out_dir}")


if __name__ == "__main__":
    main()