import os
import torch
from data.base_dataset import BaseDataset
from util.util import is_mesh_file, pad, pad_vertices
import numpy as np
from models.layers.mesh import Mesh

class AutoEncoderData(BaseDataset):

    def __init__(self, opt):
        BaseDataset.__init__(self, opt)
        self.opt = opt
        self.device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu')
        self.root = opt.dataroot
        self.dir = os.path.join(opt.dataroot)
        self.classes, self.class_to_idx = self.find_classes(self.dir)
        print(self.classes, self.class_to_idx)
        self.paths = self.make_dataset_by_class(self.dir, self.class_to_idx, opt.phase)
        self.nclasses = len(self.classes)
        self.size = len(self.paths)
        self.get_mean_std()
        # # modify for network later.
        opt.nclasses = self.nclasses
        opt.input_nc = self.ninput_channels

    def __getitem__(self, index):
        path = self.paths[index][0]
        #print(path)
        mesh = Mesh(file=path, opt=self.opt, hold_history=True, export_folder=self.opt.export_folder)
        meta = {}
        mesh.vs = (mesh.vs - np.mean(mesh.vs, 0)) / np.std(mesh.vs, 0)
        meta['mesh'] = mesh
        meta['export_folder'] = mesh.export_folder
        meta['filename'] = mesh.filename
        # get edge features
        edge_features = mesh.extract_features()
        edge_features = pad(edge_features, self.opt.ninput_edges)
        vs, pad_iter = pad_vertices(mesh.vs, 1402)
        meta['edge_features'] = (edge_features - self.mean) / self.std
        meta['label'] = vs.astype(np.float)
        meta['init_faces'] = mesh.init_faces
        meta['pad_iter'] = pad_iter
        return meta

    def __len__(self):
        return self.size

    @staticmethod
    def make_dataset(path):
        meshes = []
        assert os.path.isdir(path), '%s is not a valid directory' % path

        for root, _, fnames in sorted(os.walk(path)):
            for fname in fnames:
                if is_mesh_file(fname):
                    path = os.path.join(root, fname)
                    meshes.append(path)

        return meshes
    
    @staticmethod    
    def find_classes(dir):
        classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
        classes.sort()
        if '.ipynb_checkpoints' in classes:
            classes.remove('.ipynb_checkpoints')
        class_to_idx = {classes[i]: i for i in range(len(classes))}
        return classes, class_to_idx

    @staticmethod
    def make_dataset_by_class(dir, class_to_idx, phase):
        meshes = []
        dir = os.path.expanduser(dir)
        for target in sorted(os.listdir(dir)):
            d = os.path.join(dir, target)
            if not os.path.isdir(d):
                continue
            for root, _, fnames in sorted(os.walk(d)):
                ## LEE
                #file_num = 0
                for fname in sorted(fnames):
                    if is_mesh_file(fname) and (root.count(phase)==1):
                        path = os.path.join(root, fname)
                        item = (path, class_to_idx[target])
                        meshes.append(item)
                        
                        ## LEE
                        #file_num += 1
                        #if file_num == 100 :
                        #   break
        return meshes