"""Modified from SparseConvNet data preparation: https://github.com/facebookres
earch/SparseConvNet/blob/master/examples/ScanNet/prepare_data.py."""

import argparse
import glob
import json
import multiprocessing as mp
import numpy as np
import open3d as o3d
import plyfile
import scannet_util
from segmentator import segment_mesh
import torch

# Map relevant classes to {0,1,...,19}, and ignored classes to -100
remapper = np.ones(150) * (-100)
for i, x in enumerate([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39]):
    remapper[x] = i

parser = argparse.ArgumentParser()
parser.add_argument('--data_split', help='data split (train / val / test)', default='test')
opt = parser.parse_args()

path = './data/scannetv2/'
split = opt.data_split
print('data split: {}'.format(split))
files = sorted(glob.glob( path + split + '/*_vh_clean_2.ply'))
if opt.data_split != 'test':
    files2 = sorted(glob.glob(path + split + '/*_vh_clean_2.labels.ply'))
    files3 = sorted(glob.glob(path + split + '/*_vh_clean_2.0.010000.segs.json'))
    files4 = sorted(glob.glob(path + split + '/*[0-9].aggregation.json'))
    assert len(files) == len(files2)
    assert len(files) == len(files3)
    assert len(files) == len(files4), '{} {}'.format(len(files), len(files4))


def f_test(fn):

    f = plyfile.PlyData().read(fn)
    points = np.array([list(x) for x in f.elements[0]])
    coords = np.ascontiguousarray(points[:, :3] - points[:, :3].mean(0))
    colors = np.ascontiguousarray(points[:, 3:6]) / 127.5 - 1

    mesh = o3d.io.read_triangle_mesh(fn)
    vertices = torch.from_numpy(np.array(mesh.vertices).astype(np.float32))
    faces = torch.from_numpy(np.array(mesh.triangles).astype(np.int64))
    superpoint = segment_mesh(vertices, faces, kthr=0.01, segMinVerts=20)

    torch.save((coords, colors, superpoint), fn[:-15] + '_inst_nostuff.pth')
    print('Saving to ' + fn[:-15] + '_inst_nostuff.pth')


def f(fn):
    fn2 = fn[:-3] + 'labels.ply'
    fn3 = fn[:-15] + '_vh_clean_2.0.010000.segs.json'
    fn4 = fn[:-15] + '.aggregation.json'

    f = plyfile.PlyData().read(fn)
    points = np.array([list(x) for x in f.elements[0]])
    coords = np.ascontiguousarray(points[:, :3] - points[:, :3].mean(0))
    colors = np.ascontiguousarray(points[:, 3:6]) / 127.5 - 1

    f2 = plyfile.PlyData().read(fn2)
    sem_labels = remapper[np.array(f2.elements[0]['label'])]

    with open(fn3) as jsondata:
        d = json.load(jsondata)
        seg = d['segIndices']
    segid_to_pointid = {}
    for i in range(len(seg)):
        if seg[i] not in segid_to_pointid:
            segid_to_pointid[seg[i]] = []
        segid_to_pointid[seg[i]].append(i)

    instance_segids = []
    labels = []
    with open(fn4) as jsondata:
        d = json.load(jsondata)
        for x in d['segGroups']:
            if (scannet_util.g_raw2scannetv2[x['label']] != 'wall'
                    and scannet_util.g_raw2scannetv2[x['label']] != 'floor'):
                instance_segids.append(x['segments'])
                labels.append(x['label'])
                assert x['label'] in scannet_util.g_raw2scannetv2.keys()
    # 
    if fn == './data/scannetv2/val/scene0217_00_vh_clean_2.ply' and instance_segids[0] == instance_segids[int(len(instance_segids) / 2)]:
        instance_segids = instance_segids[:int(len(instance_segids) / 2)]
    check = []
    for i in range(len(instance_segids)):
        check += instance_segids[i]
    if not len(np.unique(check)) == len(check):
        print('problem sence:' + str(fn))    
    assert len(np.unique(check)) == len(check)

    instance_labels = np.ones(sem_labels.shape[0]) * -100
    for i in range(len(instance_segids)):
        segids = instance_segids[i]
        pointids = []
        for segid in segids:
            pointids += segid_to_pointid[segid]
        instance_labels[pointids] = i
        assert len(np.unique(sem_labels[pointids])) == 1

    mesh = o3d.io.read_triangle_mesh(fn)
    vertices = torch.from_numpy(np.array(mesh.vertices).astype(np.float32))
    faces = torch.from_numpy(np.array(mesh.triangles).astype(np.int64))

    superpoint = segment_mesh(vertices, faces, kthr=0.01, segMinVerts=20)
    # len_v = vertices.shape[0]
    # len_f = faces.shape[0]
    # pass
    # superpoint = segmentator.segment_mesh(vertices, faces).numpy()
 
    # torch.save(faces,'/home/data1/smn/SPFormer/data/scannetv2/data_temp/face_tensor.pt')
    # torch.save(vertices,'/home/data1/smn/SPFormer/data/scannetv2/data_temp/vertices_tensor.pt')
    # print('len_v :' + str(len_v))
    # print('len_f :' + str(len_f))
        
    torch.save((coords, colors, superpoint, sem_labels, instance_labels), fn[:-15] + '_inst_nostuff.pth')
    print('Saving to ' + fn[:-15] + '_inst_nostuff.pth')



p = mp.Pool(processes=mp.cpu_count())
if opt.data_split == 'test':
    p.map(f_test, files)
else:
    p.map(f, files)
# files_ = ['./data/scannetv2/val/scene0217_00_vh_clean_2.ply']
# p.map(f,files_)
p.close()
p.join()
