"""
Morfeus lab
The University of Texas
MD Anderson Cancer Center
Author - Caleb O'Connor
Email - csoconnor@mdanderson.org

Description:

Structure:

"""
import copy

import vtk
import numpy as np
import pyvista as pv
import SimpleITK as sitk

from scipy.spatial.transform import Rotation

from open3d.geometry import PointCloud
from open3d.utility import Vector3dVector
from open3d.pipelines.registration import (registration_icp, ICPConvergenceCriteria,
                                           TransformationEstimationPointToPlane, TransformationEstimationPointToPoint)


class ICP(object):
    def __init__(self, source, target, matrix=None):
        self.source = source
        self.target = target

        self.matrix = matrix

        self.icp = None

    def compute_com(self):
        translation = np.asarray(self.mov.center - self.ref.center)

        self.matrix = np.identity(4)
        self.matrix[:3, 3] = translation

    def compute_vtk(self, distance=1e-5, iterations=1000, landmarks=None, com_matching=True, inverse=False):
        if landmarks is None:
            landmarks = int(np.round(len(self.target.points) / 10))

        self.icp = vtk.vtkIterativeClosestPointTransform()
        self.icp.SetSource(self.source)
        self.icp.SetTarget(self.target)
        self.icp.GetLandmarkTransform().SetModeToRigidBody()
        self.icp.SetCheckMeanDistance(1)
        self.icp.SetMeanDistanceModeToRMS()
        if landmarks:
            self.icp.SetMaximumNumberOfLandmarks(landmarks)
        self.icp.SetMaximumMeanDistance(distance)
        self.icp.SetMaximumNumberOfIterations(iterations)
        self.icp.SetStartByMatchingCentroids(com_matching)
        self.icp.Modified()
        self.icp.Update()

        if inverse:
            self.matrix = np.linalg.inv(pv.array_from_vtkmatrix(self.icp.GetMatrix()))
        else:
            self.matrix = pv.array_from_vtkmatrix(self.icp.GetMatrix())

    def compute_o3d(self, distance=10, iterations=1000, rmse=1e-7, fitness=1e-7, method='point', com_matching=True,
                    inverse=False):
        ref_pcd = PointCloud()
        ref_pcd.points = Vector3dVector(np.asarray(self.source.points))

        mov_pcd = PointCloud()
        mov_pcd.points = Vector3dVector(np.asarray(self.target.points))
        mov_pcd.normals = Vector3dVector(np.asarray(self.target.point_normals))

        initial_transform = np.identity(4)
        if com_matching:
            translation = mov_pcd.get_center() - ref_pcd.get_center()
            initial_transform[:3, 3] = translation

        if method == 'point':
            self.icp = registration_icp(ref_pcd, mov_pcd, distance, initial_transform,
                                        TransformationEstimationPointToPoint(),
                                        ICPConvergenceCriteria(max_iteration=iterations,
                                                               relative_rmse=rmse,
                                                               relative_fitness=fitness))
        else:
            self.icp = registration_icp(ref_pcd, mov_pcd, distance, initial_transform,
                                        TransformationEstimationPointToPlane())

        if inverse:
            self.matrix = np.linalg.inv(self.icp.transformation)
        else:
            self.matrix = self.icp.transformation

    def get_matrix(self):
        return self.matrix

    def get_correspondence_set(self):
        if hasattr(self.icp, 'correspondence_set'):
            return np.asarray(self.icp.correspondence_set)

        else:
            return None
