import numpy as np
import torch


def rotation_6d_to_matrix_numpy(d6):
        a1, a2 = d6[:3], d6[3:]
        b1 = a1 / np.linalg.norm(a1)
        b2 = a2 - np.dot(b1, a2) * b1
        b2 = b2 / np.linalg.norm(b2)
        b3 = np.cross(b1, b2)
        return np.stack((b1, b2, b3), axis=-2)

class Pointcloud:
    def __init__(self, points, camera=None, name="unknown", parent=None, from_parent_index=None):
        self.points = points # Nx3 numpy array
        if camera is not None and camera.shape == (9,):
            rotation_matrix = rotation_6d_to_matrix_numpy(camera)
            translation = camera[6:]
            camera = np.eye(4)
            camera[:3, :3] = rotation_matrix
            camera[:3, 3] = translation
        self.camera = camera # 4x4 numpy array
        self.current_transform = np.eye(4)
        self.transform_history = {}
        self.original_points = points.copy()
        self.name = name
        self.parent = parent
        self.from_parent_index = from_parent_index

    # --------- basic ---------
    def __getitem__(self, item):
        return self.points[item]

    def __len__(self):
        return len(self.points)

    def __repr__(self):
        return f"Pointcloud({self.points})"

    def __str__(self):
        return f"Pointcloud with {len(self.points)} points. \n\tCenter: {np.mean(self.points, axis=0)} \n\tTransform history: {list(self.transform_history.keys())}"

    def __eq__(self, other):
        return np.array_equal(self.points, other.points)

    def __ne__(self, other):
        return not np.array_equal(self.points, other.points)

    def __contains__(self, item):
        return item in self.points

    def __iter__(self):
        return iter
    
    def concat(self, other):
        return Pointcloud(np.concatenate([self.points, other.points], axis=0))
    
    # --------- downsample ---------
    def voxel_downsample(self, voxel_size):
        voxel_indices = np.floor(self.points / voxel_size).astype(np.int32)
        _, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
        idx_sort = np.argsort(inverse)
        idx_unique = idx_sort[np.cumsum(counts)-counts]
        downsampled_points = self.points[idx_unique]
        return Pointcloud(
                          downsampled_points, 
                          name=self.name+'(voxel downsampled with voxel_size='+str(voxel_size)+')',
                          parent=self.parent,
                          from_parent_index=idx_unique
                          )
    
    def random_downsample(self, num_points):
        if self.points.shape[0] == 0:
            downsampled_points = self.points
            idx = np.array([])
        else:
            idx = np.random.choice(len(self.points), num_points, replace=True)
            downsampled_points = self.points[idx]
        return Pointcloud(
            downsampled_points, 
            name=self.name+'(random downsampled with num_points='+str(num_points)+')',
            parent=self.parent,
            from_parent_index=idx
        )
    
    def fps_downsample(self, num_points):
        N = self.points.shape[0]
        mask = np.zeros(N, dtype=bool)

        sampled_indices = np.zeros(num_points, dtype=int)
        sampled_indices[0] = np.random.randint(0, N)
        distances = np.linalg.norm(self.points - self.points[sampled_indices[0]], axis=1)
        for i in range(1, num_points):
            farthest_index = np.argmax(distances)
            sampled_indices[i] = farthest_index
            mask[farthest_index] = True

            new_distances = np.linalg.norm(self.points - self.points[farthest_index], axis=1)
            distances = np.minimum(distances, new_distances)
        
        sampled_points = self.points[sampled_indices]
        return Pointcloud(
            sampled_points, 
            name=self.name+'(fps downsampled with num_points='+str(num_points)+')',
            parent=self.parent,
            from_parent_index=sampled_indices
        )
    # --------- transform ---------
    def transform(self, transform_matrix, name=None):
        self.points = np.dot(self.points, transform_matrix[:3, :3].T) + transform_matrix[:3, 3]
        self.current_transform = np.dot(self.current_transform, transform_matrix)
        if name is None:
            name = f"transform_{len(self.transform_history)}"
        self.transform_history[name] = transform_matrix
        return self

    def translate(self, translation, name=None):
        transform_matrix = np.eye(4)
        transform_matrix[:3, 3] = translation
        self.transform(transform_matrix, name)
        return self
    
    def rotate(self, rotation, name=None):
        transform_matrix = np.eye(4)
        if rotation.shape == (3, 3):
            rotation_matrix = rotation
        elif rotation.shape == (6,):
            rotation_matrix = rotation_6d_to_matrix_numpy(rotation)
        transform_matrix[:3, :3] = rotation_matrix
        self.transform(transform_matrix, name)
        return self
    
    def scale(self, scale_factor, name=None):
        transform_matrix = np.eye(4)
        transform_matrix[:3, :3] = scale_factor * np.eye(3)
        self.transform(transform_matrix, name)
        return self
        
    def same_transform(self, other):
        return np.allclose(self.current_transform, other.current_transform)
    
    def print_transform_history(self):
        print(f"Transform history of {self.name}:")
        for name, transform_matrix in self.transform_history.items():
            print(f"\t-{name}:")
            for i in range(4):
                print(f"\t\t{transform_matrix[i]}")
    
    # --------- tensor ---------
    def get_batchlized_points(self, batch_size=1):
        return torch.tensor(self.points).unsqueeze(0).repeat(batch_size, 1, 1)
    
    # --------- visualize ---------
    def visualize(self, point_size=1, color=None):
        import plotly.graph_objects as go
        fig = go.Figure()
        if color is None:
            if self.name is not None:
                hash_value = hash(self.name)
                r = (hash_value & 0xFF) / 255.0
                g = ((hash_value >> 8) & 0xFF) / 255.0
                b = ((hash_value >> 16) & 0xFF) / 255.0
                color = f'rgb({int(r*255)},{int(g*255)},{int(b*255)})'
            else:
                color = "gray"
            

        if self.points is not None:
            fig.add_trace(go.Scatter3d(
                x=self.points[:, 0], y=self.points[:, 1], z=self.points[:, 2],
                mode='markers', marker=dict(size=point_size, color=color, opacity=0.5), name=self.name
            ))
        
        if self.camera is not None:
            origin = self.camera[:3, 3]
            z_axis = self.camera[:3, 2]
            fig.add_trace(go.Cone(
                x=[origin[0]], y=[origin[1]], z=[origin[2]],
                u=[z_axis[0]], v=[z_axis[1]], w=[z_axis[2]],
                colorscale="blues",
                sizemode="absolute", sizeref=0.05, anchor="tail", showscale=False
            ))

        title = self.name
        fig.update_layout(
            title=title,
            scene=dict(
                xaxis_title='X',
                yaxis_title='Y',
                zaxis_title='Z'
            ),
            margin=dict(l=0, r=0, b=0, t=40),
            scene_camera=dict(eye=dict(x=1.25, y=1.25, z=1.25))
        )
        fig.show()

    # --------- save and load ---------
    def save(self, file_path):
        np.save(file_path, self.points)
    
    def savetxt(self, file_path):
        np.savetxt(file_path, self.points)

    def load(self, file_path):
        self.points = np.load(file_path)

    def loadtxt(self, file_path):
        self.points = np.loadtxt(file_path)


