from __future__ import annotations

import copy
from dataclasses import dataclass

import numpy as np
import open3d as o3d
from scipy.spatial.transform import Rotation


@dataclass
class RegistrationResult:
    transformation: np.ndarray
    fitness: float
    rmse: float
    target: o3d.geometry.PointCloud
    recovered: o3d.geometry.PointCloud


class PointCloudRegistration:
    """RANSAC + ICP point-cloud registration."""

    def rotation_matrix(self, roll: float = 10.0, pitch: float = 10.0, yaw: float = 10.0) -> np.ndarray:
        return Rotation.from_euler("xyz", [roll, pitch, yaw], degrees=True).as_matrix()

    def translation_vector(self, trans_range: float = 0.0, seed: int | None = None) -> np.ndarray:
        rng = np.random.default_rng(seed)
        return rng.uniform(-trans_range, trans_range, 3)

    def transform_point_cloud(
        self,
        pointcloud: o3d.geometry.PointCloud,
        rotation: np.ndarray,
        translation: np.ndarray,
    ) -> o3d.geometry.PointCloud:
        transformed = copy.deepcopy(pointcloud)
        points = np.asarray(transformed.points)
        transformed.points = o3d.utility.Vector3dVector((rotation @ points.T).T + translation)
        return transformed

    def register(
        self,
        source: o3d.geometry.PointCloud,
        target: o3d.geometry.PointCloud,
        voxel_size: float = 10.0,
    ) -> tuple[np.ndarray, float, float]:
        source_down = source.voxel_down_sample(voxel_size)
        target_down = target.voxel_down_sample(voxel_size)

        normal_radius = voxel_size * 2
        source_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=normal_radius, max_nn=30))
        target_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=normal_radius, max_nn=30))

        feature_radius = voxel_size * 5
        source_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
            source_down,
            o3d.geometry.KDTreeSearchParamHybrid(radius=feature_radius, max_nn=100),
        )
        target_fpfh = o3d.pipelines.registration.compute_fpfh_feature(
            target_down,
            o3d.geometry.KDTreeSearchParamHybrid(radius=feature_radius, max_nn=100),
        )

        distance_threshold = voxel_size * 1.5
        result_ransac = o3d.pipelines.registration.registration_ransac_based_on_feature_matching(
            source_down,
            target_down,
            source_fpfh,
            target_fpfh,
            mutual_filter=True,
            max_correspondence_distance=distance_threshold,
            estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(),
            ransac_n=3,
            checkers=[
                o3d.pipelines.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
                o3d.pipelines.registration.CorrespondenceCheckerBasedOnDistance(distance_threshold),
            ],
            criteria=o3d.pipelines.registration.RANSACConvergenceCriteria(100000, 0.999),
        )

        source.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
        target.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))

        result_icp = o3d.pipelines.registration.registration_icp(
            source,
            target,
            distance_threshold / 2,
            result_ransac.transformation,
            o3d.pipelines.registration.TransformationEstimationPointToPlane(),
            o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=200),
        )
        return result_icp.transformation, float(result_icp.fitness), float(result_icp.inlier_rmse)

    def run_simulated_registration(
        self,
        source: o3d.geometry.PointCloud,
        voxel_size: float = 10.0,
        roll: float = 10.0,
        pitch: float = 10.0,
        yaw: float = 10.0,
        trans_range: float = 0.0,
        random_seed: int | None = None,
    ) -> RegistrationResult:
        rotation = self.rotation_matrix(roll, pitch, yaw)
        translation = self.translation_vector(trans_range, random_seed)
        target = self.transform_point_cloud(source, rotation, translation)
        transformation, fitness, rmse = self.register(source, target, voxel_size)
        recovered = copy.deepcopy(target)
        recovered.transform(np.linalg.inv(transformation))
        return RegistrationResult(transformation, fitness, rmse, target, recovered)