import numpy as np
import os
import math, random
from kaldiio import ReadHelper
from torch.utils.data.dataset import Dataset
from .spec_aug import spec_augment 
from .load_labels import get_labels, get_speakerID

#### Dataset directories: Train + Test + Valid
cwd = os.getcwd()
train_dataset='scp:'+cwd+'/data/dump/trainNet_sp/feats.scp'
test_dataset='scp:'+cwd+'/data/dump/test/feats.scp'
valid_dataset='scp:'+cwd+'/data/dump/valid/feats.scp'

### Label file
label_file=cwd+'/local/data_cleaned.xlsx'
map_dict = get_labels(label_file)

# TRAIN DATA
class Train_Dataset_height_triplet_mse(Dataset):
    # load the dataset
    def __init__(self):        
        with ReadHelper(train_dataset) as reader:
            self.dic = { u:d for u,d in reader }
        self.dic_keys = list(self.dic.keys())
        print(len(self.dic_keys))
        
    # number of rows in the dataset
    def __len__(self):
        return len(self.dic)
 
    # get a row at an index
    def __getitem__(self, idx):
        
        anchor_data = self.dic[self.dic_keys[idx]]        
        anchor_label = get_speakerID(self.dic_keys[idx])        
        n_anchor = math.floor(map_dict[anchor_label[1:5]][0] / 5) - 28
        
        positive_list = []
        negative_list = []
        for i in self.dic_keys:
            j = get_speakerID(i)

            if math.floor(map_dict[j[1:5]][0] / 5)-28 == n_anchor:
                positive_list.append(i)
            else:
                negative_list.append(i)

        positive_label = random.choice(positive_list)
        positive_data = self.dic[positive_label]

        negative_label = random.choice(negative_list)
        negative_data = self.dic[negative_label]

        if anchor_data.shape[0] < 800:
            anchor_data = np.concatenate([anchor_data, np.array([[0]*83]*(800-anchor_data.shape[0]))])            
        elif anchor_data.shape[0] > 800:
            anchor_data = anchor_data[:800]
            
        if anchor_label[0] == 'F':
            anchor_data = (np.concatenate((anchor_data, np.array([1]*800).reshape(800,1)), axis=1))
        elif anchor_label[0] == 'M':
            anchor_data = (np.concatenate((anchor_data, np.array([0]*800).reshape(800,1)), axis=1))

#####################

        if positive_data.shape[0] < 800:
            positive_data = np.concatenate([positive_data, np.array([[0]*83]*(800-positive_data.shape[0]))])
        elif positive_data.shape[0] > 800:
            positive_data = positive_data[:800]
            
        if get_speakerID(positive_label)[0] == 'F':
            positive_data = (np.concatenate((positive_data, np.array([1]*800).reshape(800,1)), axis=1))
        elif get_speakerID(positive_label)[0] == 'M':
            positive_data = (np.concatenate((positive_data, np.array([0]*800).reshape(800,1)), axis=1))

#####################

        if negative_data.shape[0] < 800:
            negative_data = np.concatenate([negative_data, np.array([[0]*83]*(800-negative_data.shape[0]))])
        elif negative_data.shape[0] > 800:
            negative_data = negative_data[:800]
            
        if get_speakerID(negative_label)[0] == 'F':
            negative_data = (np.concatenate((negative_data, np.array([1]*800).reshape(800,1)), axis=1))                
        elif get_speakerID(negative_label)[0] == 'M':
            negative_data = (np.concatenate((negative_data, np.array([0]*800).reshape(800,1)), axis=1))
        
        anchor_data, positive_data, negative_data = spec_augment(anchor_data), spec_augment(positive_data), spec_augment(negative_data)
        
        label = float(map_dict[anchor_label[1:5]][0])
        #print(label)
        #label = self.label_dict[label]
            
        return anchor_data, positive_data, negative_data, label
        
