import json
import os
import shutil
import subprocess
import math
from pathlib import Path
from typing import Dict, Any, Optional, List

# -------------------------------------------------
# Paths
# -------------------------------------------------

PROJECT_ROOT = Path(__file__).resolve().parent

# colmap.bat in project root that internally calls bin/colmap.exe
COLMAP_BIN = str(PROJECT_ROOT / "colmap.bat")

CALIBRATION_JSON = PROJECT_ROOT / "calibration_data_matlab.json"
IMAGES_ROOT = PROJECT_ROOT / "images"

WORKSPACE_ROOT = PROJECT_ROOT / "workspace"
DATABASE_PATH = WORKSPACE_ROOT / "database.db"
SPARSE_ROOT = WORKSPACE_ROOT / "sparse"
DENSE_ROOT = WORKSPACE_ROOT / "dense"

# mapping from logical label to folder name under images/
CAMERA_FOLDERS: Dict[str, str] = {
    "center": "center",
    "left": "left",
    "right": "right",
    "extra": "extra",  # moving camera, no priors
}


def run_colmap(args: List[str], cwd: Path) -> None:
    cmd = [COLMAP_BIN] + args
    print("\n[CMD]", " ".join(str(c) for c in cmd))
    subprocess.run(cmd, cwd=str(cwd), check=True)


def load_intrinsics(
    calib: Dict[str, Any],
    camera_labels: Dict[str, str],
) -> Dict[str, Optional[Dict[str, Any]]]:
    """
    Read intrinsics for each label from MATLAB-style JSON.

    For 'extra' or missing entries, return None and let COLMAP choose
    its default camera model.
    """
    intrinsics_all = calib.get("intrinsics", {})
    intrinsics_by_label: Dict[str, Optional[Dict[str, Any]]] = {}

    for label in camera_labels.keys():
        if label == "extra":
            intrinsics_by_label[label] = None
            continue

        if label not in intrinsics_all:
            print(f"[WARN] no intrinsics for camera '{label}' in JSON, "
                  f"letting COLMAP estimate them")
            intrinsics_by_label[label] = None
            continue

        intr = intrinsics_all[label]

        fx = intr["fx"]
        fy = intr["fy"]
        cx = intr["cx"]
        cy = intr["cy"]
        k1 = intr["k1"]
        k2 = intr["k2"]
        p1 = intr["p1"]
        p2 = intr["p2"]
        h = intr["h"]
        w = intr["w"]

        # OPENCV model in COLMAP uses 8 params:
        # [fx, fy, cx, cy, k1, k2, p1, p2]
        params = [fx, fy, cx, cy, k1, k2, p1, p2]

        intrinsics_by_label[label] = {
            "model": "OPENCV",
            "width": w,
            "height": h,
            "params": params,
        }

    return intrinsics_by_label


def build_image_list_file(
    images_root: Path,
    folder_name: str,
    list_file: Path,
) -> None:
    """
    Build a text file that lists images in images_root/folder_name
    as relative paths like:

        center/0.png
        center/1.png
        ...

    Sorted numerically if filenames are 0.png, 1.png, 10.png, etc.
    """
    folder_path = images_root / folder_name

    if not folder_path.exists():
        print(f"[WARN] folder {folder_path} does not exist")
        list_file.write_text("")
        return

    exts = {".png"}

    all_files = [
        fname for fname in os.listdir(folder_path)
        if Path(fname).suffix.lower() in exts
    ]

    def numeric_key(name: str):
        stem = Path(name).stem
        try:
            return int(stem)
        except ValueError:
            return stem

    all_files_sorted = sorted(all_files, key=numeric_key)
    rel_paths = [f"{folder_name}/{fname}" for fname in all_files_sorted]

    if not rel_paths:
        print(f"[WARN] no images found in {folder_path}")
        list_file.write_text("")
        return

    with list_file.open("w") as f:
        for rp in rel_paths:
            f.write(rp + "\n")


def feature_extraction_for_camera(
    images_root: Path,
    camera_label: str,
    folder_name: str,
    intr: Optional[Dict[str, Any]],
    database_path: Path,
    workspace_root: Path,
) -> None:
    """
    Run COLMAP feature_extractor for one folder.

    - If intr is provided: OPENCV intrinsics from MATLAB (center/left/right).
    - If intr is None: no camera_model is set, COLMAP uses its default.
    """
    list_file = workspace_root / f"image_list_{folder_name}.txt"
    build_image_list_file(images_root, folder_name, list_file)

    if list_file.stat().st_size == 0:
        print(f"[INFO] no images to process for camera '{camera_label}'")
        return

    print(f"\n[INFO] extracting features for '{camera_label}' "
          f"from folder '{folder_name}'")

    args = [
        "feature_extractor",
        "--database_path", str(database_path),
        "--image_path", str(images_root),
        "--image_list_path", str(list_file),
        "--ImageReader.single_camera_per_folder", "1",
    ]

    if intr is not None:
        params_str = ",".join(str(p) for p in intr["params"])
        args += [
            "--ImageReader.camera_model", intr["model"],  # OPENCV
            "--ImageReader.camera_params", params_str,
        ]
    # else: no camera_model / camera_params → COLMAP chooses default

    run_colmap(args, cwd=PROJECT_ROOT)


