import shutil
import os
import sys
import random
import numpy as np

import torch
from torch.utils.data import Dataset

import potpourri3d as pp3d

sys.path.append(os.path.join(os.path.dirname(__file__), "../../src/"))  
import fannet
from fannet.utils import toNP

class Shrec11MeshDataset_Original(Dataset):

    # NOTE: Original data from the challenge, not simplified models
  
    # The original SHREC11 models were previously distributed via NIST [here](https://www.nist.gov/itl/iad/shrec-2011-datasets), but that page seems to have been lost to the sands of time. We provide a zip of the old dataset page here: https://drive.google.com/uc?export=download&id=1O_P03aAxhjCOKQH2n71j013-EfSmEp5e. The relevant files are in `SHREC11_test_database_new.zip`, which is password protected with the password `SHREC11@NIST`. 

    # Unzip it like
    # unzip -P SHREC11@NIST SHREC11_test_database_new.zip -d [DATA_ROOT]/raw
    
    def __init__(self, root_dir, split_size, k_eig, exclude_dict=None, op_cache_dir=None, spoke_length=[0.00,0.01,0.02,0.03]):
        
        self.root_dir = root_dir
        self.n_class = 30 
        self.split_size = split_size # pass None to take all entries (except those in exclude_dict)
        self.k_eig = k_eig
        self.op_cache_dir = op_cache_dir
        self.spoke_length = spoke_length

        self.class_names = []
        self.entries = {}

        # store in memory
        self.verts_list = []
        self.faces_list = []
        self.labels_list = []

        raw_path = os.path.join(self.root_dir, 'raw')

        
        ## Parse the categories file
        cat_path = os.path.join(self.root_dir, 'categories.txt')
        with open(cat_path) as cat_file:
            cat_file.readline() # skip the first two lines
            cat_file.readline() 
            for i_class in range(30):
                cat_file.readline() 
                class_name, _, count = cat_file.readline().strip().split()
                count = int(count)
                self.class_names.append(class_name)
                mesh_list = []
                for j in range(20):
                    mesh_list.append(cat_file.readline().strip())


                # Randomly grab samples for this split. If given, disallow any samples in commmon with exclude_dict (ie making sure train set is distinct from test).
                order = np.random.permutation(len(mesh_list))
                added = 0
                self.entries[class_name] = set()
                for ind in order:
                    if(split_size is not None and added == split_size): continue

                    name = mesh_list[ind]
                    if exclude_dict is not None and name in exclude_dict[class_name]:
                        continue

                    path = os.path.join(root_dir, "raw", "T{}.off".format(name))

                    verts, faces = pp3d.read_mesh(path)
                    verts = torch.tensor(verts).float()
                    faces = torch.tensor(faces)
                    #print(verts)
                    #print(faces)

                    # center and unit scale
                    verts = fannet.geometry.normalize_positions(verts)

                    self.verts_list.append(verts)
                    self.faces_list.append(faces)
                    self.labels_list.append(i_class)
                    self.entries[class_name].add(name)

                    added += 1

                print(class_name + " -- " + " ".join([p for p in self.entries[class_name]]))

                if(split_size is not None and added < split_size):
                    raise ValueError("could not find enough entries to generate requested split")
                
        for ind, label in enumerate(self.labels_list):
            self.labels_list[ind] = torch.tensor(label)

    def __len__(self):
        return len(self.verts_list)

    def __getitem__(self, idx):
        verts = self.verts_list[idx]
        faces = self.faces_list[idx]
        label = self.labels_list[idx]
        frames, mass, L, evals, evecs, gradX, gradY, fanIndices = fannet.geometry.get_operators(verts, faces, k_eig=self.k_eig, op_cache_dir=self.op_cache_dir, spoke_length=self.spoke_length)
        return verts, faces, frames, mass, L, evals, evecs, gradX, gradY, label , fanIndices

class Shrec11MeshDataset_Simplified(Dataset):

    # NOTE: Remeshed data from MeshCNN authors.
    # Can be downloaded here (link from the MeshCNN authors): https://www.dropbox.com/s/w16st84r6wc57u7/shrec_16.tar.gz. Note that despite the filename, this really is the shapes from the SHREC 2011 dataset. Extract it to the `[ROOT_DIR]/raw/` directory.

    def __init__(self, root_dir, split_size, k_eig, exclude_dict=None, op_cache_dir=None, spoke_length=[0.00,0.01,0.02,0.03]):
        
        self.root_dir = root_dir
        self.n_class = 30 
        self.split_size = split_size # pass None to take all entries (except those in exclude_dict)
        self.k_eig = k_eig
        self.op_cache_dir = op_cache_dir
        self.spoke_length = spoke_length

        self.class_names = [ 'alien', 'ants', 'armadillo', 'bird1', 'bird2', 'camel', 'cat', 'centaur', 'dinosaur', 'dino_ske', 'dog1', 'dog2', 'flamingo', 'glasses', 'gorilla', 'hand', 'horse', 'lamp', 'laptop', 'man', 'myScissor', 'octopus', 'pliers', 'rabbit', 'santa', 'shark', 'snake', 'spiders', 'two_balls', 'woman']
        
        self.entries = {}

        # store in memory
        self.verts_list = []
        self.faces_list = []
        self.labels_list = []

        raw_path = os.path.join(self.root_dir, 'raw', "shrec_16")

        for class_idx, class_name in enumerate(self.class_names):
            
            # load both train and test subdirectories; we are manually regenerating random splits to do multiple trials
            mesh_files = []
            for t in ['test', 'train']:
                files = os.listdir(os.path.join(raw_path, class_name, t))
                for f in files:
                    full_f = os.path.join(raw_path, class_name, t, f)
                    mesh_files.append(full_f)


            # Randomly grab samples for this split. If given, disallow any samples in commmon with exclude_dict (ie making sure train set is distinct from test).
            order = np.random.permutation(len(mesh_files))
            added = 0
            self.entries[class_name] = set()
            for ind in order:
                if(split_size is not None and added == split_size): continue

                path = mesh_files[ind]
                if exclude_dict is not None and path in exclude_dict[class_name]:
                    continue

                verts, faces = pp3d.read_mesh(path)
                verts = torch.tensor(verts).float()
                faces = torch.tensor(faces)
                #print(verts.shape)
                #print(faces.shape)

                # center and unit scale
                verts = fannet.geometry.normalize_positions(verts)

                self.verts_list.append(verts)
                self.faces_list.append(faces)
                self.labels_list.append(class_idx)
                self.entries[class_name].add(path)

                added += 1

            print(class_name + " -- " + " ".join([os.path.basename(p) for p in self.entries[class_name]]))

            if(split_size is not None and added < split_size):
                raise ValueError("could not find enough entries to generate requested split")
            
        for ind, label in enumerate(self.labels_list):
            self.labels_list[ind] = torch.tensor(label)

        # Precompute operators
        self.frames_list, self.massvec_list, self.L_list, self.evals_list, self.evecs_list, self.gradX_list, self.gradY_list, self.fanIndices = fannet.geometry.get_all_operators(self.verts_list, 
                                                    self.faces_list, k_eig=self.k_eig, op_cache_dir=self.op_cache_dir, spoke_length=self.spoke_length)

    def __len__(self):
        return len(self.verts_list)

    def __getitem__(self, idx):
        return self.verts_list[idx], self.faces_list[idx], self.frames_list[idx], self.massvec_list[idx], self.L_list[idx], self.evals_list[idx], self.evecs_list[idx], self.gradX_list[idx], self.gradY_list[idx], self.labels_list[idx], self.fanIndices[idx]