class PointcloudGroup:
    def __init__(self, pointclouds: list[Pointcloud] = [], name="unknown"):
        self.pointclouds = pointclouds
        self.name = name
        
    # --------- basic ---------
    def __getitem__(self, item):
        return self.pointclouds[item]

    def __len__(self):
        return len(self.pointclouds)

    def __repr__(self):
        return f"PointcloudGroup({self.pointclouds})"

    def __str__(self):
        return f"PointcloudGroup with {len(self.pointclouds)} pointclouds."

    def __eq__(self, other):
        return np.array_equal(self.pointclouds, other.pointclouds)

    def __ne__(self, other):
        return not np.array_equal(self.pointclouds, other.pointclouds)

    def __contains__(self, item):
        return item in self.pointclouds

    def __iter__(self):
        return iter
    
    def __add__(self, pointcloud: Pointcloud):
        new_group = PointcloudGroup(self.name)
        new_group.pointclouds = self.pointclouds.copy()
        new_group.pointclouds.append(pointcloud)
        return new_group
    
    def __iadd__(self, pointcloud: Pointcloud):
        self.pointclouds.append(pointcloud)
        return self
    
    def add(self, pointcloud: Pointcloud):
        self.pointclouds.append(pointcloud)
    
    def concat(self, other):
        new_group = PointcloudGroup(self.name)
        new_group.pointclouds = self.pointclouds.copy()
        new_group.pointclouds.extend(other.pointclouds)
        return new_group
    
    # --------- merge ---------
    def merge_pointclouds(self, name="unknown"):
        points = np.concatenate([pointcloud.points for pointcloud in self.pointclouds], axis=0)
        return Pointcloud(points, name=name)
    
    # --------- transform ---------
    def transform(self, transform_matrix, name="unknown_transform"):
        for pointcloud in self.pointclouds:
            pointcloud.transform(transform_matrix, name)
    
    def translate(self, translation, name="unknown_translate"):
        transform_matrix = np.eye(4)
        transform_matrix[:3, 3] = translation
        self.transform(transform_matrix, name)
    
    def rotate(self, rotation_matrix, name="unknown_ratate"):
        transform_matrix = np.eye(4)
        transform_matrix[:3, :3] = rotation_matrix
        self.transform(transform_matrix, name)
    
    def scale(self, scale_factor, name="unknown_scale"):
        transform_matrix = np.eye(4)
        transform_matrix[:3, :3] = scale_factor * np.eye(3)
        self.transform(transform_matrix, name)
        
    # --------- visualize ---------
    def visualize(self, point_size=1, color=None):
        import plotly.graph_objects as go
        fig = go.Figure()
        for pointcloud in self.pointclouds:
            if color is None:
                if pointcloud.name is not None:
                    hash_value = hash(pointcloud.name)
                r = (hash_value & 0xFF) / 255.0
                g = ((hash_value >> 8) & 0xFF) / 255.0
                b = ((hash_value >> 16) & 0xFF) / 255.0
                color = f'rgb({int(r*255)},{int(g*255)},{int(b*255)})'
            else:
                color = "gray"
            
            if pointcloud.points is not None:
                fig.add_trace(go.Scatter3d(
                    x=pointcloud.points[:, 0], y=pointcloud.points[:, 1], z=pointcloud.points[:, 2],
                    mode='markers', marker=dict(size=point_size, color=color, opacity=0.5), name=pointcloud.name
                ))
        
            if pointcloud.camera is not None:
                origin = pointcloud.camera[:3, 3]
                z_axis = pointcloud.camera[:3, 2]
                fig.add_trace(go.Cone(
                    x=[origin[0]], y=[origin[1]], z=[origin[2]],
                    u=[z_axis[0]], v=[z_axis[1]], w=[z_axis[2]],
                    colorscale="blues",
                    sizemode="absolute", sizeref=0.05, anchor="tail", showscale=False
                ))

        title = self.name
        fig.update_layout(
            title=title,
            scene=dict(
                xaxis_title='X',
                yaxis_title='Y',
                zaxis_title='Z'
            ),
            margin=dict(l=0, r=0, b=0, t=40),
            scene_camera=dict(eye=dict(x=1.25, y=1.25, z=1.25))
        )
        fig.show()