#!/usr/bin/env python3
"""
Point Cloud Registration Demo using Open3D and 3DMatch RedKitchen
Interactive Gradio app for pairwise point cloud registration
"""

import json
import os
import copy
import numpy as np
import gradio as gr
from pathlib import Path
import tempfile
from typing import Tuple, Dict, List
import plotly.graph_objects as go

try:
    import open3d as o3d
except ImportError:
    raise ImportError("Please install open3d: pip install open3d")


def get_file_path(file_obj):
    """Convert a Gradio file output or string into a filesystem path."""
    if file_obj is None:
        return None

    if isinstance(file_obj, str):
        return file_obj

    if hasattr(file_obj, "name"):
        return file_obj.name

    if isinstance(file_obj, dict):
        return file_obj.get("path") or file_obj.get("name")

    return str(file_obj)


def point_cloud_to_arrays(pcd, max_points=30000):
    """Convert Open3D point cloud to numpy array, optionally downsampled."""
    points = np.asarray(pcd.points)
    
    if len(points) == 0:
        return points
    
    if len(points) > max_points:
        idx = np.random.choice(len(points), max_points, replace=False)
        points = points[idx]
    
    return points


class RegistrationDemo:
    def __init__(self, examples_dir: str = "examples"):
        self.examples_dir = Path(examples_dir)
        self.metadata = self._load_metadata()
        self.current_source = None
        self.current_target = None
        self.current_source_down = None
        self.current_target_down = None
    
    def _load_metadata(self) -> Dict:
        """Load pair metadata"""
        metadata_file = self.examples_dir / "pair_metadata.json"
        if metadata_file.exists():
            with open(metadata_file) as f:
                return json.load(f)
        return {}
    
    def get_pair_choices(self) -> List[str]:
        """Get list of available demo pairs"""
        if not self.metadata:
            return []
        return sorted(self.metadata.keys())
    
    def load_demo_pair(self, pair_key: str) -> Tuple[str, str]:
        """Load source and target from demo pair"""
        if pair_key not in self.metadata:
            return None, None
        
        pair_info = self.metadata[pair_key]
        source_file = self.examples_dir / pair_info["source_file"]
        target_file = self.examples_dir / pair_info["target_file"]
        
        if source_file.exists() and target_file.exists():
            return str(source_file), str(target_file)
        
        return None, None
    
    def load_point_clouds(self, source_path: str, target_path: str) -> Tuple[bool, str]:
        """Load point clouds from files"""
        try:
            self.current_source = o3d.io.read_point_cloud(source_path)
            self.current_target = o3d.io.read_point_cloud(target_path)
            
            n_source = len(self.current_source.points)
            n_target = len(self.current_target.points)
            
            return True, f"✓ Loaded: {n_source} source points, {n_target} target points"
        except Exception as e:
            return False, f"✗ Error loading point clouds: {e}"
    
    def preprocess_point_clouds(self, voxel_size: float) -> Tuple[bool, str]:
        """Preprocess loaded point clouds"""
        if self.current_source is None or self.current_target is None:
            return False, "✗ No point clouds loaded"
        
        try:
            # Remove non-finite points
            def remove_non_finite(pcd):
                pcd_clean = o3d.geometry.PointCloud()
                mask = ~np.any(~np.isfinite(np.asarray(pcd.points)), axis=1)
                pcd_clean.points = o3d.utility.Vector3dVector(np.asarray(pcd.points)[mask])
                if pcd.has_normals():
                    pcd_clean.normals = o3d.utility.Vector3dVector(np.asarray(pcd.normals)[mask])
                pcd_clean.remove_duplicated_points()
                return pcd_clean
            
            source_clean = remove_non_finite(self.current_source)
            target_clean = remove_non_finite(self.current_target)
            
            # Downsample
            self.current_source_down = source_clean.voxel_down_sample(voxel_size)
            self.current_target_down = target_clean.voxel_down_sample(voxel_size)
            
            n_source = len(self.current_source_down.points)
            n_target = len(self.current_target_down.points)
            
            if n_source < 10 or n_target < 10:
                return False, f"✗ Too few points after downsampling: {n_source}, {n_target}"
            
            return True, f"✓ Preprocessed: {n_source} source, {n_target} target points"
        
        except Exception as e:
            return False, f"✗ Preprocessing error: {e}"
    
    def estimate_normals_and_features(
        self,
        voxel_size: float,
        normal_radius_mult: float,
        fpfh_radius_mult: float,
    ) -> Tuple[bool, str, object, object]:
        """Estimate normals and compute FPFH features"""
        if self.current_source_down is None or self.current_target_down is None:
            return False, "✗ No preprocessed point clouds", None, None
        
        try:
            normal_radius = voxel_size * normal_radius_mult
            fpfh_radius = voxel_size * fpfh_radius_mult
            
            # Estimate normals
            self.current_source_down.estimate_normals(
                o3d.geometry.KDTreeSearchParamRadius(radius=normal_radius)
            )
            self.current_target_down.estimate_normals(
                o3d.geometry.KDTreeSearchParamRadius(radius=normal_radius)
            )
            
            # Compute FPFH
            source_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
                self.current_source_down,
                o3d.geometry.KDTreeSearchParamRadius(radius=fpfh_radius)
            )
            
            target_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
                self.current_target_down,
                o3d.geometry.KDTreeSearchParamRadius(radius=fpfh_radius)
            )
            
            return True, "✓ Features computed", source_fpfh, target_fpfh
        
        except Exception as e:
            return False, f"✗ Feature computation error: {e}", None, None
    
    def ransac_registration(
        self,
        voxel_size: float,
        normal_radius_mult: float = 2.0,
        fpfh_radius_mult: float = 5.0,
        distance_mult: float = 1.5,
        max_iterations: int = 50000,
    ) -> Tuple[bool, str, np.ndarray]:
        """Global registration using RANSAC"""
        if self.current_source_down is None or self.current_target_down is None:
            return False, "✗ No preprocessed point clouds", None
        
        try:
            success, msg, source_fpfh, target_fpfh = self.estimate_normals_and_features(
                voxel_size,
                normal_radius_mult,
                fpfh_radius_mult,
            )
            if not success:
                return False, msg, None
            
            distance_threshold = voxel_size * distance_mult
            
            result = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
                self.current_source_down, self.current_target_down,
                source_fpfh, target_fpfh,
                mutual_filter=False,
                max_correspondence_distance=distance_threshold,
                estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(False),
                ransac_n=3,
                checkers=[
                    o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
                    o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold)
                ],
                criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(
                    int(max_iterations),
                    0.999,
                )
            )
            
            msg = f"✓ RANSAC: fitness={result.fitness:.4f}, RMSE={result.inlier_rmse:.6f}"
            return True, msg, result.transformation
        
        except Exception as e:
            return False, f"✗ RANSAC error: {e}", None
    
    def icp_registration(self, init_transform: np.ndarray, voxel_size: float,
                        distance_mult: float = 0.4, max_iterations: int = 50) -> Tuple[bool, str, np.ndarray]:
        """Local registration using ICP"""
        if self.current_source_down is None or self.current_target_down is None:
            return False, "✗ No preprocessed point clouds", None
        
        try:
            normal_radius = voxel_size * 2.0

            if not self.current_source_down.has_normals():
                self.current_source_down.estimate_normals(
                    o3d.geometry.KDTreeSearchParamRadius(radius=normal_radius)
                )

            if not self.current_target_down.has_normals():
                self.current_target_down.estimate_normals(
                    o3d.geometry.KDTreeSearchParamRadius(radius=normal_radius)
                )
            
            distance_threshold = voxel_size * distance_mult
            
            result = o3d.pipelines.registration.registration_icp(
                self.current_source_down, self.current_target_down,
                max_correspondence_distance=distance_threshold,
                init=init_transform,
                criteria=o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=max_iterations),
                estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPlane()
            )
            
            msg = f"✓ ICP: fitness={result.fitness:.4f}, RMSE={result.inlier_rmse:.6f}"
            return True, msg, result.transformation
        
        except Exception as e:
            return False, f"✗ ICP error: {e}", None
    
    def create_before_visualization(self) -> Tuple[go.Figure, str]:
        """Create before registration Plotly visualization."""
        if self.current_source is None or self.current_target is None:
            return None, "✗ No point clouds loaded"
        
        try:
            src = point_cloud_to_arrays(self.current_source)
            tgt = point_cloud_to_arrays(self.current_target)
            
            fig = go.Figure()
            
            fig.add_trace(go.Scatter3d(
                x=src[:, 0], y=src[:, 1], z=src[:, 2],
                mode="markers",
                marker=dict(size=1.5, color="orange"),
                name="Source"
            ))
            
            fig.add_trace(go.Scatter3d(
                x=tgt[:, 0], y=tgt[:, 1], z=tgt[:, 2],
                mode="markers",
                marker=dict(size=1.5, color="cyan"),
                name="Target"
            ))
            
            fig.update_layout(
                scene=dict(aspectmode="data"),
                margin=dict(l=0, r=0, t=30, b=0),
                height=520,
                legend=dict(x=0, y=1)
            )
            
            return fig, "✓ Before visualization ready"
        
        except Exception as e:
            return None, f"✗ Visualization error: {e}"
    
    def create_after_visualization(self, transform: np.ndarray) -> Tuple[go.Figure, str]:
        """Create after registration Plotly visualization."""
        if self.current_source is None or self.current_target is None:
            return None, "✗ No point clouds loaded"
        
        if transform is None:
            return None, "✗ No transformation available"
        
        try:
            source_transformed = copy.deepcopy(self.current_source)
            source_transformed.transform(transform)
            
            src = point_cloud_to_arrays(source_transformed)
            tgt = point_cloud_to_arrays(self.current_target)
            
            fig = go.Figure()
            
            fig.add_trace(go.Scatter3d(
                x=src[:, 0], y=src[:, 1], z=src[:, 2],
                mode="markers",
                marker=dict(size=1.5, color="lime"),
                name="Aligned Source"
            ))
            
            fig.add_trace(go.Scatter3d(
                x=tgt[:, 0], y=tgt[:, 1], z=tgt[:, 2],
                mode="markers",
                marker=dict(size=1.5, color="cyan"),
                name="Target"
            ))
            
            fig.update_layout(
                scene=dict(aspectmode="data"),
                margin=dict(l=0, r=0, t=30, b=0),
                height=520,
                legend=dict(x=0, y=1)
            )
            
            return fig, "✓ After visualization ready"
        
        except Exception as e:
            return None, f"✗ Visualization error: {e}"
    
    def get_metrics_and_transform(self, fitness: float, rmse: float, 
                                  transform: np.ndarray) -> Tuple[str, str]:
        """Format metrics and transformation matrix for display"""
        metrics_text = f"""
Fitness: {fitness:.6f}
RMSE: {rmse:.6f}
        """.strip()
        
        if transform is not None:
            transform_text = "Transformation Matrix (4x4):\n"
            for row in transform:
                transform_text += " ".join([f"{x:12.6f}" for x in row]) + "\n"
        else:
            transform_text = "No transformation available"
        
        return metrics_text, transform_text
    
    def save_aligned_source(self, transform: np.ndarray) -> Tuple[str, str]:
        """Save transformed source point cloud only."""
        if self.current_source is None:
            return None, "✗ No source point cloud loaded"

        if transform is None:
            return None, "✗ No transformation available"

        try:
            aligned_source = copy.deepcopy(self.current_source)
            aligned_source.transform(transform)

            with tempfile.NamedTemporaryFile(suffix=".ply", delete=False) as f:
                o3d.io.write_point_cloud(f.name, aligned_source)
                return f.name, "✓ Aligned source saved"

        except Exception as e:
            return None, f"✗ Save aligned source error: {e}"
    
    def register(self, pair_selection: str, mode: str, voxel_size: float,
                 normal_radius_mult: float, fpfh_radius_mult: float,
                 ransac_dist_mult: float, ransac_iter: int,
                 icp_dist_mult: float, icp_iter: int,
                 source_upload, target_upload) -> Tuple:
        """Main registration pipeline"""
        
        # Load point clouds
        if pair_selection != "Upload" and pair_selection:
            source_path, target_path = self.load_demo_pair(pair_selection)
            if source_path is None:
                return None, None, "✗ Failed to load demo pair", "", "", ""
        else:
            if source_upload is None or target_upload is None:
                return None, None, "✗ Please upload both source and target", "", "", ""
            source_path = get_file_path(source_upload)
            target_path = get_file_path(target_upload)
        
        # Load
        success, msg = self.load_point_clouds(source_path, target_path)
        if not success:
            return None, None, msg, "", "", ""
        
        # Preprocess
        success, msg = self.preprocess_point_clouds(voxel_size)
        if not success:
            return None, None, msg, "", "", ""
        
        # Registration
        final_transform = np.eye(4)
        ransac_msg = ""
        icp_msg = ""
        
        try:
            if mode in ["RANSAC + ICP", "RANSAC only"]:
                success, ransac_msg, transform = self.ransac_registration(
                    voxel_size,
                    normal_radius_mult,
                    fpfh_radius_mult,
                    ransac_dist_mult,
                    int(ransac_iter),
                )
                if not success:
                    return None, None, ransac_msg, "", "", ""
                final_transform = transform
            
            if mode in ["RANSAC + ICP", "ICP only"]:
                if mode == "ICP only":
                    init = np.eye(4)
                else:
                    init = final_transform
                
                success, icp_msg, transform = self.icp_registration(
                    init, voxel_size, icp_dist_mult, icp_iter
                )
                if not success:
                    return None, None, icp_msg, "", "", ""
                final_transform = transform
            
            # Visualizations
            before_file, before_msg = self.create_before_visualization()
            after_file, after_msg = self.create_after_visualization(final_transform)
            
            if before_file is None or after_file is None:
                return None, None, before_msg or after_msg, "", "", ""
            
            # Compute final fitness/RMSE
            source_aligned = copy.deepcopy(self.current_source_down)
            source_aligned.transform(final_transform)
            
            distances = source_aligned.compute_point_cloud_distance(self.current_target_down)
            distances = np.asarray(distances)
            
            fitness = np.sum(distances < voxel_size * 1.5) / len(distances)
            rmse = np.sqrt(np.mean(distances ** 2))
            
            metrics_text, transform_text = self.get_metrics_and_transform(fitness, rmse, final_transform)
            
            status_msg = f"✓ Registration complete!\n{ransac_msg}\n{icp_msg}".strip()
            
            aligned_file, aligned_msg = self.save_aligned_source(final_transform)

            if aligned_file is None:
                return before_file, after_file, status_msg + "\n" + aligned_msg, metrics_text, transform_text, None

            return before_file, after_file, status_msg, metrics_text, transform_text, aligned_file
        
        except Exception as e:
            return None, None, f"✗ Registration error: {e}", "", "", ""


