import open3d as o3d
import numpy as np
from pclpy import pcl
from tqdm import tqdm
import os


def o3d_meshes(dataset_name: str, path_to_data_folder='../../data'):
    """
    Read in mesh (.ply, .stl, .off) files. The function assumes that each class of objects is in
    a separate folder and highly spefified for our needs, which is why a list of all objects of
    each data set are provided as a hard-coded array of strings.
    You can download the data sets referenced in [#1]_ and [#2]_ and use the path to them.

    References
    ----------
    .. [#1] http://modelnet.cs.princeton.edu/ Z. Wu, S. Song, A. Khosla, F. Yu, L. Zhang, X. Tang and J. Xiao. 3D ShapeNets: A Deep Representation for Volumetric Shapes. Proceedings of 28th IEEE Conference on Computer Vision and Pattern Recognition (CVPR2015)
    .. [#2] http://www.cim.mcgill.ca/~shape/benchMark/ K. Siddiqi, J. Zhang, D. Macrini, A. Shokoufandeh, S. Bouix & S. Dickinson. Retrieving Articulated 3D Models Using Medial Surfaces. Machine Vision and Applications, 19(4), 261--274, 2008.

    Parameters
    ----------
    dataset_name: Should correspond to the name of the folder with the data set
    path_to_data_folder

    Returns
    -------
    o3d_meshes : array-like, shape (number of objects)
    The Open3D mesh representations of the objects as a list.

    labels :  array-like, shape (number of objects)
    The labels of the objects as a list of integers starting from 0.

    """

    # http://modelnet.cs.princeton.edu/
    if dataset_name == 'modelnet10':
        objects = ['bathtub', 'bed', 'chair', 'desk', 'dresser', 'monitor', 'night_stand', 'sofa', 'table', 'toilet']

    elif dataset_name == 'tali_15':
        objects = ['Manching', 'Milet']

    elif dataset_name == 'mixed_bones':
        objects = ['capra', 'ovis_aries']

    # http://www.cim.mcgill.ca/~shape/benchMark/
    elif dataset_name == 'mc_gill':
        objects = ['airplanes_ply', 'dinosaurs_ply', 'fishes_ply']

    else:
        raise ValueError('Unknown dataset')

    o3d_meshes = []
    labels = []
    print('Read in %d classes of mesh files...' % len(objects))
    for i, obj in enumerate(tqdm(objects)):
        if dataset_name == 'modelnet10':

            objects_o3d = [o3d.io.read_triangle_mesh(file) for file in
                           list_files(path_to_data_folder + '/' + dataset_name + '/' + obj + '/test')]
        else:
            objects_o3d = [o3d.io.read_triangle_mesh(file) for file in
                           list_files(path_to_data_folder + '/' + dataset_name + '/' + obj)]
        # print('class ', i, ': ', len(objects_o3d), ' objects')
        o3d_meshes.extend(objects_o3d)
        labels.extend([i] * len(objects_o3d))

    return o3d_meshes, labels


def save_o3d_arrays(o3d_meshes, labels, data_set_name: str):
    os.makedirs('../generated/' + data_set_name, exist_ok=True)

    np.save('../generated/' + data_set_name + '/' + data_set_name + '_labels', labels)
    np.save('../generated/' + data_set_name + '/' + data_set_name + '_vertices',
            [np.asarray(obj.vertices_list, dtype=object) for obj in o3d_meshes])


def save_meshes_as_pcd_pointclouds(mesh_list, data_set_name: str):
    path = '../generated/' + data_set_name + '/pcd_pointclouds/'
    os.makedirs(path, exist_ok=True)

    for i, mesh in enumerate(mesh_list):
        pcd = o3d.geometry.PointCloud()
        pcd.points = mesh.vertices
        pcd.normals = mesh.vertex_normals

        o3d.io.write_point_cloud(path + data_set_name + '_' + str(i) + '.pcd', pcd)


