import sys
import os
import shutil
import math
from dataset.shapenet import ShapeNetPartSegDataset
from model.pointnet2_part_seg import PointNet2PartSegmentNet
import torch_geometric.transforms as GT
import torch
import argparse
from distutils.util import strtobool
import numpy as np
import pointcloud as pc
import operator


sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../')  # add project root directory


def eval_sample(net, sample):
    '''
    sample: { 'points': tensor(n, 3), 'labels': tensor(n,) }
    return: (pred_label, gt_label) with labels shape (n,)
    '''
    net.eval()
    with torch.no_grad():
        # points: (n, 3)
        points, gt_label = sample['points'], sample['labels']
        n = points.shape[0]
        f = points.shape[1]

        points = points.view(1, n, f)  # make a batch
        points = points.transpose(1, 2).contiguous()
        points = points.to(device, dtype)

        pred = net(points)  # (batch_size, n, num_classes)
        pred_label = pred.max(2)[1]
        pred_label = pred_label.view(-1).cpu()  # (n,)

        assert pred_label.shape == gt_label.shape
        return (pred_label, gt_label)


def append_normal_angles(data):

    def func(x):
        theta = math.acos(x[2]) / math.pi
        phi = (math.atan2(x[1], x[0]) + math.pi) / (2.0 * math.pi)
        return (theta, phi)

    res = np.array([func(xi) for xi in data[:, 3:6]])

    print(res)

    return np.column_stack((data, res))


def recreate_folder(folder):
    if os.path.exists(folder) and os.path.isdir(folder):
        shutil.rmtree(folder)
    os.mkdir(folder)

sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../')  # add project root directory

parser = argparse.ArgumentParser()
parser.add_argument('--npoints', type=int, default=2048, help='resample points number')
parser.add_argument('--model', type=str, default='./checkpoint/seg_model_custom_0.pth', help='model path')
parser.add_argument('--sample_idx', type=int, default=0, help='select a sample to segment and view result')
parser.add_argument('--headers', type=strtobool, default=True, help='if raw files come with headers')
parser.add_argument('--with_normals', type=strtobool, default=True, help='if training will include normals')
parser.add_argument('--collate_per_segment', type=strtobool, default=True, help='whether to look at pointclouds or sub')
parser.add_argument('--has_variations', type=strtobool, default=False,
                    help='whether a single pointcloud has variations '
                         'named int(id)_pc.(xyz|dat) look at pointclouds or sub')
opt = parser.parse_args()
print(opt)