# TEST DATA 
class Test_Dataset_height_triplet_mse(Dataset):
    # load the dataset
    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = { u:d for u,d in reader }
        self.dic_keys = list(self.dic.keys())
        
    # number of rows in the dataset
    def __len__(self):
        return len(self.dic)
 
    # get a row at an index
    def __getitem__(self, idx):
        
        data = self.dic[self.dic_keys[idx]]
        label = get_speakerID(self.dic_keys[idx])
        
        if data.shape[0] < 800:
            data = np.concatenate([data, np.array([[0]*83]*(800-data.shape[0]))])              
        elif data.shape[0] > 800:
            data = data[:800]
            
        if label[0] == 'F':
            data = (np.concatenate((data, np.array([1]*800).reshape(800,1)), axis=1))                
        elif label[0] == 'M':
            data = (np.concatenate((data, np.array([0]*800).reshape(800,1)), axis=1))
            
        if label[0] == 'M':  
            gender = 0
        else:
            gender = 1

        label_ht = float(map_dict[label[1:5]][0])
        label_age = float(map_dict[label[1:5]][1])
        labels = [label_ht, label_age, gender]
            
        return data, label_ht, gender

        
# VALIDATION DATA
class Val_Dataset_height_triplet_mse(Dataset):
    # load the dataset
    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = { u:d for u,d in reader }
        self.dic_keys = list(self.dic.keys())
        
    # number of rows in the dataset
    def __len__(self):
        return len(self.dic)
 
    # get a row at an index
    def __getitem__(self, idx):
        
        anchor_data = self.dic[self.dic_keys[idx]]
        anchor_label = get_speakerID(self.dic_keys[idx])
        n_anchor = math.floor(map_dict[anchor_label[1:5]][0] / 5) - 28
        
        positive_list = []
        negative_list = []
        for i in self.dic_keys:            
            j = get_speakerID(i)
            if math.floor(map_dict[j[1:5]][0] / 5)-28 == n_anchor:
                positive_list.append(i)
            else:
                negative_list.append(i)

        positive_label = random.choice(positive_list)
        positive_data = self.dic[positive_label]

        negative_label = random.choice(negative_list)
        negative_data = self.dic[negative_label]

        if anchor_data.shape[0] < 800:
            anchor_data = np.concatenate([anchor_data, np.array([[0]*83]*(800-anchor_data.shape[0]))])              
        elif anchor_data.shape[0] > 800:
            anchor_data = anchor_data[:800]
            
        if anchor_label[0] == 'F':
            anchor_data = (np.concatenate((anchor_data, np.array([1]*800).reshape(800,1)), axis=1))
        elif anchor_label[0] == 'M':
            anchor_data = (np.concatenate((anchor_data, np.array([0]*800).reshape(800,1)), axis=1))

#####################

        if positive_data.shape[0] < 800:
            positive_data = np.concatenate([positive_data, np.array([[0]*83]*(800-positive_data.shape[0]))])              
        elif positive_data.shape[0] > 800:
            positive_data = positive_data[:800]
            
        if get_speakerID(positive_label)[0] == 'F':
            positive_data = (np.concatenate((positive_data, np.array([1]*800).reshape(800,1)), axis=1))                
        elif get_speakerID(positive_label)[0] == 'M':
            positive_data = (np.concatenate((positive_data, np.array([0]*800).reshape(800,1)), axis=1))

#####################

        if negative_data.shape[0] < 800:
            negative_data = np.concatenate([negative_data, np.array([[0]*83]*(800-negative_data.shape[0]))])              
        elif negative_data.shape[0] > 800:
            negative_data = negative_data[:800]
            
        if get_speakerID(negative_label)[0] == 'F':
            negative_data = (np.concatenate((negative_data, np.array([1]*800).reshape(800,1)), axis=1))                
        elif get_speakerID(negative_label)[0] == 'M':
            negative_data = (np.concatenate((negative_data, np.array([0]*800).reshape(800,1)), axis=1))
            
        #print(data.shape)
        
        #anchor_data, positive_data, negative_data = spec_augment(anchor_data), spec_augment(positive_data), spec_augment(negative_data)
        
        label = float(map_dict[anchor_label[1:5]][0])
            
        return anchor_data, positive_data, negative_data, label
9d39928186ced761cde1863b3b9ce171180578c8
<mask token>