def save_meshes_as_poisson_sampled_pointclouds(mesh_list, data_set_name: str, number_of_points, init_factor):
    path = '../../generated/' + data_set_name + '/pcd_pointclouds_poisson/' + str(number_of_points) + '/'
    os.makedirs(path, exist_ok=True)

    print('Save %d meshes as .pcd point clouds...' % len(mesh_list))
    for i, mesh in enumerate(tqdm(mesh_list)):
        file_path = path + data_set_name + '_' + str(i).zfill(2) + '.pcd'

        pcd = mesh.sample_points_poisson_disk(number_of_points=number_of_points, init_factor=init_factor)
        o3d.io.write_point_cloud(file_path, pcd)


def save_meshes_as_ply_files(mesh_list, data_set_name: str):
    path = '../generated/' + data_set_name + '/ply_files/'
    os.makedirs(path, exist_ok=True)

    for i, mesh in enumerate(mesh_list):
        o3d.io.write_triangle_mesh(path + data_set_name + '_' + str(i) + '.ply', mesh)


def read_pcd_pointclouds(data_set_name: str, max_point_cloud_size=50000, path_to_generated_folder='../generated'):
    pcd_files = list_files(path_to_generated_folder + '/' + data_set_name + '/pcd_pointclouds/')
    print('Start reading %d pcd pointclouds.' % len(pcd_files))

    reader = pcl.io.PCDReader()

    pcl_point_clouds = []
    normals = []
    ignored = []
    for i, file in enumerate(pcd_files):
        pc = pcl.PointCloud.PointXYZ()
        reader.read(file, pc)
        p3d_pc = o3d.io.read_point_cloud(file)

        if len(pc.points) < max_point_cloud_size:
            pcl_point_clouds.append(pc)
            normals.append(p3d_pc.normals)
        else:
            ignored.append(i)

    print('%d files were ignored because they are too large.' % len(ignored))
    return pcl_point_clouds, normals, ignored


def read_pcd_pointclouds_poisson(data_set_name: str, number_of_points: int, path_to_generated_folder='../../generated'):
    pcd_files = list_files(path_to_generated_folder + '/' + data_set_name + '/pcd_pointclouds_poisson/'
                           + str(number_of_points) + '/')
    reader = pcl.io.PCDReader()
    pcl_point_clouds = []
    open3d_point_clouds = []
    print('Start reading %d poisson sampled pcd pointclouds...' % len(pcd_files))
    for i, file in enumerate(tqdm(pcd_files)):
        pc = pcl.PointCloud.PointXYZ()
        reader.read(file, pc)
        pcl_point_clouds.append(pc)

        open3d_point_clouds.append(o3d.io.read_point_cloud(file))
        # print(file)

    return pcl_point_clouds, open3d_point_clouds


def read_ply_files_with_pcl(data_set_name: str, max_point_cloud_size=50000):
    pcd_files = list_files('../generated/' + data_set_name + '/ply_files/')
    print('Start reading %d pcd pointclouds.' % len(pcd_files))

    reader = pcl.io.PLYReader()

    pcl_point_clouds = []
    normals = []
    ignored = []
    for i, file in enumerate(pcd_files):
        pc = pcl.PointCloud.PointXYZ()
        mesh = pcl.PolygonMesh()
        ply = pcl.io.loadPLYFile(file, mesh)
        reader.read(file, pc)
        p3d_pc = o3d.io.read_point_cloud(file)

        if len(pc.points) < max_point_cloud_size:
            pcl_point_clouds.append(pc)
            normals.append(p3d_pc.normals)
        else:
            ignored.append(i)

    print('%d files were ignored because they are too large.' % len(ignored))
    return pcl_point_clouds, normals, ignored


def list_files(directory):
    r = []
    for root, dirs, files in os.walk(directory):
        for name in files:
            r.append(os.path.join(root, name))
    return r