if __name__ == '__main__':

    # ------------------------------------------------------------------------------------------------------------------
    # Load point cloud, cluster it and store clusters as point cloud cluster files again for later prediction
    # ------------------------------------------------------------------------------------------------------------------

    # Create dataset
    print('Create data set ..')

    dataset_folder = './data/raw/predict/'

    # pointcloud_file = './pointclouds/0_0.xyz'

    # pointcloud_file = 'C:/Projekte/csg-fitter/csg-fitter/models/0/0.xyz'

    param_set = {"test1": {"cells": 2, "eps": 0.15, "min_samples": 100},
                 "test2": {"cells": 2, "eps": 0.2, "min_samples": 100},
                 "test8": {"cells": 2, "eps": 0.1, "min_samples": 100},
                 "test12": {"cells": 2, "eps": 0.1, "min_samples": 100},
                 "test15": {"cells": 2, "eps": 0.2, "min_samples": 100}

                 }

    file = "test15"

    used_params = param_set[file]

    pointcloud_file = 'C:/Projekte/visigrapp2020/data/' + file + '/pc.txt'


    # Load and pre-process point cloud
    pcloud = pc.read_pointcloud(pointcloud_file)
    pcloud = pc.normalize_pointcloud(pcloud, 1)

    #pc_clusters = pc.hierarchical_clustering(pcloud, selected_indices_0=[0, 1, 2, 3, 4, 5],
    #                                         selected_indices_1=[0, 1, 2, 3, 4, 5], eps=0.7, min_samples=5)

    n_cells = used_params["cells"]
    pc_clusters = pc.cluster_cubes(pcloud, [n_cells, n_cells, n_cells])
    print("Pre-Processing: Clustering")
    pc.draw_clusters(pc_clusters)

    recreate_folder(dataset_folder)
    for idx, pcc in enumerate(pc_clusters):
         pcc = pc.farthest_point_sampling(pcc, opt.npoints)
         recreate_folder(dataset_folder + str(idx) + '/')
         pc.write_pointcloud(dataset_folder + str(idx) + '/pc.xyz', pcc)
         #draw_sample_data(pcc, True)

    # ------------------------------------------------------------------------------------------------------------------
    # Load point cloud clusters and model.
    # ------------------------------------------------------------------------------------------------------------------

    # Load dataset
    print('load dataset ..')
    test_transform = GT.Compose([GT.NormalizeScale(), ])

    test_dataset = ShapeNetPartSegDataset(
         mode='predict',
         root_dir='data',
         with_normals=opt.with_normals,
         npoints=opt.npoints,
         refresh=True,
         collate_per_segment=opt.collate_per_segment,
         has_variations=opt.has_variations,
         headers=opt.headers
    )
    num_classes = test_dataset.num_classes()

    # Load model
    print('Construct model ..')
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    dtype = torch.float

    net = PointNet2PartSegmentNet(num_classes, with_normals=opt.with_normals)
    net.load_state_dict(torch.load(opt.model, map_location=device.type))
    net = net.to(device, dtype)
    net.eval()

    # ------------------------------------------------------------------------------------------------------------------
    # Predict per cluster.
    # ------------------------------------------------------------------------------------------------------------------

    labeled_dataset = None
    result_clusters = []

    # Iterate over all the samples and predict
    for idx, sample in enumerate(test_dataset):

         predicted_label, _ = eval_sample(net, sample)
         if opt.with_normals:
             sample_data = np.column_stack((sample["points"].numpy(), predicted_label.numpy()))
         else:
             sample_data = np.column_stack((sample["points"].numpy(), sample["normals"], predicted_label.numpy()))

         result_clusters.append(sample_data)

         if labeled_dataset is None:
             labeled_dataset = sample_data
         else:
             labeled_dataset = np.vstack((labeled_dataset, sample_data))

         print("Prediction done for cluster " + str(idx+1) + "/" + str(len(test_dataset)))

    # ------------------------------------------------------------------------------------------------------------------
    # Remove cluster rows if the amount of points for a particular primitive type is below a threshold.
    # ------------------------------------------------------------------------------------------------------------------

    min_cluster_size = 10
    contamination = 0.01

    filtered_clusters = filter(lambda c : c.shape[0] > min_cluster_size, result_clusters)
    type_filtered_clusters = []
    for c in filtered_clusters:

        prim_types = np.unique(c[:, 6])
        pt_count = {}
        for pt in prim_types:
            pt_count[pt] = len(c[c[:, 6] == pt])

        max_pt = max(pt_count.items(), key=operator.itemgetter(1))[0]
        min_size = pt_count[max_pt] * contamination

        valid_types = []
        for pt in prim_types:
            if pt_count[pt] > min_size:
                valid_types.append(pt)

        filtered_c = c[np.isin(c[:, 6], valid_types)]
        type_filtered_clusters.append(filtered_c)

    result_clusters = type_filtered_clusters

    labeled_dataset = np.vstack(result_clusters)

    np.savetxt('labeled_dataset.txt', labeled_dataset)

    # ------------------------------------------------------------------------------------------------------------------
    # Clustering that results in per-primitive type clusters
    # ------------------------------------------------------------------------------------------------------------------

    #labeled_dataset = np.loadtxt('labeled_dataset.txt')
    pc.draw_sample_data(labeled_dataset)

    # Try to get rid of outliers.
    # labeled_dataset,outliers = pc.split_outliers(labeled_dataset, columns=[0,1,2,3,4,5])
    # pc.draw_sample_data(outliers, False)

    print("Final clustering..")

    labeled_dataset = pc.append_onehotencoded_type(labeled_dataset)

    print("Test row: ", labeled_dataset[:1, :])

    total_clusters = []

    clusters = pc.cluster_dbscan(labeled_dataset, [0,1,2,3,4,5], eps=used_params["eps"], min_samples=used_params["min_samples"])
    print("Pre-clustering done. Clusters: ", len(clusters))
    #pc.draw_clusters(clusters)

    for cluster in clusters:
        #cluster = pc.normalize_pointcloud(cluster)

        print("2nd level clustering ..")

        prim_types_in_cluster = len(np.unique(cluster[:, 6], axis=0))
        if prim_types_in_cluster == 1:
            print("No need for 2nd level clustering since there is only a single primitive type in the cluster.")
            total_clusters.append(cluster)
        else:
            sub_clusters = pc.cluster_dbscan(cluster, [0,1,2,7,8,9,10], eps=used_params["eps"], min_samples=used_params["min_samples"])
            print("Sub clusters: ", len(sub_clusters))
            total_clusters.extend(sub_clusters)

    result_clusters = list(filter(lambda c: c.shape[0] > used_params["min_samples"], total_clusters))

    for cluster in result_clusters:
        print("Cluster: ", cluster.shape[0])

        # pc.draw_sample_data(cluster, False)

    print("Number of clusters: ", len(result_clusters))

    pc.draw_clusters(result_clusters)

    # ------------------------------------------------------------------------------------------------------------------
    # Write clusters to file.
    # ------------------------------------------------------------------------------------------------------------------
    pc.write_clusters("clusters.txt", result_clusters)