if __name__ == "__main__":
    # clean workspace like the .bat script
    if WORKSPACE_ROOT.exists():
        print(f"[INFO] deleting existing workspace: {WORKSPACE_ROOT}")
        shutil.rmtree(WORKSPACE_ROOT)

    WORKSPACE_ROOT.mkdir(parents=True, exist_ok=True)
    SPARSE_ROOT.mkdir(parents=True, exist_ok=True)
    DENSE_ROOT.mkdir(parents=True, exist_ok=True)

    # load MATLAB calibration (intrinsics only; extrinsics will be ignored)
    with CALIBRATION_JSON.open("r") as f:
        calib = json.load(f)

    intrinsics_by_label = load_intrinsics(calib, CAMERA_FOLDERS)

    # 1) feature extraction per camera (center, left, right, extra)
    for label, folder in CAMERA_FOLDERS.items():
        intr = intrinsics_by_label.get(label)
        feature_extraction_for_camera(
            images_root=IMAGES_ROOT,
            camera_label=label,
            folder_name=folder,
            intr=intr,
            database_path=DATABASE_PATH,
            workspace_root=WORKSPACE_ROOT,
        )

    # 2) matching (no rig_configurator, so no extrinsic priors)
    print("\n[INFO] running exhaustive_matcher")
    run_colmap(
        [
            "exhaustive_matcher",
            "--database_path", str(DATABASE_PATH),
        ],
        cwd=PROJECT_ROOT,
    )

    # 3) sparse reconstruction (extrinsics fully estimated by COLMAP)
    print("\n[INFO] running mapper")
    run_colmap(
        [
            "mapper",
            "--database_path", str(DATABASE_PATH),
            "--image_path", str(IMAGES_ROOT),
            "--output_path", str(SPARSE_ROOT),
        ],
        cwd=PROJECT_ROOT,
    )

    # pick sparse model folder (usually sparse/0)
    model_bin_path = SPARSE_ROOT / "0"
    if not model_bin_path.exists():
        subdirs = [d for d in SPARSE_ROOT.iterdir() if d.is_dir()]
        if not subdirs:
            raise RuntimeError("no sparse model found in 'sparse' folder")
        model_bin_path = subdirs[0]

    print(f"\n[INFO] using sparse model at: {model_bin_path}")

    # 4) undistort images into a COLMAP MVS workspace
    print("\n[INFO] running image_undistorter")
    run_colmap(
        [
            "image_undistorter",
            "--image_path", str(IMAGES_ROOT),
            "--input_path", str(model_bin_path),
            "--output_path", str(DENSE_ROOT),
            "--output_type", "COLMAP",
        ],
        cwd=PROJECT_ROOT,
    )

    # 5) dense stereo (PatchMatch)
    print("\n[INFO] running patch_match_stereo")
    run_colmap(
        [
            "patch_match_stereo",
            "--workspace_path", str(DENSE_ROOT),
            "--workspace_format", "COLMAP",
            "--PatchMatchStereo.geom_consistency", "true",
        ],
        cwd=PROJECT_ROOT,
    )

    # 6) fuse depth maps into a point cloud
    fused_path = DENSE_ROOT / "fused.ply"
    print("\n[INFO] running stereo_fusion")
    run_colmap(
        [
            "stereo_fusion",
            "--workspace_path", str(DENSE_ROOT),
            "--workspace_format", "COLMAP",
            "--input_type", "geometric",
            "--output_path", str(fused_path),
        ],
        cwd=PROJECT_ROOT,
    )

    # 7) mesh the fused cloud
    print("\n[INFO] running poisson_mesher")
    run_colmap(
        [
            "poisson_mesher",
            "--input_path", str(fused_path),
            "--output_path", str(DENSE_ROOT / "meshed-poisson.ply"),
        ],
        cwd=PROJECT_ROOT,
    )

    # 8) convert sparse model to TXT for inspection (like the .bat)
    print("\n[INFO] converting sparse model to TXT")
    run_colmap(
        [
            "model_converter",
            "--input_path", str(model_bin_path),
            "--output_path", str(model_bin_path),
            "--output_type", "TXT",
        ],
        cwd=PROJECT_ROOT,
    )

    print("\nDone.")
