# -*- coding: utf-8 -*-
# Euclidean Cluster Extraction
# http://pointclouds.org/documentation/tutorials/cluster_extraction.php#cluster-extraction
import numpy as np
import pptk
import pcl
import sys
import os
import glob

def proc_cloud(cloud_path):

    try:
        cloud = pcl.load(cloud_path)
        cloud_filtered = cloud
    except:
        print('Error loading pcd file, quitting.')
        exit()

    # Creating the KdTree object for the search method of the extraction
    # pcl::search::KdTree<pcl::PointXYZ>::Ptr tree (new pcl::search::KdTree<pcl::PointXYZ>);
    # tree->setInputCloud (cloud_filtered);
    #tree = cloud_filtered.make_octree(128)
    tree = cloud_filtered.make_kdtree()


    # std::vector<pcl::PointIndices> cluster_indices;
    # pcl::EuclideanClusterExtraction<pcl::PointXYZ> ec;
    # ec.setClusterTolerance (0.02); // 2cm
    # ec.setMinClusterSize (100);
    # ec.setMaxClusterSize (25000);
    # ec.setSearchMethod (tree);
    # ec.setInputCloud (cloud_filtered);
    # ec.extract (cluster_indices);
    ec = cloud_filtered.make_EuclideanClusterExtraction()
    ec.set_ClusterTolerance(.5)
    ec.set_MinClusterSize(20)
    ec.set_MaxClusterSize(250000)
    ec.set_SearchMethod(tree)
    cluster_indices = ec.Extract()

    print('cluster_indices : ' + str(cluster_indices.count) + " count.")
    # print('cluster_indices : ' + str(cluster_indices.indices.max_size) + " count.")

    #   int j = 0;
    #   for (std::vector<pcl::PointIndices>::const_iterator it = cluster_indices.begin (); it != cluster_indices.end (); ++it)
    #   {
    #     pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_cluster (new pcl::PointCloud<pcl::PointXYZ>);
    #     for (std::vector<int>::const_iterator pit = it->indices.begin (); pit != it->indices.end (); ++pit)
    #       cloud_cluster->points.push_back (cloud_filtered->points[*pit]); //*
    #     cloud_cluster->width = cloud_cluster->points.size ();
    #     cloud_cluster->height = 1;
    #     cloud_cluster->is_dense = true;
    #
    #     std::cout << "PointCloud representing the Cluster: " << cloud_cluster->points.size () << " data points." << std::endl;
    #     std::stringstream ss;
    #     ss << "cloud_cluster_" << j << ".pcd";
    #     writer.write<pcl::PointXYZ> (ss.str (), *cloud_cluster, false); //*
    #     j++;
    #   }
    #

    cloud_cluster = pcl.PointCloud()
    avg_points = []

    f = open(cloud_path+'.csv','w+')

    min_depth = 0
    max_depth = 150


    for j, indices in enumerate(cluster_indices):
        # cloudsize = indices
        #print('cluster = ' + str(j) + ', indices = ' + str(len(indices)))
        # cloudsize = len(indices)
        points = np.zeros((len(indices), 3), dtype=np.float32)
        # points = np.zeros((cloudsize, 3), dtype=np.float32)

        # for indice in range(len(indices)):
        for i, indice in enumerate(indices):
            # print('dataNum = ' + str(i) + ', data point[x y z]: ' + str(cloud_filtered[indice][0]) + ' ' + str(cloud_filtered[indice][1]) + ' ' + str(cloud_filtered[indice][2]))
            #print('PointCloud representing the Cluster: ' + str(cloud_cluster.size) + " data points.")
            points[i][0] = cloud_filtered[indice][0]
            points[i][1] = cloud_filtered[indice][1]
            points[i][2] = cloud_filtered[indice][2]

        C = np.cov(np.transpose(points[:,0:3]))
        E, v = np.linalg.eig(C)
        #print(str(np.mean(points,0)) + ' : ' + str(np.mean(E)))

        cent = np.mean(points,0)
        if cent[2] >= min_depth and cent[2] <= max_depth:

            esd = 2*np.sqrt(1/(np.sqrt(E[1])*np.sqrt(E[2])))

            tmp = [cent[0],cent[1],cent[2],E[0],E[1],E[2],esd]
            avg_points.append(tmp)
            #avg_points[j,0:3] = np.mean(points,0)
            #avg_points[j,3:] = E
            #eig_values[j,:] = E
            #cloud_cluster.from_array(points)
            #ss = "cloud_cluster_" + str(j) + ".pcd";
            #pcl.save(cloud_cluster, ss)
            output = ''
            for val in tmp:
                output = output + str(val) + '\t'
            f.write(output[:-1]+'\n')

    f.close()

    # show the cloud
    pc = np.array(avg_points)
    v = pptk.viewer(pc[:,:3])
    v.set(point_size=0.4)
    print(len(avg_points))

if __name__ == '__main__':

    if len(sys.argv) < 2:
        print('Please supply a data set dir as first argument.')
        exit()

    dirs = ['blank1','x128','x64','x32','x16','x8','x4','x2','base']

    for dir in dirs:
        p = os.path.join(sys.argv[1],dir,'*.pcd')
        clouds = glob.glob(p)
        if len(clouds) > 0:
            print(clouds[0])
            proc_cloud(clouds[0])