#!/usr/bin/env python3
"""Merge multiple PLY point clouds from lingbot-map jobs into a single PLY.

Each PLY is in its own local reference frame. This script:
  1. Loads all input PLYs as Open3D PointClouds.
  2. Runs voxel downsampling + FPFH feature extraction on each.
  3. Uses RANSAC global registration between adjacent pairs to find initial
     alignment (useful when cameras don't share a scene).
  4. Refines with ICP (point-to-plane).
  5. Merges all aligned clouds and saves the result.

Usage:
    python3 stitch.py out.ply input1.ply input2.ply [input3.ply ...]
    python3 stitch.py out.ply ~/cosma-qc-frames/job_*/reconstruction.ply
    python3 stitch.py out.ply --voxel 0.02 --no-ransac *.ply

The first PLY is the reference frame; all others are aligned to it.

Requires: open3d  (pip install open3d)
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np


def load_ply(path: str):
    import open3d as o3d
    pcd = o3d.io.read_point_cloud(path)
    if len(pcd.points) == 0:
        raise ValueError(f"Empty point cloud: {path}")
    return pcd


def preprocess(pcd, voxel_size: float):
    import open3d as o3d
    pcd_down = pcd.voxel_down_sample(voxel_size)
    pcd_down.estimate_normals(
        o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 2, max_nn=30)
    )
    fpfh = o3d.pipelines.registration.compute_fpfh_feature(
        pcd_down,
        o3d.geometry.KDTreeSearchParamHybrid(radius=voxel_size * 5, max_nn=100),
    )
    return pcd_down, fpfh


def ransac_registration(src_down, dst_down, src_fpfh, dst_fpfh, voxel_size: float):
    import open3d as o3d
    dist_thr = voxel_size * 1.5
    result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
        src_down, dst_down, src_fpfh, dst_fpfh,
        mutual_filter=True,
        max_correspondence_distance=dist_thr,
        estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(False),
        ransac_n=4,
        checkers=[
            o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
            o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(dist_thr),
        ],
        criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(4_000_000, 500),
    )
    return result.transformation


def icp_refine(src, dst, init_transform, voxel_size: float):
    import open3d as o3d
    result = o3d.pipelines.registration.registration_icp(
        src, dst,
        max_correspondence_distance=voxel_size * 0.4,
        init=init_transform,
        estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPlane(),
    )
    return result.transformation


def _load_world_poses(h5_path: str, n_plys: int) -> list[np.ndarray]:
    """Load world-frame transforms from trajectory_world.h5, one per PLY.

    Divides the pose sequence into n_plys equal chunks.
    Returns T_i_to_ref (4x4) for each PLY, where T_0_to_ref = I.
    """
    import h5py
    with h5py.File(h5_path, "r") as f:
        if "poses_world" not in f or "T_4x4" not in f["poses_world"]:
            raise ValueError(f"{h5_path}: missing poses_world/T_4x4")
        T_all = f["poses_world/T_4x4"][:]   # (M, 4, 4)

    M = len(T_all)
    chunk = max(1, M // n_plys)

    avg_T = []
    for i in range(n_plys):
        start = i * chunk
        end   = min(start + chunk, M)
        chunk_T = T_all[start:end]
        avg_t = chunk_T[:, :3, 3].mean(0)
        T_rep = chunk_T[0].copy()
        T_rep[:3, 3] = avg_t
        avg_T.append(T_rep)

    T0_inv = np.linalg.inv(avg_T[0])
    return [T0_inv @ avg_T[i] for i in range(n_plys)]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("output", type=Path, help="Output merged PLY")
    ap.add_argument("inputs", nargs="+", type=Path, help="Input PLY files")
    ap.add_argument("--voxel", type=float, default=0.05,
                    help="Voxel size for downsampling / feature extraction (default 0.05 = 5 cm)")
    ap.add_argument("--no-ransac", action="store_true",
                    help="Skip RANSAC global registration (use if clouds share common poses)")
    ap.add_argument("--icp-only", action="store_true",
                    help="Use identity as init transform and refine with ICP only")
    ap.add_argument("--merge-voxel", type=float, default=0.02,
                    help="Final voxel downsampling on merged cloud (0 = no downsample)")
    ap.add_argument("--poses", type=str, default=None,
                    help="Path to trajectory_world.h5 — use world poses as T_init "
                         "for ICP (replaces RANSAC). Requires h5py.")
    args = ap.parse_args()

    try:
        import open3d as o3d
    except ImportError:
        sys.exit("open3d not found. Install: pip install open3d")

    if len(args.inputs) < 2:
        sys.exit("Need at least 2 input PLY files.")

    print(f"Loading {len(args.inputs)} PLYs...")
    clouds = [load_ply(str(p)) for p in args.inputs]
    for p, c in zip(args.inputs, clouds):
        print(f"  {p.name}: {len(c.points):,} pts")

    # Reference = first cloud
    merged = clouds[0]
    ref_down, ref_fpfh = preprocess(clouds[0], args.voxel)

    # Load pose-guided transforms if available
    world_transforms = None
    if args.poses:
        print(f"Loading world poses from {args.poses}...")
        world_transforms = _load_world_poses(args.poses, len(clouds))
        print(f"Pose-guided init: {len(world_transforms)} transforms loaded")

    for i, src_pcd in enumerate(clouds[1:], start=1):
        print(f"\nAligning {args.inputs[i].name} → {args.inputs[0].name}...")
        src_down, src_fpfh = preprocess(src_pcd, args.voxel)

        if world_transforms is not None:
            init_tf = world_transforms[i]
            print("  Using world pose T_init (no RANSAC)")
        elif args.icp_only or args.no_ransac:
            init_tf = np.eye(4)
        else:
            print("  RANSAC global registration...")
            init_tf = ransac_registration(src_down, ref_down, src_fpfh, ref_fpfh, args.voxel)

        print("  ICP refinement...")
        src_pcd_for_icp = src_pcd.voxel_down_sample(args.voxel * 0.2)
        src_pcd_for_icp.estimate_normals(
            o3d.geometry.KDTreeSearchParamHybrid(radius=args.voxel, max_nn=30)
        )
        ref_for_icp = merged.voxel_down_sample(args.voxel * 0.2)
        ref_for_icp.estimate_normals(
            o3d.geometry.KDTreeSearchParamHybrid(radius=args.voxel, max_nn=30)
        )
        final_tf = icp_refine(src_pcd_for_icp, ref_for_icp, init_tf, args.voxel)

        src_pcd.transform(final_tf)
        merged = merged + src_pcd
        print(f"  Merged total: {len(merged.points):,} pts")

    if args.merge_voxel > 0:
        print(f"\nFinal downsample (voxel={args.merge_voxel})...")
        merged = merged.voxel_down_sample(args.merge_voxel)
        print(f"Final cloud: {len(merged.points):,} pts")

    args.output.parent.mkdir(parents=True, exist_ok=True)
    o3d.io.write_point_cloud(str(args.output), merged)
    print(f"\nSaved → {args.output}")


if __name__ == "__main__":
    main()