"""
Run COLMAP SfM (Structure-from-Motion) on each part's images using pycolmap.
Produces sparse reconstruction (cameras, images, points3D) needed for 3DGS.

If COLMAP fails (likely for textureless CAD parts with few views), we report
the failure clearly so we can fall back to an alternative pose estimator.
"""
import os
import sys
import shutil
import pycolmap
import numpy as np
from PIL import Image

PARTS_DIR = "data/parts"
PARTS = sorted(os.listdir(PARTS_DIR))

def run_colmap_for_part(part):
    input_dir = os.path.abspath(os.path.join(PARTS_DIR, part, "input"))
    workspace = os.path.abspath(os.path.join(PARTS_DIR, part, "colmap_ws"))
    db_path = os.path.join(workspace, "database.db")
    image_dir = os.path.join(workspace, "images")
    sparse_dir = os.path.join(workspace, "sparse", "0")

    # fresh workspace
    if os.path.exists(workspace):
        shutil.rmtree(workspace)
    os.makedirs(image_dir, exist_ok=True)
    os.makedirs(os.path.join(workspace, "sparse"), exist_ok=True)

    # copy images
    images = sorted([f for f in os.listdir(input_dir) if f.endswith(".png")])
    for img in images:
        shutil.copy(os.path.join(input_dir, img), os.path.join(image_dir, img))

    print(f"\n{'='*60}")
    print(f"PART: {part}  ({len(images)} images)")
    print(f"{'='*60}")

    # Step 1: feature extraction
    print("[1/4] Feature extraction...")
    try:
        pycolmap.extract_features(database_path=db_path, image_path=image_dir)
    except Exception as e:
        print(f"  FAILED: {e}")
        return False, str(e)

    # Step 2: matching
    print("[2/4] Feature matching (exhaustive)...")
    try:
        pycolmap.match_exhaustive(database_path=db_path)
    except Exception as e:
        print(f"  FAILED: {e}")
        return False, str(e)

    # Step 3: incremental SfM
    print("[3/4] Incremental mapping (SfM)...")
    try:
        os.makedirs(sparse_dir, exist_ok=True)
        reconstruction = pycolmap.incremental_mapping(
            database_path=db_path,
            image_path=image_dir,
            output_path=os.path.join(workspace, "sparse"),
            options=pycolmap.IncrementalPipelineOptions(
                min_num_matches=5,
                min_focal_length_ratio=0.1,
                max_focal_length_ratio=10.0,
            ),
        )
    except Exception as e:
        print(f"  FAILED: {e}")
        return False, str(e)

    if reconstruction is None or len(reconstruction) == 0:
        print("  FAILED: No reconstruction produced (0 models)")
        return False, "No reconstruction produced"

    # Check how many images were registered
    rec = reconstruction[0]
    n_reg = sum(1 for im in rec.images.values() if im.has_pose)
    n_points = len(rec.points3D)
    n_cameras = len(rec.cameras)
    print(f"[4/4] Reconstruction: {n_reg}/{len(images)} images registered, {n_points} 3D points, {n_cameras} cameras")

    if n_reg < 3:
        print(f"  WARNING: only {n_reg} images registered - insufficient for 3DGS")
        return False, f"Only {n_reg} images registered"

    # Export to text + PLY for inspection
    rec.write(sparse_dir)
    pycolmap.export_points3D(
        output_path=os.path.join(workspace, "sparse", "points3D.ply"),
        database_path=db_path,
        image_path=image_dir,
        output_type=pycolmap.OutputType.PLY,
    )
    print(f"  Saved sparse model to {sparse_dir}")
    return True, f"{n_reg}/{len(images)} registered, {n_points} points"


def main():
    results = {}
    for part in PARTS:
        input_dir = os.path.join(PARTS_DIR, part, "input")
        if not os.path.isdir(input_dir):
            continue
        ok, msg = run_colmap_for_part(part)
        results[part] = (ok, msg)

    print(f"\n{'='*60}")
    print("COLMAP SUMMARY")
    print(f"{'='*60}")
    for part, (ok, msg) in results.items():
        status = "OK" if ok else "FAILED"
        print(f"  {part}: {status} - {msg}")

    return results

if __name__ == "__main__":
    main()