#! /usr/bin/python
# -*- encoding: utf-8 -*-

import torch
import numpy
import random
import pdb
import glob
import os
from torch.utils.data import Dataset, DataLoader
from PIL import Image

def round_down(num, divisor):
    return num - (num%divisor)

def worker_init_fn(worker_id):
    numpy.random.seed(numpy.random.get_state()[1][0] + worker_id)

class meta_loader(Dataset):
    def __init__(self, train_path, train_ext, transform):
        
        ## Read Training Files
        files = glob.glob('%s/*/*.%s'%(train_path,train_ext))

        ## Make a mapping from Class Name to Class Number
        dictkeys = list(set([x.split('/')[-2] for x in files]))
        dictkeys.sort()
        dictkeys = { key : ii for ii, key in enumerate(dictkeys) }

        self.transform  = transform

        self.label_dict = {}
        self.data_list  = []
        self.data_label = []
        
        for lidx, file in enumerate(files):
            speaker_name = file.split('/')[-2]
            speaker_label = dictkeys[speaker_name];

            if not (speaker_label in self.label_dict):
                self.label_dict[speaker_label] = [];

            self.label_dict[speaker_label].append(lidx);
            
            self.data_label.append(speaker_label)
            self.data_list.append(file)

        print('%d files from %d classes found.'%(len(self.data_list),len(self.label_dict)))

    def __getitem__(self, indices):

        feat = []
        for index in indices:
            feat.append(self.transform(Image.open(self.data_list[index])));
        feat = numpy.stack(feat, axis=0)

        return torch.FloatTensor(feat), self.data_label[index]

    def __len__(self):

        return len(self.data_list)

class test_dataset_loader(Dataset):
    def __init__(self, test_list, test_path, transform, **kwargs):
        self.test_path  = test_path
        self.data_list  = test_list
        self.transform  = transform

    def __getitem__(self, index):
        img = Image.open(os.path.join(self.test_path, self.data_list[index]))
        return self.transform(img), self.data_list[index]

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


class meta_sampler(torch.utils.data.Sampler):
    def __init__(self, data_source, nPerClass, max_img_per_cls, batch_size):

        self.label_dict         = data_source.label_dict
        self.nPerClass          = nPerClass
        self.max_img_per_cls    = max_img_per_cls;
        self.batch_size         = batch_size;
        
    def __iter__(self):
        
        ## Get a list of identities
        dictkeys = list(self.label_dict.keys());
        dictkeys.sort()

        lol = lambda lst, sz: [lst[i:i+sz] for i in range(0, len(lst), sz)]

        flattened_list = []
        flattened_label = []

        ## Data for each class
        for findex, key in enumerate(dictkeys):
            data    = self.label_dict[key]
            numSeg  = round_down(min(len(data),self.max_img_per_cls),self.nPerClass)
            
            rp      = lol(numpy.random.permutation(len(data))[:numSeg],self.nPerClass)
            flattened_label.extend([findex] * (len(rp)))
            for indices in rp:
                flattened_list.append([data[i] for i in indices])

        ## Data in random order
        mixid           = numpy.random.permutation(len(flattened_label))
        mixlabel        = []
        mixmap          = []

        ## Prevent two pairs of the same speaker in the same batch
        for ii in mixid:
            startbatch = len(mixlabel) - len(mixlabel) % self.batch_size
            if flattened_label[ii] not in mixlabel[startbatch:]:
                mixlabel.append(flattened_label[ii])
                mixmap.append(ii)
        
        return iter([flattened_list[i] for i in mixmap])
    
    def __len__(self):
        return len(self.data_source)

def get_data_loader(batch_size, max_img_per_cls, nDataLoaderThread, nPerClass, train_path, train_ext, transform, **kwargs):
    
    train_dataset = meta_loader(train_path, train_ext, transform)

    train_sampler = meta_sampler(train_dataset, nPerClass, max_img_per_cls, batch_size)

    train_loader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=batch_size,
        num_workers=nDataLoaderThread,
        sampler=train_sampler,
        pin_memory=False,
        worker_init_fn=worker_init_fn,
        drop_last=True,
    )
    
    return train_loader