def create_gradio_interface():
    """Create Gradio interface"""
    demo = RegistrationDemo(examples_dir="examples")
    
    pair_choices = demo.get_pair_choices()
    if not pair_choices:
        pair_choices = ["(No demo pairs loaded)"]
    
    with gr.Blocks(title="Point Cloud Registration Demo") as app:
        gr.Markdown("""
        # Point Cloud Registration Demo
        
        An interactive demo for pairwise 3D point cloud registration using **Open3D**.
        
        This demo includes pre-selected point cloud pairs from the **3DMatch Geometric Registration Benchmark**, specifically the **RedKitchen** scene.  
        These point clouds are 3D fragments reconstructed from depth frames using TSDF fusion.
        
        Below, you can choose a provided demo pair or upload your own source and target point clouds for registration.
        
        The registration pipeline uses:
        - **RANSAC + FPFH** for global registration to estimate the initial transformation.
        - **ICP** for local refinement to improve the alignment.
        - Supported modes include **RANSAC + ICP**, **RANSAC only**, and **ICP only**.
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Input")
                
                pair_selection = gr.Dropdown(
                    choices=["Upload"] + pair_choices,
                    value="Upload" if not pair_choices else pair_choices[0],
                    label="Dataset Pair"
                )
                
                source_upload = gr.File(
                    label="Source Point Cloud (PLY)",
                    file_types=[".ply"],
                    type="filepath",
                )

                target_upload = gr.File(
                    label="Target Point Cloud (PLY)",
                    file_types=[".ply"],
                    type="filepath",
                )
            
            with gr.Column(scale=1):
                gr.Markdown("### Registration Settings")
                
                mode = gr.Radio(
                    choices=["RANSAC + ICP", "RANSAC only", "ICP only"],
                    value="RANSAC + ICP",
                    label="Registration Mode"
                )
                
                voxel_size = gr.Slider(0.01, 0.2, value=0.05, step=0.01, label="Voxel Size")
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### RANSAC Parameters")
                ransac_dist_mult = gr.Slider(0.5, 3.0, value=1.5, step=0.1, label="Distance Multiplier")
                ransac_iter = gr.Slider(1000, 100000, value=50000, step=5000, label="Max Iterations")
            
            with gr.Column(scale=1):
                gr.Markdown("### ICP Parameters")
                icp_dist_mult = gr.Slider(0.1, 1.0, value=0.4, step=0.1, label="Distance Multiplier")
                icp_iter = gr.Slider(10, 100, value=50, step=5, label="Max Iterations")
            
            with gr.Column(scale=1):
                gr.Markdown("### Feature Parameters")
                normal_radius_mult = gr.Slider(1.0, 5.0, value=2.0, step=0.5, label="Normal Radius Mult")
                fpfh_radius_mult = gr.Slider(2.0, 10.0, value=5.0, step=0.5, label="FPFH Radius Mult")
        
        run_button = gr.Button("Run Registration", variant="primary", size="lg")
        
        with gr.Row():
            status = gr.Textbox(label="Status", lines=2, interactive=False)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### Before Registration")
                before_viewer = gr.Plot(label="Before Registration")
            
            with gr.Column(scale=1):
                gr.Markdown("### After Registration")
                after_viewer = gr.Plot(label="After Registration")

                download_aligned = gr.File(
                    label="Download Aligned Source",
                    type="filepath",
                )
        
        with gr.Row():
            with gr.Column(scale=1):
                metrics = gr.Textbox(label="Metrics", lines=3, interactive=False)
            
            with gr.Column(scale=1):
                transform_matrix = gr.Textbox(label="Transformation Matrix", lines=5, interactive=False)
        
        # Event handlers
        run_button.click(
            fn=demo.register,
            inputs=[
                pair_selection, mode, voxel_size,
                normal_radius_mult, fpfh_radius_mult,
                ransac_dist_mult, ransac_iter,
                icp_dist_mult, icp_iter,
                source_upload, target_upload
            ],
            outputs=[before_viewer, after_viewer, status, metrics, transform_matrix, download_aligned]
        )
    
    return app


if __name__ == "__main__":
    app = create_gradio_interface()
    app.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        ssr_mode=False,
    )