class Val_Dataset_height_triplet_mse(Dataset):

    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = {u: d for u, d in reader}
        self.dic_keys = list(self.dic.keys())

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

    def __getitem__(self, idx):
        anchor_data = self.dic[self.dic_keys[idx]]
        anchor_label = get_speakerID(self.dic_keys[idx])
        n_anchor = math.floor(map_dict[anchor_label[1:5]][0] / 5) - 28
        positive_list = []
        negative_list = []
        for i in self.dic_keys:
            j = get_speakerID(i)
            if math.floor(map_dict[j[1:5]][0] / 5) - 28 == n_anchor:
                positive_list.append(i)
            else:
                negative_list.append(i)
        positive_label = random.choice(positive_list)
        positive_data = self.dic[positive_label]
        negative_label = random.choice(negative_list)
        negative_data = self.dic[negative_label]
        if anchor_data.shape[0] < 800:
            anchor_data = np.concatenate([anchor_data, np.array([[0] * 83] *
                (800 - anchor_data.shape[0]))])
        elif anchor_data.shape[0] > 800:
            anchor_data = anchor_data[:800]
        if anchor_label[0] == 'F':
            anchor_data = np.concatenate((anchor_data, np.array([1] * 800).
                reshape(800, 1)), axis=1)
        elif anchor_label[0] == 'M':
            anchor_data = np.concatenate((anchor_data, np.array([0] * 800).
                reshape(800, 1)), axis=1)
        if positive_data.shape[0] < 800:
            positive_data = np.concatenate([positive_data, np.array([[0] * 
                83] * (800 - positive_data.shape[0]))])
        elif positive_data.shape[0] > 800:
            positive_data = positive_data[:800]
        if get_speakerID(positive_label)[0] == 'F':
            positive_data = np.concatenate((positive_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(positive_label)[0] == 'M':
            positive_data = np.concatenate((positive_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        if negative_data.shape[0] < 800:
            negative_data = np.concatenate([negative_data, np.array([[0] * 
                83] * (800 - negative_data.shape[0]))])
        elif negative_data.shape[0] > 800:
            negative_data = negative_data[:800]
        if get_speakerID(negative_label)[0] == 'F':
            negative_data = np.concatenate((negative_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(negative_label)[0] == 'M':
            negative_data = np.concatenate((negative_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        label = float(map_dict[anchor_label[1:5]][0])
        return anchor_data, positive_data, negative_data, label

<mask token>


class Train_Dataset_height_triplet_mse(Dataset):
    <mask token>
    <mask token>
    <mask token>


class Test_Dataset_height_triplet_mse(Dataset):

    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = {u: d for u, d in reader}
        self.dic_keys = list(self.dic.keys())

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

    def __getitem__(self, idx):
        data = self.dic[self.dic_keys[idx]]
        label = get_speakerID(self.dic_keys[idx])
        if data.shape[0] < 800:
            data = np.concatenate([data, np.array([[0] * 83] * (800 - data.
                shape[0]))])
        elif data.shape[0] > 800:
            data = data[:800]
        if label[0] == 'F':
            data = np.concatenate((data, np.array([1] * 800).reshape(800, 1
                )), axis=1)
        elif label[0] == 'M':
            data = np.concatenate((data, np.array([0] * 800).reshape(800, 1
                )), axis=1)
        if label[0] == 'M':
            gender = 0
        else:
            gender = 1
        label_ht = float(map_dict[label[1:5]][0])
        label_age = float(map_dict[label[1:5]][1])
        labels = [label_ht, label_age, gender]
        return data, label_ht, gender


class Val_Dataset_height_triplet_mse(Dataset):

    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = {u: d for u, d in reader}
        self.dic_keys = list(self.dic.keys())

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

    def __getitem__(self, idx):
        anchor_data = self.dic[self.dic_keys[idx]]
        anchor_label = get_speakerID(self.dic_keys[idx])
        n_anchor = math.floor(map_dict[anchor_label[1:5]][0] / 5) - 28
        positive_list = []
        negative_list = []
        for i in self.dic_keys:
            j = get_speakerID(i)
            if math.floor(map_dict[j[1:5]][0] / 5) - 28 == n_anchor:
                positive_list.append(i)
            else:
                negative_list.append(i)
        positive_label = random.choice(positive_list)
        positive_data = self.dic[positive_label]
        negative_label = random.choice(negative_list)
        negative_data = self.dic[negative_label]
        if anchor_data.shape[0] < 800:
            anchor_data = np.concatenate([anchor_data, np.array([[0] * 83] *
                (800 - anchor_data.shape[0]))])
        elif anchor_data.shape[0] > 800:
            anchor_data = anchor_data[:800]
        if anchor_label[0] == 'F':
            anchor_data = np.concatenate((anchor_data, np.array([1] * 800).
                reshape(800, 1)), axis=1)
        elif anchor_label[0] == 'M':
            anchor_data = np.concatenate((anchor_data, np.array([0] * 800).
                reshape(800, 1)), axis=1)
        if positive_data.shape[0] < 800:
            positive_data = np.concatenate([positive_data, np.array([[0] * 
                83] * (800 - positive_data.shape[0]))])
        elif positive_data.shape[0] > 800:
            positive_data = positive_data[:800]
        if get_speakerID(positive_label)[0] == 'F':
            positive_data = np.concatenate((positive_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(positive_label)[0] == 'M':
            positive_data = np.concatenate((positive_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        if negative_data.shape[0] < 800:
            negative_data = np.concatenate([negative_data, np.array([[0] * 
                83] * (800 - negative_data.shape[0]))])
        elif negative_data.shape[0] > 800:
            negative_data = negative_data[:800]
        if get_speakerID(negative_label)[0] == 'F':
            negative_data = np.concatenate((negative_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(negative_label)[0] == 'M':
            negative_data = np.concatenate((negative_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        label = float(map_dict[anchor_label[1:5]][0])
        return anchor_data, positive_data, negative_data, label

<mask token>


class Train_Dataset_height_triplet_mse(Dataset):

    def __init__(self):
        with ReadHelper(train_dataset) as reader:
            self.dic = {u: d for u, d in reader}
        self.dic_keys = list(self.dic.keys())
        print(len(self.dic_keys))

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

    def __getitem__(self, idx):
        anchor_data = self.dic[self.dic_keys[idx]]
        anchor_label = get_speakerID(self.dic_keys[idx])
        n_anchor = math.floor(map_dict[anchor_label[1:5]][0] / 5) - 28
        positive_list = []
        negative_list = []
        for i in self.dic_keys:
            j = get_speakerID(i)
            if math.floor(map_dict[j[1:5]][0] / 5) - 28 == n_anchor:
                positive_list.append(i)
            else:
                negative_list.append(i)
        positive_label = random.choice(positive_list)
        positive_data = self.dic[positive_label]
        negative_label = random.choice(negative_list)
        negative_data = self.dic[negative_label]
        if anchor_data.shape[0] < 800:
            anchor_data = np.concatenate([anchor_data, np.array([[0] * 83] *
                (800 - anchor_data.shape[0]))])
        elif anchor_data.shape[0] > 800:
            anchor_data = anchor_data[:800]
        if anchor_label[0] == 'F':
            anchor_data = np.concatenate((anchor_data, np.array([1] * 800).
                reshape(800, 1)), axis=1)
        elif anchor_label[0] == 'M':
            anchor_data = np.concatenate((anchor_data, np.array([0] * 800).
                reshape(800, 1)), axis=1)
        if positive_data.shape[0] < 800:
            positive_data = np.concatenate([positive_data, np.array([[0] * 
                83] * (800 - positive_data.shape[0]))])
        elif positive_data.shape[0] > 800:
            positive_data = positive_data[:800]
        if get_speakerID(positive_label)[0] == 'F':
            positive_data = np.concatenate((positive_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(positive_label)[0] == 'M':
            positive_data = np.concatenate((positive_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        if negative_data.shape[0] < 800:
            negative_data = np.concatenate([negative_data, np.array([[0] * 
                83] * (800 - negative_data.shape[0]))])
        elif negative_data.shape[0] > 800:
            negative_data = negative_data[:800]
        if get_speakerID(negative_label)[0] == 'F':
            negative_data = np.concatenate((negative_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(negative_label)[0] == 'M':
            negative_data = np.concatenate((negative_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        anchor_data, positive_data, negative_data = spec_augment(anchor_data
            ), spec_augment(positive_data), spec_augment(negative_data)
        label = float(map_dict[anchor_label[1:5]][0])
        return anchor_data, positive_data, negative_data, label


class Test_Dataset_height_triplet_mse(Dataset):

    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = {u: d for u, d in reader}
        self.dic_keys = list(self.dic.keys())

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

    def __getitem__(self, idx):
        data = self.dic[self.dic_keys[idx]]
        label = get_speakerID(self.dic_keys[idx])
        if data.shape[0] < 800:
            data = np.concatenate([data, np.array([[0] * 83] * (800 - data.
                shape[0]))])
        elif data.shape[0] > 800:
            data = data[:800]
        if label[0] == 'F':
            data = np.concatenate((data, np.array([1] * 800).reshape(800, 1
                )), axis=1)
        elif label[0] == 'M':
            data = np.concatenate((data, np.array([0] * 800).reshape(800, 1
                )), axis=1)
        if label[0] == 'M':
            gender = 0
        else:
            gender = 1
        label_ht = float(map_dict[label[1:5]][0])
        label_age = float(map_dict[label[1:5]][1])
        labels = [label_ht, label_age, gender]
        return data, label_ht, gender


class Val_Dataset_height_triplet_mse(Dataset):

    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = {u: d for u, d in reader}
        self.dic_keys = list(self.dic.keys())

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

    def __getitem__(self, idx):
        anchor_data = self.dic[self.dic_keys[idx]]
        anchor_label = get_speakerID(self.dic_keys[idx])
        n_anchor = math.floor(map_dict[anchor_label[1:5]][0] / 5) - 28
        positive_list = []
        negative_list = []
        for i in self.dic_keys:
            j = get_speakerID(i)
            if math.floor(map_dict[j[1:5]][0] / 5) - 28 == n_anchor:
                positive_list.append(i)
            else:
                negative_list.append(i)
        positive_label = random.choice(positive_list)
        positive_data = self.dic[positive_label]
        negative_label = random.choice(negative_list)
        negative_data = self.dic[negative_label]
        if anchor_data.shape[0] < 800:
            anchor_data = np.concatenate([anchor_data, np.array([[0] * 83] *
                (800 - anchor_data.shape[0]))])
        elif anchor_data.shape[0] > 800:
            anchor_data = anchor_data[:800]
        if anchor_label[0] == 'F':
            anchor_data = np.concatenate((anchor_data, np.array([1] * 800).
                reshape(800, 1)), axis=1)
        elif anchor_label[0] == 'M':
            anchor_data = np.concatenate((anchor_data, np.array([0] * 800).
                reshape(800, 1)), axis=1)
        if positive_data.shape[0] < 800:
            positive_data = np.concatenate([positive_data, np.array([[0] * 
                83] * (800 - positive_data.shape[0]))])
        elif positive_data.shape[0] > 800:
            positive_data = positive_data[:800]
        if get_speakerID(positive_label)[0] == 'F':
            positive_data = np.concatenate((positive_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(positive_label)[0] == 'M':
            positive_data = np.concatenate((positive_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        if negative_data.shape[0] < 800:
            negative_data = np.concatenate([negative_data, np.array([[0] * 
                83] * (800 - negative_data.shape[0]))])
        elif negative_data.shape[0] > 800:
            negative_data = negative_data[:800]
        if get_speakerID(negative_label)[0] == 'F':
            negative_data = np.concatenate((negative_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(negative_label)[0] == 'M':
            negative_data = np.concatenate((negative_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        label = float(map_dict[anchor_label[1:5]][0])
        return anchor_data, positive_data, negative_data, label

<mask token>
cwd = os.getcwd()
train_dataset = 'scp:' + cwd + '/data/dump/trainNet_sp/feats.scp'
test_dataset = 'scp:' + cwd + '/data/dump/test/feats.scp'
valid_dataset = 'scp:' + cwd + '/data/dump/valid/feats.scp'
label_file = cwd + '/local/data_cleaned.xlsx'
map_dict = get_labels(label_file)


class Train_Dataset_height_triplet_mse(Dataset):

    def __init__(self):
        with ReadHelper(train_dataset) as reader:
            self.dic = {u: d for u, d in reader}
        self.dic_keys = list(self.dic.keys())
        print(len(self.dic_keys))

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

    def __getitem__(self, idx):
        anchor_data = self.dic[self.dic_keys[idx]]
        anchor_label = get_speakerID(self.dic_keys[idx])
        n_anchor = math.floor(map_dict[anchor_label[1:5]][0] / 5) - 28
        positive_list = []
        negative_list = []
        for i in self.dic_keys:
            j = get_speakerID(i)
            if math.floor(map_dict[j[1:5]][0] / 5) - 28 == n_anchor:
                positive_list.append(i)
            else:
                negative_list.append(i)
        positive_label = random.choice(positive_list)
        positive_data = self.dic[positive_label]
        negative_label = random.choice(negative_list)
        negative_data = self.dic[negative_label]
        if anchor_data.shape[0] < 800:
            anchor_data = np.concatenate([anchor_data, np.array([[0] * 83] *
                (800 - anchor_data.shape[0]))])
        elif anchor_data.shape[0] > 800:
            anchor_data = anchor_data[:800]
        if anchor_label[0] == 'F':
            anchor_data = np.concatenate((anchor_data, np.array([1] * 800).
                reshape(800, 1)), axis=1)
        elif anchor_label[0] == 'M':
            anchor_data = np.concatenate((anchor_data, np.array([0] * 800).
                reshape(800, 1)), axis=1)
        if positive_data.shape[0] < 800:
            positive_data = np.concatenate([positive_data, np.array([[0] * 
                83] * (800 - positive_data.shape[0]))])
        elif positive_data.shape[0] > 800:
            positive_data = positive_data[:800]
        if get_speakerID(positive_label)[0] == 'F':
            positive_data = np.concatenate((positive_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(positive_label)[0] == 'M':
            positive_data = np.concatenate((positive_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        if negative_data.shape[0] < 800:
            negative_data = np.concatenate([negative_data, np.array([[0] * 
                83] * (800 - negative_data.shape[0]))])
        elif negative_data.shape[0] > 800:
            negative_data = negative_data[:800]
        if get_speakerID(negative_label)[0] == 'F':
            negative_data = np.concatenate((negative_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(negative_label)[0] == 'M':
            negative_data = np.concatenate((negative_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        anchor_data, positive_data, negative_data = spec_augment(anchor_data
            ), spec_augment(positive_data), spec_augment(negative_data)
        label = float(map_dict[anchor_label[1:5]][0])
        return anchor_data, positive_data, negative_data, label


class Test_Dataset_height_triplet_mse(Dataset):

    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = {u: d for u, d in reader}
        self.dic_keys = list(self.dic.keys())

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

    def __getitem__(self, idx):
        data = self.dic[self.dic_keys[idx]]
        label = get_speakerID(self.dic_keys[idx])
        if data.shape[0] < 800:
            data = np.concatenate([data, np.array([[0] * 83] * (800 - data.
                shape[0]))])
        elif data.shape[0] > 800:
            data = data[:800]
        if label[0] == 'F':
            data = np.concatenate((data, np.array([1] * 800).reshape(800, 1
                )), axis=1)
        elif label[0] == 'M':
            data = np.concatenate((data, np.array([0] * 800).reshape(800, 1
                )), axis=1)
        if label[0] == 'M':
            gender = 0
        else:
            gender = 1
        label_ht = float(map_dict[label[1:5]][0])
        label_age = float(map_dict[label[1:5]][1])
        labels = [label_ht, label_age, gender]
        return data, label_ht, gender


class Val_Dataset_height_triplet_mse(Dataset):

    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = {u: d for u, d in reader}
        self.dic_keys = list(self.dic.keys())

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

    def __getitem__(self, idx):
        anchor_data = self.dic[self.dic_keys[idx]]
        anchor_label = get_speakerID(self.dic_keys[idx])
        n_anchor = math.floor(map_dict[anchor_label[1:5]][0] / 5) - 28
        positive_list = []
        negative_list = []
        for i in self.dic_keys:
            j = get_speakerID(i)
            if math.floor(map_dict[j[1:5]][0] / 5) - 28 == n_anchor:
                positive_list.append(i)
            else:
                negative_list.append(i)
        positive_label = random.choice(positive_list)
        positive_data = self.dic[positive_label]
        negative_label = random.choice(negative_list)
        negative_data = self.dic[negative_label]
        if anchor_data.shape[0] < 800:
            anchor_data = np.concatenate([anchor_data, np.array([[0] * 83] *
                (800 - anchor_data.shape[0]))])
        elif anchor_data.shape[0] > 800:
            anchor_data = anchor_data[:800]
        if anchor_label[0] == 'F':
            anchor_data = np.concatenate((anchor_data, np.array([1] * 800).
                reshape(800, 1)), axis=1)
        elif anchor_label[0] == 'M':
            anchor_data = np.concatenate((anchor_data, np.array([0] * 800).
                reshape(800, 1)), axis=1)
        if positive_data.shape[0] < 800:
            positive_data = np.concatenate([positive_data, np.array([[0] * 
                83] * (800 - positive_data.shape[0]))])
        elif positive_data.shape[0] > 800:
            positive_data = positive_data[:800]
        if get_speakerID(positive_label)[0] == 'F':
            positive_data = np.concatenate((positive_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(positive_label)[0] == 'M':
            positive_data = np.concatenate((positive_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        if negative_data.shape[0] < 800:
            negative_data = np.concatenate([negative_data, np.array([[0] * 
                83] * (800 - negative_data.shape[0]))])
        elif negative_data.shape[0] > 800:
            negative_data = negative_data[:800]
        if get_speakerID(negative_label)[0] == 'F':
            negative_data = np.concatenate((negative_data, np.array([1] * 
                800).reshape(800, 1)), axis=1)
        elif get_speakerID(negative_label)[0] == 'M':
            negative_data = np.concatenate((negative_data, np.array([0] * 
                800).reshape(800, 1)), axis=1)
        label = float(map_dict[anchor_label[1:5]][0])
        return anchor_data, positive_data, negative_data, label

import numpy as np
import os
import math, random
from kaldiio import ReadHelper
from torch.utils.data.dataset import Dataset
from .spec_aug import spec_augment 
from .load_labels import get_labels, get_speakerID

#### Dataset directories: Train + Test + Valid
cwd = os.getcwd()
train_dataset='scp:'+cwd+'/data/dump/trainNet_sp/feats.scp'
test_dataset='scp:'+cwd+'/data/dump/test/feats.scp'
valid_dataset='scp:'+cwd+'/data/dump/valid/feats.scp'

### Label file
label_file=cwd+'/local/data_cleaned.xlsx'
map_dict = get_labels(label_file)

# TRAIN DATA
class Train_Dataset_height_triplet_mse(Dataset):
    # load the dataset
    def __init__(self):        
        with ReadHelper(train_dataset) as reader:
            self.dic = { u:d for u,d in reader }
        self.dic_keys = list(self.dic.keys())
        print(len(self.dic_keys))
        
    # number of rows in the dataset
    def __len__(self):
        return len(self.dic)
 
    # get a row at an index
    def __getitem__(self, idx):
        
        anchor_data = self.dic[self.dic_keys[idx]]        
        anchor_label = get_speakerID(self.dic_keys[idx])        
        n_anchor = math.floor(map_dict[anchor_label[1:5]][0] / 5) - 28
        
        positive_list = []
        negative_list = []
        for i in self.dic_keys:
            j = get_speakerID(i)

            if math.floor(map_dict[j[1:5]][0] / 5)-28 == n_anchor:
                positive_list.append(i)
            else:
                negative_list.append(i)

        positive_label = random.choice(positive_list)
        positive_data = self.dic[positive_label]

        negative_label = random.choice(negative_list)
        negative_data = self.dic[negative_label]

        if anchor_data.shape[0] < 800:
            anchor_data = np.concatenate([anchor_data, np.array([[0]*83]*(800-anchor_data.shape[0]))])            
        elif anchor_data.shape[0] > 800:
            anchor_data = anchor_data[:800]
            
        if anchor_label[0] == 'F':
            anchor_data = (np.concatenate((anchor_data, np.array([1]*800).reshape(800,1)), axis=1))
        elif anchor_label[0] == 'M':
            anchor_data = (np.concatenate((anchor_data, np.array([0]*800).reshape(800,1)), axis=1))

#####################

        if positive_data.shape[0] < 800:
            positive_data = np.concatenate([positive_data, np.array([[0]*83]*(800-positive_data.shape[0]))])
        elif positive_data.shape[0] > 800:
            positive_data = positive_data[:800]
            
        if get_speakerID(positive_label)[0] == 'F':
            positive_data = (np.concatenate((positive_data, np.array([1]*800).reshape(800,1)), axis=1))
        elif get_speakerID(positive_label)[0] == 'M':
            positive_data = (np.concatenate((positive_data, np.array([0]*800).reshape(800,1)), axis=1))

#####################

        if negative_data.shape[0] < 800:
            negative_data = np.concatenate([negative_data, np.array([[0]*83]*(800-negative_data.shape[0]))])
        elif negative_data.shape[0] > 800:
            negative_data = negative_data[:800]
            
        if get_speakerID(negative_label)[0] == 'F':
            negative_data = (np.concatenate((negative_data, np.array([1]*800).reshape(800,1)), axis=1))                
        elif get_speakerID(negative_label)[0] == 'M':
            negative_data = (np.concatenate((negative_data, np.array([0]*800).reshape(800,1)), axis=1))
        
        anchor_data, positive_data, negative_data = spec_augment(anchor_data), spec_augment(positive_data), spec_augment(negative_data)
        
        label = float(map_dict[anchor_label[1:5]][0])
        #print(label)
        #label = self.label_dict[label]
            
        return anchor_data, positive_data, negative_data, label
        
# TEST DATA 
class Test_Dataset_height_triplet_mse(Dataset):
    # load the dataset
    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = { u:d for u,d in reader }
        self.dic_keys = list(self.dic.keys())
        
    # number of rows in the dataset
    def __len__(self):
        return len(self.dic)
 
    # get a row at an index
    def __getitem__(self, idx):
        
        data = self.dic[self.dic_keys[idx]]
        label = get_speakerID(self.dic_keys[idx])
        
        if data.shape[0] < 800:
            data = np.concatenate([data, np.array([[0]*83]*(800-data.shape[0]))])              
        elif data.shape[0] > 800:
            data = data[:800]
            
        if label[0] == 'F':
            data = (np.concatenate((data, np.array([1]*800).reshape(800,1)), axis=1))                
        elif label[0] == 'M':
            data = (np.concatenate((data, np.array([0]*800).reshape(800,1)), axis=1))
            
        if label[0] == 'M':  
            gender = 0
        else:
            gender = 1

        label_ht = float(map_dict[label[1:5]][0])
        label_age = float(map_dict[label[1:5]][1])
        labels = [label_ht, label_age, gender]
            
        return data, label_ht, gender

        
# VALIDATION DATA
class Val_Dataset_height_triplet_mse(Dataset):
    # load the dataset
    def __init__(self):
        with ReadHelper(test_dataset) as reader:
            self.dic = { u:d for u,d in reader }
        self.dic_keys = list(self.dic.keys())
        
    # number of rows in the dataset
    def __len__(self):
        return len(self.dic)
 
    # get a row at an index
    def __getitem__(self, idx):
        
        anchor_data = self.dic[self.dic_keys[idx]]
        anchor_label = get_speakerID(self.dic_keys[idx])
        n_anchor = math.floor(map_dict[anchor_label[1:5]][0] / 5) - 28
        
        positive_list = []
        negative_list = []
        for i in self.dic_keys:            
            j = get_speakerID(i)
            if math.floor(map_dict[j[1:5]][0] / 5)-28 == n_anchor:
                positive_list.append(i)
            else:
                negative_list.append(i)

        positive_label = random.choice(positive_list)
        positive_data = self.dic[positive_label]

        negative_label = random.choice(negative_list)
        negative_data = self.dic[negative_label]

        if anchor_data.shape[0] < 800:
            anchor_data = np.concatenate([anchor_data, np.array([[0]*83]*(800-anchor_data.shape[0]))])              
        elif anchor_data.shape[0] > 800:
            anchor_data = anchor_data[:800]
            
        if anchor_label[0] == 'F':
            anchor_data = (np.concatenate((anchor_data, np.array([1]*800).reshape(800,1)), axis=1))
        elif anchor_label[0] == 'M':
            anchor_data = (np.concatenate((anchor_data, np.array([0]*800).reshape(800,1)), axis=1))

#####################

        if positive_data.shape[0] < 800:
            positive_data = np.concatenate([positive_data, np.array([[0]*83]*(800-positive_data.shape[0]))])              
        elif positive_data.shape[0] > 800:
            positive_data = positive_data[:800]
            
        if get_speakerID(positive_label)[0] == 'F':
            positive_data = (np.concatenate((positive_data, np.array([1]*800).reshape(800,1)), axis=1))                
        elif get_speakerID(positive_label)[0] == 'M':
            positive_data = (np.concatenate((positive_data, np.array([0]*800).reshape(800,1)), axis=1))

#####################

        if negative_data.shape[0] < 800:
            negative_data = np.concatenate([negative_data, np.array([[0]*83]*(800-negative_data.shape[0]))])              
        elif negative_data.shape[0] > 800:
            negative_data = negative_data[:800]
            
        if get_speakerID(negative_label)[0] == 'F':
            negative_data = (np.concatenate((negative_data, np.array([1]*800).reshape(800,1)), axis=1))                
        elif get_speakerID(negative_label)[0] == 'M':
            negative_data = (np.concatenate((negative_data, np.array([0]*800).reshape(800,1)), axis=1))
            
        #print(data.shape)
        
        #anchor_data, positive_data, negative_data = spec_augment(anchor_data), spec_augment(positive_data), spec_augment(negative_data)
        
        label = float(map_dict[anchor_label[1:5]][0])
            
        return anchor_data, positive_data, negative_data, label