# MIT LICENSE
# https://github.com/bponsler/kinectToPly/blob/master/kinectToPly.py
import numpy as np


class Ply(object):
    '''The Ply class provides the ability to write a point cloud represented
    by two arrays: an array of points (num points, 3), and an array of colors
    (num points, 3) to a PLY file.
    '''

    def __init__(self, points, colors):
        '''
        * points -- The matrix of points (num points, 3)
        * colors -- The matrix of colors (num points, 3)
        '''
        self.__points = points
        self.__colors = colors

    def write(self, filename):
        '''Write the point cloud data to a PLY file of the given name.
        * filename -- The PLY file
        '''
        # Write the headers
        lines = self.__getLinesForHeader()

        fd = open(filename, "w")
        for line in lines:
            fd.write("%s\n" % line)

        # Write the points
        self.__writePoints(fd, self.__points, self.__colors)

        fd.close()

    def __getLinesForHeader(self):
        '''Get the list of lines for the PLY header.'''
        lines = [
            "ply",
            "format ascii 1.0",
            "comment generated by: kinectToPly",
            "element vertex %s" % len(self.__points),
            "property float x",
            "property float y",
            "property float z",
            "property uchar red",
            "property uchar green",
            "property uchar blue",
            "end_header",
            ]

        return lines

    def __writePoints(self, fd, points, colors):
        '''Write the point cloud points to a file.
        * fd -- The file descriptor
        * points -- The matrix of points (num points, 3)
        * colors -- The matrix of colors (num points, 3)
        '''
        # Stack the two arrays together
        stacked = np.column_stack((points, colors))

        # Write the array to the file
        np.savetxt(
            fd,
            stacked,
            delimiter='\n',
            fmt="%f %f %f %d %d %d")


def write_xyz_rgb_as_ply(point_cloud, rgb_image, path):
    """Write a point cloud with associated rgb image to a ply file
    # Arguments
        point_cloud: xyz point cloud in format [height, width, channels]
        rgb_image: uint8 image in format [height, width, channels]
        path: Where to save the file, ex: '/path/to/folder/file.ply'
    """
    xyz = point_cloud.reshape([point_cloud.size/3, 3])
    rgb = np.squeeze(rgb_image).reshape([point_cloud.size/3, 3])
    ply = Ply(xyz, rgb)
    ply.write(path)