#!/usr/bin/env python3
"""光度标定入口脚本

用法:
    python run_calibration.py \
        --model_dir /path/to/colmap_model \
        --rgb_dir /path/to/rgb \
        --mask_dir /path/to/masks \
        --output_dir /path/to/output
"""

import argparse
import sys
from pathlib import Path

from calib_io import run_calibration
from calibration import calibrate, CalibrationOptions


def parse_args():
    parser = argparse.ArgumentParser(
        description="Photometric calibration using COLMAP reconstruction",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--model_dir", "-m",
        type=str, required=True,
        help="COLMAP sparse reconstruction directory (contains cameras.bin/images.bin/points3D.bin)",
    )
    parser.add_argument(
        "--rgb_dir", "-r",
        type=str, required=True,
        help="Directory containing input RGB images",
    )
    parser.add_argument(
        "--mask_dir", "-mk",
        type=str, required=True,
        help="Directory containing foreground masks (.png, same names as RGB)",
    )
    parser.add_argument(
        "--output_dir", "-o",
        type=str, required=True,
        help="Output directory for corrected images and report",
    )

    # ---- 可选参数 ----
    parser.add_argument("--mask_erosion_radius", type=int, default=2)
    parser.add_argument("--min_shared_tracks", type=int, default=3)
    parser.add_argument("--valid_color_min", type=int, default=1)
    parser.add_argument("--valid_color_max", type=int, default=254)
    parser.add_argument("--huber_k", type=float, default=1.345)
    parser.add_argument("--max_iterations", type=int, default=20)
    parser.add_argument("--tolerance", type=float, default=1e-6)

    return parser.parse_args()


def main():
    args = parse_args()

    # 路径校验
    model_dir = Path(args.model_dir)
    rgb_dir = Path(args.rgb_dir)
    mask_dir = Path(args.mask_dir)
    output_dir = Path(args.output_dir)

    if not model_dir.exists():
        print(f"[ERROR] model_dir does not exist: {model_dir}", file=sys.stderr)
        sys.exit(1)
    if not rgb_dir.exists():
        print(f"[ERROR] rgb_dir does not exist: {rgb_dir}", file=sys.stderr)
        sys.exit(1)
    if not mask_dir.exists():
        print(f"[ERROR] mask_dir does not exist: {mask_dir}", file=sys.stderr)
        sys.exit(1)

    # 参数
    options = CalibrationOptions(
        mask_erosion_radius=args.mask_erosion_radius,
        min_shared_tracks=args.min_shared_tracks,
        valid_color_min=args.valid_color_min,
        valid_color_max=args.valid_color_max,
        huber_k=args.huber_k,
        max_iterations=args.max_iterations,
        tolerance=args.tolerance,
    )

    print("=" * 50)
    print("Photometric Calibration")
    print("=" * 50)
    print(f"  model_dir     : {model_dir}")
    print(f"  rgb_dir       : {rgb_dir}")
    print(f"  mask_dir      : {mask_dir}")
    print(f"  output_dir    : {output_dir}")
    print(f"  min_shared_tracks: {options.min_shared_tracks}")
    print(f"  mask_erosion_radius: {options.mask_erosion_radius}")
    print("-" * 50)

    # 执行标定
    print("Starting calibration...")
    print(f"model_dir={model_dir}")
    outputs = run_calibration(
        model_dir=str(model_dir),
        rgb_dir=str(rgb_dir),
        mask_dir=str(mask_dir),
        output_dir=str(output_dir),
        calibrator=lambda dataset: calibrate(dataset, options=options),
    )

    print("=" * 50)
    print(f"[DONE] Corrected images: {len(outputs.image_paths)}")
    print(f"        Output dir    : {output_dir}")
    print(f"        Report        : {outputs.report_path}")
    print("=" * 50)


if __name__ == "__main__":
    main()