from typing import Iterator, List, Tuple

import numpy as np
import trimesh

from ..classes import (DatasetInterface, MeshDataDescription,
                       PcnPartialPointCloudData)


def _copy(array: np.ndarray):
    if array is None:
        return None
    return np.copy(array)


def _read_surface(path: str, num_samples: int) -> np.ndarray:
    data = np.load(path)
    vertices = data["vertices"]
    offset = data["offset"]
    scale = data["scale"]
    rand_indices = np.random.choice(len(vertices),
                                    size=num_samples,
                                    replace=len(vertices) < num_samples)
    vertices = vertices[rand_indices]
    return vertices, offset, scale


class Dataset(DatasetInterface):
    def __init__(self,
                 path_list: List[str],
                 num_coarse_points: int = 1024,
                 num_dense_points: int = 16384):
        self.path_list = path_list
        self.num_coarse_points = num_coarse_points
        self.num_dense_points = num_dense_points
        self.cache = [None] * len(path_list)

    def load_datapoint(self, path_index: int) -> PcnPartialPointCloudData:
        data = self.cache[path_index]
        if data is None:
            path = self.path_list[path_index]
            data = np.load(path)
            gt_dense_points = data["gt"].astype(np.float32)
            partial_points_list = []
            for k in range(8):
                partial_points = data[f"partial_input_{k}"]
                partial_points_list.append(partial_points.astype(np.float32))

            rand_indices = np.random.choice(len(gt_dense_points),
                                            size=self.num_coarse_points,
                                            replace=False)
            gt_coarse_points = gt_dense_points[rand_indices]

            data = PcnPartialPointCloudData(
                gt_coarse_points=gt_coarse_points,
                gt_dense_points=gt_dense_points,
                partial_points_list=partial_points_list)

            self.cache[path_index] = data

        return data

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

    def __iter__(self) -> Iterator[PcnPartialPointCloudData]:
        for path_index in range(len(self.path_list)):
            yield self.load_datapoint(path_index)

    def __getitem__(self, indices) -> PcnPartialPointCloudData:
        if isinstance(indices, int):
            return self.load_datapoint(indices)
        ret = []
        for index in indices:
            ret.append(self.load_datapoint(index))
        return ret

    def shuffle(self) -> Iterator[PcnPartialPointCloudData]:
        indices = np.random.permutation(len(self.path_list))
        for index in indices:
            yield self[int(index)]


class PointCloudAndMeshPairDataset(DatasetInterface):
    def __init__(self,
                 path_pair_list: List[Tuple[str]],
                 num_coarse_points: int = 1024,
                 num_dense_points: int = 16384):
        self.path_pair_list = path_pair_list
        self.num_coarse_points = num_coarse_points
        self.num_dense_points = num_dense_points

    def load_datapoint(self, path_index: int):
        surface_path, obj_path = self.path_pair_list[path_index]

        partial_pc_data = _read_surface(surface_path, self.num_dense_points)

        rand_indices = np.random.choice(len(partial_pc_data.surface_points),
                                        size=self.num_dense_points,
                                        replace=False)
        gt_dense_points = partial_pc_data.surface_points[rand_indices]

        rand_indices = np.random.choice(len(gt_dense_points),
                                        size=self.num_coarse_points,
                                        replace=False)
        gt_coarse_points = gt_dense_points[rand_indices]

        pointcloud_data = GtPartialSamplingData(
            gt_coarse_points=gt_coarse_points,
            gt_dense_points=gt_dense_points,
            surface_points=partial_pc_data.surface_points,
            num_viewpoints=partial_pc_data.num_viewpoints,
            partial_point_indices_list=partial_pc_data.
            partial_point_indices_list,
            scale=partial_pc_data.scale,
            offset=partial_pc_data.offset)

        mesh = trimesh.load_mesh(str(obj_path))
        if isinstance(mesh, trimesh.Scene):
            mesh = trimesh.util.concatenate(
                tuple(
                    trimesh.Trimesh(vertices=g.vertices, faces=g.faces)
                    for g in mesh.geometry.values()))
        vertices = mesh.vertices
        vertex_indices = mesh.faces
        mesh_data = MeshDataDescription(vertices=vertices,
                                        vertex_indices=vertex_indices,
                                        path=obj_path)

        return pointcloud_data, mesh_data

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

    def __iter__(self):
        for path_index in range(len(self.path_pair_list)):
            yield self.load_datapoint(path_index)

    def __getitem__(self, indices):
        if isinstance(indices, int):
            return self.load_datapoint(indices)
        ret = []
        for index in indices:
            ret.append(self.load_datapoint(index))
        return ret

    def shuffle(self, seed=0):
        rst = np.random.RandomState(seed)
        indices = rst.permutation(len(self.path_pair_list))
        for index in indices:
            yield self[int(index)]
