#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Sep 15 17:29:27 2020

@author: fa19
"""

import nibabel as nb

import numpy as np
import torch
import random
import torch.nn.functional as F

means = np.load('../dHCP_mean_seg.npy')
std = np.load('../dHCP_std_seg.npy')
means = torch.from_numpy(means)
stds = torch.from_numpy(std)

### SEGMENTATION FILES ###

unwarped_files_directory = '/data/Data/benchmarking/fsaverage_32k_30_01_2021/ico6'
unwarped_labels_directory = '/data/Data/dHCP/M-CRIB-S/template_space/ico6L'
warped_files_directory = '/data/Data/benchmarking/fsaverage_32k_30_01_2021/ico6_warped'
warped_labels_directory = '/data/Data/dHCP/M-CRIB-S/template_space/ico6L_warp'

#unwarped_files_directory='/data/Data/derivatives_native_ico6_seg/features'
#warped_files_directory='/data/Data/derivatives_native_ico6_seg/features_warp'
#unwarped_labels_directory ='/data/Data/derivatives_native_ico6_seg/labels'
#warped_labels_directory ='/data/Data/derivatives_native_ico6_seg/labels_warp'

#means = torch.Tensor([0.0345])
#stds = torch.Tensor([0.1906])

test_rotation_arr = np.load('../GraphMethods/data/unseen_rots.npy').astype(int)
rotation_arr = np.load('../GraphMethods/data/rotations_array.npy').astype(int)
reversing_arr = np.load('../GraphMethods/data/reversing_arr.npy')

#smoothing_arr = [[0, 10], [-12,34], [11,200], [-1,1]]

# minima and maxima defines as mean +/- 4*std

lower_bound = torch.Tensor([ -0.2819,  -0.7279,  -0.5199, -16.1347])
upper_bound = torch.Tensor([ 2.5354,  0.7970,  2.5550, 16.2459])

minima = torch.Tensor([  0.0000,  -0.7279,  -0.3271, -14.8748])
maxima = torch.Tensor([ 2.5354,  0.7970,  2.5550, 12.1209])

xy_points = np.load('equirectangular_ico_6_points.npy')
xy_points[:,0] = (xy_points[:,0] + 0.1)%1
grid = np.load('grid_170_square.npy')

grid_x, grid_y = np.meshgrid(np.linspace(0.02, 0.98, 170), np.linspace(0.02, 0.98, 170))
grid[:,0] = grid_x.flatten()
grid[:,1] = grid_y.flatten()

from scipy.interpolate import griddata

"""

unwarped_files_directory: the directory of all the files in input_arr. BOTH L and R
aarped_files_directory: the directory of all the warped files.

warped directory could be the same as unwarped directory

"""

class My_dHCP_Data(torch.utils.data.Dataset):

    def __init__(self, input_arr, rotations = False,
                 number_of_warps = 0, parity_choice = 'left', smoothing = False, normalisation = None, sample_only = True, output_as_torch = True ):

        """

        A Full Dataset for the dHCP Data. Can include warps, rotations and parity flips.

        Fileanme style:

            in the array: only 'sub-X-ses-Y'
            but for the filenames themselves
                Left = 'sub-X_ses-Y_L'
                Right = 'sub-X_ses-Y_R'
                if warped:
                    'sub-X_ses-Y_L_W1'

        INPUT ARGS:

            1. input_arr:
                Numpy array size Nx2
                FIRST index MUST be the filename (excluding directory AND L or R ) of MERGED nibabel files
                LAST index must be the (float) label
                (OPTIONAL) Middle index if size 3 (optional) is any confounding metadata (Also float, e.g scan age for predicting birth age)

            2 . rotations - boolean: to add rotations or not to add rotations

            3. number of warps to include - INT
                NB WARPED AR INCLUDED AS FILENAME CHANGES. WARP NUMBER X IS WRITTEN AS filename_WX
                NUMBER OF WARPS CANNOT EXCEED NUMBER OF WARPES PRESENT IN FILES

            4. Particy Choice (JMPORTANT!) - defines left and right-ness

                If: 'left'- will output ONLY LEFT
                If: 'both' - will randomly choose L or R
                If 'combined' - will output a combined array (left first), will be eventually read as a file with twice the number of input channels. as they will be stacked together

            5. smoothing - boolean, will clip extremal values according to the smoothing_array

            6. normalisation - str. Will normalise according to 'range', 'std' or 'None'
                Range is from -1 to 1
                Std is mean = 0, std = 1

            7. output_as_torch - boolean:
                outputs values as torch Tensors if you want (usually yes)

        """

        self.input_arr = input_arr

        self.image_files = input_arr[:,0]
        self.label = input_arr[:,-1]

        self.sample_only = sample_only
        self.rotations = rotations

        self.number_of_warps = number_of_warps

        self.parity = parity_choice

        self.smoothing = smoothing
        self.normalisation = normalisation

        self.output_as_torch = output_as_torch
        if self.number_of_warps != 0 and self.number_of_warps != None:
            self.directory = warped_files_directory
        else:
            self.directory = unwarped_files_directory

    def __len__(self):

        L = len(self.input_arr)

        if self.number_of_warps !=0:
            if self.sample_only == False:
                L = L*self.number_of_warps

        return L

    def __test_input_params__(self):
        assert self.input_arr.shape[1] >=2, 'check your input array is a nunpy array of files and labels'
        assert type(self.number_of_warps) == int, "number of warps must be an in integer (can be 0)"
        assert self.parity in ['left', 'both', 'combined'], "parity choice must be either left, combined or both"
        if self.number_of_rotations != 0:
            assert self.rotation_arr != None,'Must specify a rotation file containing rotation vertex ids if rotations are non-zero'
        assert self.rotations == bool, 'rotations must be boolean'
        assert self.normalisation in [None, 'none', 'std', 'range'], 'Normalisation must be either std or range'

    def __genfilename__(self,idx):

        """
        gets the appropriate file based on input parameters on PARITY and on WARPS

        """
        # grab raw filename

        if self.number_of_warps != 0:
            warp_choice = str(1 + idx//len(self.input_arr))

            idx = idx%len(self.input_arr)

        raw_filename = self.image_files[idx]

        # add parity to it. IN THE FORM OF A LIST!  If requries both will output a list of length 2
        filename = []

        if self.parity == 'left':
            filename.append(raw_filename + '_L')
        elif self.parity == 'both':
            coin_flip = random.randint(0,1)
            if coin_flip == 0:
                filename.append(raw_filename + '_L')
            elif coin_flip == 1:
                filename.append(raw_filename + '_R')
        elif self.parity == 'combined':
            filename.append(raw_filename + '_L')
            filename.append(raw_filename+'_R')

        # filename is now a list of the correct filenames.

        # now add warps if required

        if self.number_of_warps != 0:

            filename = [s + '_W'+warp_choice for s in filename ]

        return filename

    def __getitem__(self, idx):

        """
        First load the images and collect them as numpy arrays

        then collect the label

        then collect the metadata (though might be None)
        """

        filename = self.__genfilename__(idx)

        image_gifti = [nb.load(self.directory + '/'+individual_filename+'.shape.gii').darrays for individual_filename in filename]

        image = []
        if self.rotations == True:

            rotation_choice = random.randint(0, len(rotation_arr)-1)
            if rotation_choice !=0:
                for file in image_gifti:
                    image.extend(item.data[rotation_arr[rotation_choice]] for item in file)
            else:
                for file in image_gifti:
                    image.extend(item.data for item in file)
        else:
            for file in image_gifti:
                image.extend(item.data for item in file)

        ### labels
        if self.number_of_warps != 0:

            idx = idx%len(self.input_arr)
        label = self.label[idx]

        ###### metadata grabbing if necessary

        if self.input_arr.shape[1] > 2:

            self.metadata = input_arr[:,1:-1]

        else:
            self.metadata = None

        if self.smoothing != False:
            for i in range(len(image)):
                image[i] = np.clip(image[i], lower_bound[i%len(lower_bound)].item(), upper_bound[i%len(upper_bound)].item())

        # torchify if required:

        if self.normalisation != None:
            if self.normalisation == 'std':
                for i in range(len(image)):

                    image[i] = ( image[i] - means[i%len(means)].item( )) / stds[i%len(stds)].item()

            elif self.normalisation == 'range':
                for i in range(len(image)):

                    image[i] = image[i] - minima[i%len(minima)].item()
                    image[i] = image[i] / (maxima[i%len(maxima)].item()- minima[i%len(minima)].item())

        if self.output_as_torch:
            image = torch.Tensor( image )

            label = torch.Tensor( [label] )

            if self.metadata != None:

                metadata = torch.Tensor( [self.metadata] )

        if self.metadata != None:
            sample = {'image': image, 'metadata' : self.metadata, 'label': label}

        else:
            sample = {'image': image,'label': label}

        return sample

"""

examples:

file_arr = np.load('/home/fa19/Documents/dHCP_Data_merged/scan_age_regression_full_shuffled_18-08-2020.npy', allow_pickle = True)

My_dHCP_Data(file_arr, rotations=True, smoothing = True, parity_choice='both')

My_dHCP_Data(file_arr, rotations=True, smoothing = True, parity_choice='combined')

My_dHCP_Data(file_arr, number_of_warps = 5, rotations=True, smoothing = False, parity_choice='left')

"""
def get_global_mean_and_std_from_ds(ds):
    nb_samples = 0
    num_channels = ds[0]['image'].size(0)
    channel_mean =  torch.zeros(num_channels)
    channel_var = torch.zeros(num_channels)
    #channel_std = torch.Tensor([0., 0., 0.])
    for samples in ds:
        # scale image to be between 0 and 1
        images = samples['image']

        channel_mean += images.mean(1)
        channel_var += images.var(1)
        nb_samples += 1

    channel_mean /= nb_samples
    channel_var /= nb_samples
    channel_std = np.sqrt(channel_var)

    return channel_mean, channel_std

def get_global_min_and_max_from_ds(ds):

    num_channels = ds[0]['image'].size(0)

    running_minima = torch.ones(num_channels)*100
    running_maxima = torch.ones(num_channels)*-100

    for samples in ds:
        # scale image to be between 0 and 1
        images = samples['image']

        image_minima = torch.min(images, dim=1)[0]
        image_maxima = torch.max(images, dim=1)[0]
        for i in range(len(image_minima)):
            if image_minima[i] < running_minima[i]:
                running_minima[i] = image_minima[i].item()
            if image_maxima[i] > running_maxima[i]:
                running_maxima[i] = image_maxima[i].item()

    return running_minima, running_maxima

class My_Projected_dHCP_Data(torch.utils.data.Dataset):

    def __init__(self, input_arr, rotations = False,
                 number_of_warps = 0, parity_choice = 'left', smoothing = False, normalisation = None, projected =False, sample_only = True, output_as_torch = True ):

        """

        A Full Dataset for the dHCP Data. Can include warps, rotations and parity flips.

        Fileanme style:

            in the array: only 'sub-X-ses-Y'
            but for the filenames themselves
                Left = 'sub-X_ses-Y_L'
                Right = 'sub-X_ses-Y_R'
                if warped:
                    'sub-X_ses-Y_L_W1'

        INPUT ARGS:

            1. input_arr:
                Numpy array size Nx2
                FIRST index MUST be the filename (excluding directory AND L or R ) of MERGED nibabel files
                LAST index must be the (float) label
                (OPTIONAL) Middle index if size 3 (optional) is any confounding metadata (Also float, e.g scan age for predicting birth age)

            2 . rotations - boolean: to add rotations or not to add rotations

            3. number of warps to include - INT
                NB WARPED AR INCLUDED AS FILENAME CHANGES. WARP NUMBER X IS WRITTEN AS filename_WX
                NUMBER OF WARPS CANNOT EXCEED NUMBER OF WARPES PRESENT IN FILES

            4. Particy Choice (JMPORTANT!) - defines left and right-ness

                If: 'left'- will output ONLY LEFT
                If: 'both' - will randomly choose L or R
                If 'combined' - will output a combined array (left first), will be eventually read as a file with twice the number of input channels. as they will be stacked together

            5. smoothing - boolean, will clip extremal values according to the smoothing_array

            6. normalisation - str. Will normalise according to 'range', 'std' or 'None'
                Range is from -1 to 1
                Std is mean = 0, std = 1

            7. output_as_torch - boolean:
                outputs values as torch Tensors if you want (usually yes)

        """

        self.input_arr = input_arr

        self.image_files = input_arr[:,0]
        self.label = input_arr[:,-1]

        self.rotations = rotations

        self.projected = projected

        self.number_of_warps = number_of_warps

        self.parity = parity_choice

        self.smoothing = smoothing
        self.normalisation = normalisation
        self.sample_only = sample_only

        self.output_as_torch = output_as_torch
        if self.number_of_warps != 0 and self.number_of_warps != None:
            self.directory = warped_files_directory
        else:
            self.directory = unwarped_files_directory

    def __len__(self):

        L = len(self.input_arr)

        if self.number_of_warps !=0:
            if self.sample_only == False:
                L = L*self.number_of_warps

        return L

    def __test_input_params__(self):
        assert self.input_arr.shape[1] >=2, 'check your input array is a nunpy array of files and labels'
        assert type(self.number_of_warps) == int, "number of warps must be an in integer (can be 0)"
        assert self.parity in ['left', 'both', 'combined'], "parity choice must be either left, combined or both"
        if self.number_of_rotations != 0:
            assert self.rotation_arr != None,'Must specify a rotation file containing rotation vertex ids if rotations are non-zero'
        assert self.rotations == bool, 'rotations must be boolean'
        assert self.normalisation in [None, 'none', 'std', 'range'], 'Normalisation must be either std or range'

    def __genfilename__(self,idx):

        """
        gets the appropriate file based on input parameters on PARITY and on WARPS

        """
        # grab raw filename

        raw_filename = self.image_files[idx]

        # add parity to it. IN THE FORM OF A LIST!  If requries both will output a list of length 2
        filename = []

        if self.parity == 'left':
            filename.append(raw_filename + '_L')
        elif self.parity == 'both':
            coin_flip = random.randint(0,1)
            if coin_flip == 0:
                filename.append(raw_filename + '_L')
            elif coin_flip == 1:
                filename.append(raw_filename + '_R')
        elif self.parity == 'combined':
            filename.append(raw_filename + '_L')
            filename.append(raw_filename+'_R')

        # filename is now a list of the correct filenames.

        # now add warps if required

        if self.number_of_warps != 0:
            warp_choice = str(random.randint(1,self.number_of_warps))
            filename = [s + '_W'+warp_choice for s in filename ]

        return filename

    def __getitem__(self, idx):

        """
        First load the images and collect them as numpy arrays

        then collect the label

        then collect the metadata (though might be None)
        """

        filename = self.__genfilename__(idx)

        image_gifti = [nb.load(self.directory + '/'+individual_filename+'.shape.gii').darrays for individual_filename in filename]

        image = []
        if self.rotations == True:

            rotation_choice = random.randint(0, len(rotation_arr)-1)
            if rotation_choice !=0:
                for file in image_gifti:
                    image.extend(item.data[rotation_arr[rotation_choice]] for item in file)
            else:
                for file in image_gifti:
                    image.extend(item.data for item in file)
        else:
            for file in image_gifti:
                image.extend(item.data for item in file)

        ### labels
        if self.number_of_warps != 0:

            idx = idx%len(self.input_arr)
        label = self.label[idx]

        ###### metadata grabbing if necessary

        if self.input_arr.shape[1] > 2:

            self.metadata = input_arr[:,1:-1]

        else:
            self.metadata = None

        if self.smoothing != False:
            for i in range(len(image)):
                image[i] = np.clip(image[i], lower_bound[i%len(lower_bound)].item(), upper_bound[i%len(upper_bound)].item())

        # torchify if required:

        if self.normalisation != None:
            if self.normalisation == 'std':
                for i in range(len(image)):

                    image[i] = ( image[i] - means[i%len(means)].item( )) / stds[i%len(stds)].item()

            elif self.normalisation == 'range':
                for i in range(len(image)):

                    image[i] = image[i] - minima[i%len(minima)].item()
                    image[i] = image[i] / (maxima[i%len(maxima)].item()- minima[i%len(minima)].item())

        if self.output_as_torch:
            image = torch.Tensor( image )

            label = torch.Tensor( [label] )

            if self.metadata != None:

                metadata = torch.Tensor( [self.metadata] )

        if self.projected == True:
            image = griddata(xy_points, image.T, grid, 'linear')
            image = torch.Tensor(image.reshape(170,170,4)).permute(2,0,1)

        if self.metadata != None:
            sample = {'image': image, 'metadata' : self.metadata, 'label': label}

        else:
            sample = {'image': image,'label': label}

        return sample

class My_Projected_dHCP_Data_Segmentation(torch.utils.data.Dataset):

    def __init__(self, input_arr, rotations = False,
                 number_of_warps = 0, parity_choice = 'left', smoothing = False, normalisation = None, projected =False, sample_only = True, output_as_torch = True ):

        """

        A Full Dataset for the dHCP Data. Can include warps, rotations and parity flips.

        Fileanme style:

            in the array: only 'sub-X-ses-Y'
            but for the filenames themselves
                Left = 'sub-X_ses-Y_L'
                Right = 'sub-X_ses-Y_R'
                if warped:
                    'sub-X_ses-Y_L_W1'

        INPUT ARGS:

            1. input_arr:
                Numpy array size Nx2
                FIRST index MUST be the filename (excluding directory AND L or R ) of MERGED nibabel files
                LAST index must be the (float) label
                (OPTIONAL) Middle index if size 3 (optional) is any confounding metadata (Also float, e.g scan age for predicting birth age)

            2 . rotations - boolean: to add rotations or not to add rotations

            3. number of warps to include - INT
                NB WARPED AR INCLUDED AS FILENAME CHANGES. WARP NUMBER X IS WRITTEN AS filename_WX
                NUMBER OF WARPS CANNOT EXCEED NUMBER OF WARPES PRESENT IN FILES

            4. Particy Choice (JMPORTANT!) - defines left and right-ness

                If: 'left'- will output ONLY LEFT
                If: 'both' - will randomly choose L or R
                If 'combined' - will output a combined array (left first), will be eventually read as a file with twice the number of input channels. as they will be stacked together

            5. smoothing - boolean, will clip extremal values according to the smoothing_array

            6. normalisation - str. Will normalise according to 'range', 'std' or 'None'
                Range is from -1 to 1
                Std is mean = 0, std = 1

            7. output_as_torch - boolean:
                outputs values as torch Tensors if you want (usually yes)

        """

        self.input_arr = input_arr

        self.image_files = input_arr[:,0]

        self.rotations = rotations

        self.projected = projected

        self.number_of_warps = number_of_warps

        self.parity = parity_choice

        self.smoothing = smoothing
        self.normalisation = normalisation
        self.sample_only = sample_only

        self.output_as_torch = output_as_torch
        if self.number_of_warps != 0 and self.number_of_warps != None:
            self.directory = warped_files_directory
        else:
            self.directory = unwarped_files_directory

        if self.number_of_warps != 0 and self.number_of_warps != None:
            self.label_directory = warped_labels_directory
        else:
            self.label_directory = unwarped_labels_directory

    def __len__(self):

        L = len(self.input_arr)

        if self.number_of_warps !=0:
            if self.sample_only == False:
                L = L*self.number_of_warps

        return L

    def __test_input_params__(self):
        assert self.input_arr.shape[1] >=2, 'check your input array is a nunpy array of files and labels'
        assert type(self.number_of_warps) == int, "number of warps must be an in integer (can be 0)"
        assert self.parity in ['left', 'both', 'combined'], "parity choice must be either left, combined or both"
        if self.number_of_rotations != 0:
            assert self.rotation_arr != None,'Must specify a rotation file containing rotation vertex ids if rotations are non-zero'
        assert self.rotations == bool, 'rotations must be boolean'
        assert self.normalisation in [None, 'none', 'std', 'range'], 'Normalisation must be either std or range'

    def __genfilename__(self,idx, right):

        """
        gets the appropriate file based on input parameters on PARITY and on WARPS

        """
        # grab raw filename

        raw_filename = self.image_files[idx]

        # add parity to it. IN THE FORM OF A LIST!  If requries both will output a list of length 2
        filename = []

        if self.parity != 'combined':
            if right == True:
                filename.append(raw_filename + '_R')

            else:

               filename.append(raw_filename + '_L')

#        if self.parity == 'left':
#            filename.append(raw_filename + '_L')
#
#        elif self.parity == 'both':
#            coin_flip = random.randint(0,1)
#            if coin_flip == 0:
#                filename.append(raw_filename + '_L')
#            elif coin_flip == 1:
#                filename.append(raw_filename + '_R')
#                right = True

        if self.parity == 'combined':
            filename.append(raw_filename + '_L')
            filename.append(raw_filename+'_R')

        # filename is now a list of the correct filenames.

        # now add warps if required

        if self.number_of_warps != 0:
            warp_choice = str(random.randint(0,self.number_of_warps))
            if warp_choice !='0':

                filename = [s + '_W'+warp_choice for s in filename ]

        return filename

    def __getitem__(self, idx):

        """
        First load the images and collect them as numpy arrays

        then collect the label

        then collect the metadata (though might be None)
        """

        if self.parity == 'both':
            T = self.__len__()//2

            idx, right  = idx % T, idx // T
            filename = self.__genfilename__(idx, right)
        else:
            right = False
            filename = self.__genfilename__(idx, right)

        image_gifti = [nb.load(self.directory + '/'+individual_filename+'.shape.gii').darrays for individual_filename in filename]
        label_gifti = [nb.load(self.label_directory + '/'+individual_filename+'.label.gii').darrays for individual_filename in filename]

        image = []
        label = []
        if self.rotations == True:

            rotation_choice = random.randint(0, len(rotation_arr)-1)
            if rotation_choice !=0:
                for file in image_gifti:
                    image.extend(item.data[rotation_arr[rotation_choice]] for item in file)

                for file in label_gifti:
                    label.extend(item.data[rotation_arr[rotation_choice]] for item in file)

            else:
                for file in image_gifti:
                    image.extend(item.data for item in file)
                for file in label_gifti:
                    label.extend(item.data for item in file)
        else:
            for file in image_gifti:
                image.extend(item.data for item in file)
            for file in label_gifti:
                label.extend(item.data for item in file)

        if right == True:
            image = [item[reversing_arr] for item in image]
            label = [item[reversing_arr] for item in label]

        ###### metadata grabbing if necessary

        if self.input_arr.shape[1] > 2:

            self.metadata = input_arr[:,1:-1]

        else:
            self.metadata = None

        if self.smoothing != False:
            for i in range(len(image)):
                image[i] = np.clip(image[i], lower_bound[i%len(lower_bound)].item(), upper_bound[i%len(upper_bound)].item())

        # torchify if required:

        if self.normalisation != None:
            if self.normalisation == 'std':
                for i in range(len(image)):

                    image[i] = ( image[i] - means[i%len(means)].item( )) / stds[i%len(stds)].item()

            elif self.normalisation == 'range':
                for i in range(len(image)):

                    image[i] = image[i] - minima[i%len(minima)].item()
                    image[i] = image[i] / (maxima[i%len(maxima)].item()- minima[i%len(minima)].item())

        if self.output_as_torch:
            image = torch.Tensor( image )

            label = torch.Tensor( label )

            if self.metadata != None:

                metadata = torch.Tensor( [self.metadata] )

        if self.projected == True:
            image = griddata(xy_points, image.T, grid, 'nearest')
            image = torch.Tensor(image.reshape(170,170,4)).permute(2,0,1)

            label = griddata(xy_points, label.T, grid, 'nearest')
            label = torch.Tensor(label.reshape(170,170,1))#.permute(2,0,1)
            label = F.one_hot(label.to(torch.int64), 37).contiguous()
            label = label.squeeze()
            label = label.permute(2,0,1)

        if self.metadata != None:
            sample = {'image': image, 'metadata' : self.metadata, 'label': label}

        else:
            sample = {'image': image,'label': label}

        return sample

class My_Projected_dHCP_Data_Segmentation_Test(torch.utils.data.Dataset):

    def __init__(self, input_arr, rotations = False,
                 number_of_warps = 0, parity_choice = 'left', smoothing = False, normalisation = None, projected =False, sample_only = True, output_as_torch = True ):

        """

        A Full Dataset for the dHCP Data. Can include warps, rotations and parity flips.

        Fileanme style:

            in the array: only 'sub-X-ses-Y'
            but for the filenames themselves
                Left = 'sub-X_ses-Y_L'
                Right = 'sub-X_ses-Y_R'
                if warped:
                    'sub-X_ses-Y_L_W1'

        INPUT ARGS:

            1. input_arr:
                Numpy array size Nx2
                FIRST index MUST be the filename (excluding directory AND L or R ) of MERGED nibabel files
                LAST index must be the (float) label
                (OPTIONAL) Middle index if size 3 (optional) is any confounding metadata (Also float, e.g scan age for predicting birth age)

            2 . rotations - boolean: to add rotations or not to add rotations

            3. number of warps to include - INT
                NB WARPED AR INCLUDED AS FILENAME CHANGES. WARP NUMBER X IS WRITTEN AS filename_WX
                NUMBER OF WARPS CANNOT EXCEED NUMBER OF WARPES PRESENT IN FILES

            4. Particy Choice (JMPORTANT!) - defines left and right-ness

                If: 'left'- will output ONLY LEFT
                If: 'both' - will randomly choose L or R
                If 'combined' - will output a combined array (left first), will be eventually read as a file with twice the number of input channels. as they will be stacked together

            5. smoothing - boolean, will clip extremal values according to the smoothing_array

            6. normalisation - str. Will normalise according to 'range', 'std' or 'None'
                Range is from -1 to 1
                Std is mean = 0, std = 1

            7. output_as_torch - boolean:
                outputs values as torch Tensors if you want (usually yes)

        """

        self.input_arr = input_arr

        self.image_files = input_arr[:,0]

        self.rotations = rotations

        self.projected = projected

        self.number_of_warps = number_of_warps

        self.parity = parity_choice

        self.smoothing = smoothing
        self.normalisation = normalisation
        self.sample_only = sample_only

        self.output_as_torch = output_as_torch
        if self.number_of_warps != 0 and self.number_of_warps != None:
            self.directory = warped_files_directory
        else:
            self.directory = unwarped_files_directory

        if self.number_of_warps != 0 and self.number_of_warps != None:
            self.label_directory = warped_labels_directory
        else:
            self.label_directory = unwarped_labels_directory

    def __len__(self):

        L = len(self.input_arr)

        if self.number_of_warps !=0:
            if self.sample_only == False:
                L = L*self.number_of_warps

        return L

    def __test_input_params__(self):
        assert self.input_arr.shape[1] >=2, 'check your input array is a nunpy array of files and labels'
        assert type(self.number_of_warps) == int, "number of warps must be an in integer (can be 0)"
        assert self.parity in ['left', 'both', 'combined'], "parity choice must be either left, combined or both"
        if self.number_of_rotations != 0:
            assert self.rotation_arr != None,'Must specify a rotation file containing rotation vertex ids if rotations are non-zero'
        assert self.rotations == bool, 'rotations must be boolean'
        assert self.normalisation in [None, 'none', 'std', 'range'], 'Normalisation must be either std or range'

    def __genfilename__(self,idx, right):

        """
        gets the appropriate file based on input parameters on PARITY and on WARPS

        """
        # grab raw filename

        raw_filename = self.image_files[idx]

        # add parity to it. IN THE FORM OF A LIST!  If requries both will output a list of length 2
        filename = []

        if self.parity != 'combined':
            if right == True:
                filename.append(raw_filename + '_R')

            else:

               filename.append(raw_filename + '_L')

#        if self.parity == 'left':
#            filename.append(raw_filename + '_L')
#
#        elif self.parity == 'both':
#            coin_flip = random.randint(0,1)
#            if coin_flip == 0:
#                filename.append(raw_filename + '_L')
#            elif coin_flip == 1:
#                filename.append(raw_filename + '_R')
#                right = True

        if self.parity == 'combined':
            filename.append(raw_filename + '_L')
            filename.append(raw_filename+'_R')

        # filename is now a list of the correct filenames.

        # now add warps if required

        if self.number_of_warps != 0:
            warp_choice = str(random.randint(0,self.number_of_warps))
            if warp_choice !='0':

                filename = [s + '_W'+warp_choice for s in filename ]

        return filename

    def __getitem__(self, idx):

        """
        First load the images and collect them as numpy arrays

        then collect the label

        then collect the metadata (though might be None)
        """

        if self.parity == 'both':
            T = self.__len__()//2

            idx, right  = idx % T, idx // T
            filename = self.__genfilename__(idx, right)
        else:
            right = False
            filename = self.__genfilename__(idx, right)

        image_gifti = [nb.load(self.directory + '/'+individual_filename+'.shape.gii').darrays for individual_filename in filename]
        label_gifti = [nb.load(self.label_directory + '/'+individual_filename+'.label.gii').darrays for individual_filename in filename]

        image = []
        label = []
        if self.rotations == True:

            rotation_choice = random.randint(1, len(test_rotation_arr)-1)
            if rotation_choice !=0:
                for file in image_gifti:
                    image.extend(item.data[test_rotation_arr[rotation_choice]] for item in file)

                for file in label_gifti:
                    label.extend(item.data[test_rotation_arr[rotation_choice]] for item in file)

            else:
                for file in image_gifti:
                    image.extend(item.data for item in file)
                for file in label_gifti:
                    label.extend(item.data for item in file)
        else:
            for file in image_gifti:
                image.extend(item.data for item in file)
            for file in label_gifti:
                label.extend(item.data for item in file)

        if right == True:
            image = [item[reversing_arr] for item in image]
            label = [item[reversing_arr] for item in label]

        ###### metadata grabbing if necessary

        if self.input_arr.shape[1] > 2:

            self.metadata = input_arr[:,1:-1]

        else:
            self.metadata = None

        if self.smoothing != False:
            for i in range(len(image)):
                image[i] = np.clip(image[i], lower_bound[i%len(lower_bound)].item(), upper_bound[i%len(upper_bound)].item())

        # torchify if required:

        if self.normalisation != None:
            if self.normalisation == 'std':
                for i in range(len(image)):

                    image[i] = ( image[i] - means[i%len(means)].item( )) / stds[i%len(stds)].item()

            elif self.normalisation == 'range':
                for i in range(len(image)):

                    image[i] = image[i] - minima[i%len(minima)].item()
                    image[i] = image[i] / (maxima[i%len(maxima)].item()- minima[i%len(minima)].item())

        if self.output_as_torch:
            image = torch.Tensor( image )

            label = torch.Tensor( label )

            if self.metadata != None:

                metadata = torch.Tensor( [self.metadata] )

        if self.projected == True:
            image = griddata(xy_points, image.T, grid, 'nearest')
            image = torch.Tensor(image.reshape(170,170,4)).permute(2,0,1)

            label = griddata(xy_points, label.T, grid, 'nearest')
            label = torch.Tensor(label.reshape(170,170,1))#.permute(2,0,1)
            label = F.one_hot(label.to(torch.int64), 37).contiguous()
            label = label.squeeze()
            label = label.permute(2,0,1)

        if self.metadata != None:
            sample = {'image': image, 'metadata' : self.metadata, 'label': label}

        else:
            sample = {'image': image,'label': label}

        return sample

class My_Linear_Projected_dHCP_Data(torch.utils.data.Dataset):

    def __init__(self, input_arr, rotations = False,
                 number_of_warps = 0, parity_choice = 'left', smoothing = False, normalisation = None, projected =False, output_as_torch = True ):

        """

        A Full Dataset for the dHCP Data. Can include warps, rotations and parity flips.

        Fileanme style:

            in the array: only 'sub-X-ses-Y'
            but for the filenames themselves
                Left = 'sub-X_ses-Y_L'
                Right = 'sub-X_ses-Y_R'
                if warped:
                    'sub-X_ses-Y_L_W1'

        INPUT ARGS:

            1. input_arr:
                Numpy array size Nx2
                FIRST index MUST be the filename (excluding directory AND L or R ) of MERGED nibabel files
                LAST index must be the (float) label
                (OPTIONAL) Middle index if size 3 (optional) is any confounding metadata (Also float, e.g scan age for predicting birth age)

            2 . rotations - boolean: to add rotations or not to add rotations

            3. number of warps to include - INT
                NB WARPED AR INCLUDED AS FILENAME CHANGES. WARP NUMBER X IS WRITTEN AS filename_WX
                NUMBER OF WARPS CANNOT EXCEED NUMBER OF WARPES PRESENT IN FILES

            4. Particy Choice (JMPORTANT!) - defines left and right-ness

                If: 'left'- will output ONLY LEFT
                If: 'both' - will randomly choose L or R
                If 'combined' - will output a combined array (left first), will be eventually read as a file with twice the number of input channels. as they will be stacked together

            5. smoothing - boolean, will clip extremal values according to the smoothing_array

            6. normalisation - str. Will normalise according to 'range', 'std' or 'None'
                Range is from -1 to 1
                Std is mean = 0, std = 1

            7. output_as_torch - boolean:
                outputs values as torch Tensors if you want (usually yes)

        """

        self.input_arr = input_arr

        self.image_files = input_arr[:,0]
        self.label = input_arr[:,-1]

        self.rotations = rotations

        self.projected = projected

        self.number_of_warps = number_of_warps

        self.parity = parity_choice

        self.smoothing = smoothing
        self.normalisation = normalisation

        self.output_as_torch = output_as_torch
        if self.number_of_warps != 0 and self.number_of_warps != None:
            self.directory = warped_files_directory
        else:
            self.directory = unwarped_files_directory

    def __len__(self):

        L = len(self.input_arr)

        if self.number_of_warps !=0:
            L = L*self.number_of_warps

        return L

    def __test_input_params__(self):
        assert self.input_arr.shape[1] >=2, 'check your input array is a nunpy array of files and labels'
        assert type(self.number_of_warps) == int, "number of warps must be an in integer (can be 0)"
        assert self.parity in ['left', 'both', 'combined'], "parity choice must be either left, combined or both"
        if self.number_of_rotations != 0:
            assert self.rotation_arr != None,'Must specify a rotation file containing rotation vertex ids if rotations are non-zero'
        assert self.rotations == bool, 'rotations must be boolean'
        assert self.normalisation in [None, 'none', 'std', 'range'], 'Normalisation must be either std or range'

    def __genfilename__(self,idx):

        """
        gets the appropriate file based on input parameters on PARITY and on WARPS

        """
        # grab raw filename

        raw_filename = self.image_files[idx]

        # add parity to it. IN THE FORM OF A LIST!  If requries both will output a list of length 2
        filename = []

        if self.parity == 'left':
            filename.append(raw_filename + '_L')
        elif self.parity == 'both':
            coin_flip = random.randint(0,1)
            if coin_flip == 0:
                filename.append(raw_filename + '_L')
            elif coin_flip == 1:
                filename.append(raw_filename + '_R')
        elif self.parity == 'combined':
            filename.append(raw_filename + '_L')
            filename.append(raw_filename+'_R')

        # filename is now a list of the correct filenames.

        # now add warps if required

        if self.number_of_warps != 0:
            warp_choice = str(random.randint(1,self.number_of_warps))
            filename = [s + '_W'+warp_choice for s in filename ]

        return filename

    def __getitem__(self, idx):

        """
        First load the images and collect them as numpy arrays

        then collect the label

        then collect the metadata (though might be None)
        """

        filename = self.__genfilename__(idx)

        image_gifti = [nb.load(self.directory + '/'+individual_filename+'.shape.gii').darrays for individual_filename in filename]

        image = []
        if self.rotations == True:

            rotation_choice = random.randint(0, len(rotation_arr)-1)
            if rotation_choice !=0:
                for file in image_gifti:
                    image.extend(item.data[rotation_arr[rotation_choice]] for item in file)
            else:
                for file in image_gifti:
                    image.extend(item.data for item in file)
        else:
            for file in image_gifti:
                image.extend(item.data for item in file)

        ### labels
        if self.number_of_warps != 0:

            idx = idx%len(self.input_arr)
        label = self.label[idx]

        ###### metadata grabbing if necessary

        if self.input_arr.shape[1] > 2:

            self.metadata = input_arr[:,1:-1]

        else:
            self.metadata = None

        if self.smoothing != False:
            for i in range(len(image)):
                image[i] = np.clip(image[i], lower_bound[i%len(lower_bound)].item(), upper_bound[i%len(upper_bound)].item())

        # torchify if required:

        if self.normalisation != None:
            if self.normalisation == 'std':
                for i in range(len(image)):

                    image[i] = ( image[i] - means[i%len(means)].item( )) / stds[i%len(stds)].item()

            elif self.normalisation == 'range':
                for i in range(len(image)):

                    image[i] = image[i] - minima[i%len(minima)].item()
                    image[i] = image[i] / (maxima[i%len(maxima)].item()- minima[i%len(minima)].item())

        if self.output_as_torch:
            image = torch.Tensor( image )

            label = torch.Tensor( [label] )

            if self.metadata != None:

                metadata = torch.Tensor( [self.metadata] )

        if self.projected == True:

            image = torch.Tensor(image.reshape(170,170,4)).permute(2,0,1)

        if self.metadata != None:
            sample = {'image': image, 'metadata' : self.metadata, 'label': label}

        else:
            sample = {'image': image,'label': label}

        return sample


import matplotlib.pyplot as plt
import numpy as np
import torch as tc
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D

def show_multiple_images_v1(imgs, lxy=None, titles=None, save_name=None, cmap=None):
    if cmap is None:
        cmap = plt.cm.gray
    ni = len(imgs)
    if lxy is None:
        lx = int(np.sqrt(ni)) + 1
        ly = int(ni / lx) + 1
    else:
        lx, ly = tuple(lxy)
    plt.figure()
    for n in range(ni):
        plt.subplot(lx, ly, n + 1)
        tmp = imgs[n].cpu().numpy()
        if tmp.ndim == 2:
            plt.imshow(tmp, cmap=cmap)
        else:
            plt.imshow(tmp)
        if titles is not None:
            plt.title(str(titles[n]))
        plt.axis('off')
        plt.xticks([])
        plt.yticks([])
    if type(save_name) is str:
        plt.savefig(save_name)
    plt.show()

def scatter3d(x, y, z, if_plot=True):
    if if_plot:
        plt.close()
    if type(x) is tc.Tensor:
        x = x.cpu().numpy()
        y = y.cpu().numpy()
        z = z.cpu().numpy()
    ax1 = plt.figure().add_subplot(111, projection='3d')
    ax1.set_title('Scatter Plot')
    plt.xlabel('X')
    plt.ylabel('Y')
    ax1.scatter(x, y, z, c='r', marker='o')
    plt.legend('x1')
    ax1.view_init(elev=90, azim=0)
    if if_plot:
        plt.draw()
        plt.pause(1)
    return ax1

def surf(x=None, y=None, z=None, xlabel='x label', ylabel='y label', zlabel='z label', title=''):
    fig = plt.figure()
    ax = Axes3D(fig)
    if x is None:
        x = np.array(range(z.shape[1]))
    elif type(x) is tc.Tensor:
        x = x.cpu().numpy()
    if y is None:
        y = np.array(range(z.shape[0]))
    elif type(y) is tc.Tensor:
        y = y.cpu().numpy()
    if type(z) is tc.Tensor:
        z = z.cpu().numpy()
    x, y = np.meshgrid(x, y)
    surf = ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap=cm.jet,
                           linewidth=0, antialiased=False)
    ax.set_xlabel(xlabel, color='r')
    ax.set_ylabel(ylabel, color='r')
    ax.set_zlabel(zlabel)
    # ax.zaxis.set_major_locator(LinearLocator(10))
    # ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))
    fig.colorbar(surf, shrink=0.5, aspect=5)
    fig.suptitle(title)
    plt.show()

def plot(x, *y, marker='s'):
    if type(x) is tc.Tensor:
        if x.device != 'cpu':
            x = x.cpu()
        x = x.numpy()
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    if len(y) > 0.5:
        for y0 in y:
            if type(y0) is tc.Tensor:
                if y0.device != 'cpu':
                    y0 = y0.cpu()
                y0 = y0.numpy()
            ax.plot(x, y0, marker=marker)
    else:
        ax.plot(x, marker=marker)
    plt.show()

def plot_v1(x, *y, options=None, save=None):
    if options is None:
        options = dict()
    num_curves = max(1, y.__len__())
    # Default font
    font1 = {'family': 'Times New Roman',
             'weight': 'normal',
             'size': 40}
    # Default values
    label_names = list()
    for n in range(num_curves):
        label_names.append('curve-' + str(n))
    default_ops = ['labelsize', 'axfontname', 'labelfont', 'axnames',
                   'legendfont', 'markers', 'labelnames']
    default_val = [32, 'Times New Roman', font1, ['', ''], font1,
                   ['o', 'v', '^', '<', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+'],
                   label_names]

    save_opts = dict()
    save_opts['name'] = 'img.png'
    if type(save) is dict:
        for s in save:
            save_opts[s] = save[s]

    opts = dict()
    for n in range(default_ops.__len__()):
        if default_ops[n] in options:
            opts[default_ops[n]] = options[default_ops[n]]
        else:
            opts[default_ops[n]] = default_val[n]
    while num_curves > opts['markers'].__len__():
        opts['markers'] = opts['markers'] * 2
    opts['markers'] = opts['markers'][:num_curves]

    # start plotting
    if type(x) is tc.Tensor:
        if x.device != 'cpu':
            x = x.cpu()
        x = x.numpy()
    x = np.array(x).reshape(-1,)
    fig = plt.figure(figsize=(13, 10))
    ax = fig.add_subplot(1, 1, 1)
    if len(y) > 0.5:
        n = 0
        for y0 in y:
            if type(y0) is tc.Tensor:
                if y0.device != 'cpu':
                    y0 = y0.cpu()
                y0 = y0.numpy()
            ax.plot(x, np.array(y0).reshape(-1,), marker=opts['markers'][n], markerfacecolor='w', markersize=12,
                    markeredgewidth=2, label=opts['labelnames'][n])
            n += 1
    else:
        ax.plot(x, marker=opts['markers'][0], markerfacecolor='w', markersize=12,
                markeredgewidth=2, label=opts['labelnames'][0])

    plt.tick_params(labelsize=opts['labelsize'])
    labels = ax.get_xticklabels() + ax.get_yticklabels()
    [label.set_fontname(opts['axfontname']) for label in labels]

    plt.xlabel(opts['axnames'][0], opts['labelfont'])
    plt.ylabel(opts['axnames'][1], opts['labelfont'])

    plt.legend(prop=opts['legendfont'])
    if type(save) is dict:
        # plt.subplots_adjust(left=0.09, right=1, wspace=0.25, hspace=0.25, bottom=0.13, top=0.91)
        plt.savefig(save_opts['name'])
    plt.show()


##
# @author : Francony Steven
# @brief : A file testing
##

import pytest
import numpy
import copy
import S1_algotools as algo

# Fixtures LISTS region
@pytest.fixture
def int_fixture ():
	return [1,2,3,4,-7]

@pytest.fixture
def positive_list_fixture () :
	return [2,4,8,10,12]

@pytest.fixture
def negative_list_fixture () :
	return [-2,-4,-8,-10,-12]

@pytest.fixture
def zero_list_fixture () :
 	return [0,0,0,0,0,0,0,0,0,0]

@pytest.fixture
def mix_list_fixture () :
 	return ['a',0,"string",0,0,0,0,0,None]
# END Fixtures LISTS region

# Fixtures MATRIX region
@pytest.fixture
def mat_fixture() :
	myMat = numpy.zeros([10, 10], dtype = int);
	myMat[2][2] = 1
	return myMat

@pytest.fixture
def mat_error_fixture() :
	return numpy.zeros([10, 10], dtype = int);

@pytest.fixture
def mat_little_fixture() :
	return numpy.zeros([5, 5], dtype = int);
 # END fixtures MATRIX region

 # Fixtures RANDOM ARRAY region
@pytest.fixture
def array_empty() :
 	size = 0
	myMat = numpy.full([size,size],'',dtype='str')
	return myMat

@pytest.fixture
def array_high() :
	size = 2
	myMat = numpy.full([size,size],'',dtype='str')
	return myMat

@pytest.fixture
def array_normal() :
 	size = 6
	myMat = numpy.full([size,size],'',dtype='str')
	return myMat
# END fixtures RANDOM ARRAY region

# Fixtures STRING region
@pytest.fixture
def string_empty() :
	return ""

@pytest.fixture
def string_spaced() :
	return "There is a little string"

@pytest.fixture
def string_unspaced() :
	return "Thereisalittlestring"
# END Fixtures STRING region

# Fixtures LIST region
@pytest.fixture
def list_empty() :
	list = []
	return list

@pytest.fixture
def list_normal() :
	list = range(10)
	return list
# END Fixtures LIST region

# AVERAGING
# Basic testing average function
def test_average_above_zero (int_fixture, positive_list_fixture, negative_list_fixture, zero_list_fixture, mix_list_fixture) :
	assert algo.average_above_zero(int_fixture) == 2.5
	assert algo.average_above_zero(positive_list_fixture) == 7.2

	with pytest.raises(ZeroDivisionError) :
		assert algo.average_above_zero(negative_list_fixture)

	with pytest.raises(ZeroDivisionError) :
		assert algo.average_above_zero(zero_list_fixture)

	with pytest.raises(TypeError) :
		assert algo.average_above_zero(mix_list_fixture)
# END AVERAGING TESTING

# TABLE MAXIMUM VALUE
# Basic max_value testing function
def test_max_value (int_fixture, positive_list_fixture, negative_list_fixture, zero_list_fixture, mix_list_fixture) :
	assert algo.max_value(int_fixture) == 4
	assert algo.max_value(positive_list_fixture) == 12
	assert algo.max_value(negative_list_fixture) == -2
	assert algo.max_value(zero_list_fixture) == 0

	with pytest.raises(ValueError) :
		assert algo.max_value(mix_list_fixture)
# END TABLE MAXIMUM TESTING

# TABLE MINIMUM VALUE
# Basic min_value testing function
def test_min_value (int_fixture, positive_list_fixture, negative_list_fixture, zero_list_fixture, mix_list_fixture) :
	assert algo.min_value(int_fixture) == -7
	assert algo.min_value(positive_list_fixture) == 2
	assert algo.min_value(negative_list_fixture) == -12
	assert algo.min_value(zero_list_fixture) == 0

	with pytest.raises(TypeError) :
		assert algo.min_value(mix_list_fixture)

# END TABLE MINIMUM TESTING

# REVERSE TABLE
# Basic reverse_table testing function
def test_reverse_table (int_fixture, positive_list_fixture, negative_list_fixture, zero_list_fixture, mix_list_fixture) :
	 assert algo.reverse_table(int_fixture) == [-7,4,3,2,1]
	 assert algo.reverse_table(positive_list_fixture) == [12,10,8,4,2]
	 assert algo.reverse_table(negative_list_fixture) == [-12,-10,-8,-4,-2]
	 assert algo.reverse_table(zero_list_fixture) == [0,0,0,0,0,0,0,0,0,0]
	 assert algo.reverse_table(mix_list_fixture) == [None,0,0,0,0,0,"string",0,'a']
# END REVERSE TABLE TESTING

# ROI BBOX
# Basic roi_bbox testing function
def test_roi_bbox (mat_fixture, mat_error_fixture, mat_little_fixture) :

	assert algo.roi_bbox(mat_fixture).all() == numpy.array([[2,2],[2,2],[2,2],[2,2]]).all()

	with pytest.raises(ValueError) :
		assert algo.roi_bbox(mat_error_fixture)

	with pytest.raises(IndexError) :
		assert algo.roi_bbox(mat_little_fixture)
# END ROI BBOX TESTING

# RANDOM ARRAY FILLING
# Basic random_fill_sparse testing function
def test_random_fill(array_empty, array_high, array_normal):

	# Testing with normal array defined in fixtures
	algo.random_fill_sparse(array_normal,6)

	count = 0
	rows_length = array_normal.shape[0]
	cols_length = array_normal.shape[1]

	for row in range(rows_length):
		for col in range(cols_length):
			if array_normal[row][col] == 'X':
				count += 1

	assert count == 6
	# End testing with normal array defined in fixtures

	with pytest.raises(ValueError):
		assert algo.random_fill_sparse(array_empty,5)

	with pytest.raises(ValueError):
		assert algo.random_fill_sparse(array_high,30)
# END RANDOM ARRAY FILLING TESTING

# REMOVE WHITESPACE
# Basic remove_whitespace testing function
def test_remove_whitespace(string_empty, string_unspaced, string_spaced):
	# Empty string testing
	assert algo.remove_whitespace(string_empty) == ""
	# Spaced string testing
	assert algo.remove_whitespace(string_spaced) == copy.deepcopy(string_spaced).replace(" ","")
	# Unspaced string testing
	assert algo.remove_whitespace(string_unspaced) == copy.deepcopy(string_unspaced).replace(" ","")
# END REMOVE WHITESPACE

# RANDOM ITEM SELECTION
# Basic shuffle testing function
def test_shuffle(list_empty, list_normal):
	# Empty list testing
	assert algo.shuffle(list_empty) == []
	# Normal list testing
	assert len(set(algo.shuffle(list_normal)).intersection(copy.deepcopy(list_normal))) == len(algo.shuffle(list_normal))
		#listCopy = copy.deepcopy(list_normal)
		#list_normal = algo.shuffle(list_normal)
		#assert len(set(list_normal).intersection(listCopy)) == len(list_normal)
# END RANDOM ITEM SELECTION TESTING

# SORT SELECTIVE BUBBLE
# Basic sort_selective testing function
def test_sort_selective(list_empty, positive_list_fixture, int_fixture, negative_list_fixture, zero_list_fixture, mix_list_fixture):
	# Empty list testing
	assert algo.sort_selective(list_empty) == []
	# Positive values list testing
	assert algo.sort_selective(positive_list_fixture) == sorted(copy.deepcopy(positive_list_fixture))
	# Positive and negative values list testing
	assert algo.sort_selective(int_fixture) == sorted(copy.deepcopy(int_fixture))
	# Negative list testing
	assert algo.sort_selective(negative_list_fixture) == sorted(copy.deepcopy(negative_list_fixture))
	# Zeros list testing
	assert algo.sort_selective(zero_list_fixture) == sorted(copy.deepcopy(zero_list_fixture))
	# Mixing type list testing
	assert algo.sort_selective(mix_list_fixture) == sorted(copy.deepcopy(mix_list_fixture))
# END SORT SELECTIVE BUBBLE

# SORT BUBBLE
# Basic sort_bubble testing function
def test_sort_bubble(list_empty, positive_list_fixture, int_fixture, negative_list_fixture, zero_list_fixture, mix_list_fixture):
	# Empty list testing
	assert algo.sort_bubble(list_empty) == []
	# Positive values list testing
	assert algo.sort_bubble(positive_list_fixture) == sorted(copy.deepcopy(positive_list_fixture))
	# Positive and negative values list testing
	assert algo.sort_bubble(int_fixture) == sorted(copy.deepcopy(int_fixture))
	# Negative list testing
	assert algo.sort_bubble(negative_list_fixture) == sorted(copy.deepcopy(negative_list_fixture))
	# Zeros list testing
	assert algo.sort_bubble(zero_list_fixture) == sorted(copy.deepcopy(zero_list_fixture))
	# Mixing type list testing
	assert algo.sort_bubble(mix_list_fixture) == sorted(copy.deepcopy(mix_list_fixture))
# END SORT BUBBLE


"""
    lantz.drivers.keysight.e8364b
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    Drivers for the E8364B Network Analyzer using RAW SOCKETS

    Authors: Alexandre Bourassa
    Date: 24/03/2016
"""
import numpy as _np

from lantz import Feat, DictFeat, Action
from lantz.feat import MISSING
from lantz.errors import InstrumentError
from lantz.messagebased import MessageBasedDriver

class E8364B(MessageBasedDriver):
    """E8364B Network Analyzer
    """

    DEFAULTS = {'COMMON': {'write_termination': '\n',
                           'read_termination': '\n',
                           'timeout':10000,}}

    @Feat(read_once=True)
    def idn(self):
        return self.query('*IDN?')

# ----------------------------------------------------
#       Sweep Settings functions
# ----------------------------------------------------

    @Feat(units='Hz')
    def center_frequency(self):
        """Center Frequency
        """
        return float(self.query('SENS:FREQ:CENT?'))

    @center_frequency.setter
    def center_frequency(self, value):
        self.write('SENS:FREQ:CENT {}'.format(int(value)))

    @Feat(units='Hz')
    def span(self):
        """Span of the sweep
        """
        return float(self.query('SENS:FREQ:SPAN?'))

    @span.setter
    def span(self, value):
        self.write('SENS:FREQ:SPAN {}'.format(int(value)))

    @Feat(units='Hz')
    def if_bandwidth(self):
        """Bandwidth of the digital IF filter
        """
        return float(self.query('SENS:BAND?'))

    @if_bandwidth.setter
    def if_bandwidth(self, value):
        self.write('SENS:BAND {}'.format(int(value)))

    @Feat(limits=(1, 16001, 1))
    def nb_points(self):
        """The number of data points for the measurement
        """
        return int(self.query('SENS:SWE:POIN?'))

    @nb_points.setter
    def nb_points(self, points):
        self.write('SENS:SWE:POIN {}'.format(int(points)))

    ALLOWED_MEAS_TYPE = {"S11", "S12", "S21", "S22"}
    @Feat(values=ALLOWED_MEAS_TYPE)
    def meas_type(self):
        return self.get_measurement_catalog()['CH1_S11_1']

    @meas_type.setter
    def meas_type(self, meas_type):
        self.write("CALC:PAR:MOD {}".format(meas_type))

# ----------------------------------------------------
#       Power functions
# ----------------------------------------------------

    @Feat(values={True: 1, False: 0})
    def power_on(self):
        """RF Power State (True==ON or False==OFF)
        """
        return int(self.query('OUTP?'))

    @power_on.setter
    def power_on(self, state=True):
        self.write('OUTP {}'.format(state))

    @Feat()
    def power_level(self):
        """RF Power level in dBm
        """
        return float(self.query("SOUR:POW?"))

    @power_level.setter
    def power_level(self, level):
        self.write("SOUR:POW {}".format(level))

    ## Possibly not available on this model...
    # @Feat(values={True: 1, False: 0})
    # def manual_noise_on(self):
    #     """RF Power State (True==ON or False==OFF)
    #     """
    #     return int(self.query('OUTP:MAN:NOIS?'))
    #
    # @manual_noise_on.setter
    # def manual_noise_on(self, state=True):
    #     if state:   state_str = 'ON'
    #     else:       state_str = 'OFF'
    #     self.write('OUTP:MAN:NOIS ' + state_str)

#----------------------------------------------------
#       Averaging functions
#----------------------------------------------------

    @Feat(values={True: 1, False: 0})
    def average_on(self):
        """Averaging state (True==ON or False==OFF)
        """
        return int(self.query('SENS:AVER:STAT?'))

    @average_on.setter
    def average_on(self, state=True):
        self.write('SENS:AVER:STAT {}'.format(state))

    @Feat(limits=(1,65536, 1))
    def average_count(self):
        """Averaging count
        """
        return int(self.query('SENS:AVER:COUN?'))

    @average_count.setter
    def average_count(self, counts):
        self.write('SENS:AVER:COUN {}'.format(int(counts)))

    @Action()
    def clear_average(self):
        """Reset the averaging
        """
        self.write('SENS:AVER:CLE')

#----------------------------------------------------
#       Data Query functions
#----------------------------------------------------
    @Feat(values={"REAL32":"REAL,+32", "REAL64":"REAL,+64", "ASCII":"ASC,+0"})
    def data_format(self):
        return self.query("FORM:DATA?")

    @data_format.setter
    def data_format(self, value):
        self.write('FORM:DATA {}'.format(value))

    def query_data(self, command, use_cached=True):
        # For quick data acquisition, let's assume data format as not changed since last query
        form = self.recall('data_format') if use_cached else self.data_format
        if form is MISSING: return _np.array([])
        if form == "REAL64":
            return _np.array(self.resource.query_binary_values(command, datatype='d', is_big_endian=True))
        elif form == "REAL32":
            return _np.array(self.resource.query_binary_values(command, datatype='f', is_big_endian=True))
        elif form == "ASCII":
            return _np.array(list(map(float, self.query(command).split(','))))
        else:
            raise Exception(str(form) + "Invalid data format")

    @Action()
    def x_data(self):
        return self.query_data('SENS:X?')

    @Action()
    def y_data(self):
        data = self.query_data("CALC:DATA? SDATA")
        data = data[::2]+data[1::2]*1j
        return data

# ----------------------------------------------------
#       Traces and Channel functions (Complicates the remote use of the device
#       so don't use unless you know what you are doing...)
#       For simplicity, we will always use CH1 and parameter name 'CH1_S11_1'
# ----------------------------------------------------
    def clear_all_traces(self):
        self.write('SYST:FPR')
        self.write('DISP:WIND:STAT ON')

    def create_new_measurement(self, name='CH1_S11_1', meas_type="S11"):
        if not meas_type in self.ALLOWED_MEAS_TYPE: raise Exception("Invalid meas_type: "+str(meas_type))
        self.write("CALC:PAR:DEF '"+name+"',"+meas_type)
        self.write("DISP:WIND:TRAC:FEED '"+name+"'")

    def get_measurement_catalog(self):
        response = self.query("CALC:PAR:CAT?").translate({ord('"'):None}).split(',')
        ans = dict()
        for i in range(int(len(response)/2)):
            ans[response[2*i]]=response[2*i+1]
        return ans

    def select_measurement(self, name):
        self.write("CALC:PAR:SEL '{}'".format(name))

# ----------------------------------------------------
#       Misc. functions
# ----------------------------------------------------

    @Action()
    def initialize(self):
        super(E8364B, self).initialize()
        meas = self.get_measurement_catalog()
        if not 'CH1_S11_1' in meas:
            self.clear_all_traces()
            self.create_new_measurement(name='CH1_S11_1', meas_type="S11")

        self.select_measurement('CH1_S11_1')
        self.data_format = 'REAL64'


"""
Points

TESTS::

    sage: E = EllipticCurve('37a')
    sage: P = E(0,0)
    sage: def get_points(n): return sum([point(list(i*P)[:2], size=3) for i in range(-n,n) if i != 0 and (i*P)[0] < 3])
    sage: sum([get_points(15*n).plot3d(z=n) for n in range(1,10)])
    Graphics3d Object
"""

#*****************************************************************************
#       Copyright (C) 2006 Alex Clemesha <clemesha@gmail.com>,
#                          William Stein <wstein@gmail.com>,
#                     2008 Mike Hansen <mhansen@gmail.com>,
#
#  Distributed under the terms of the GNU General Public License (GPL)
#
#    This code is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#    General Public License for more details.
#
#  The full text of the GPL is available at:
#
#                  http://www.gnu.org/licenses/
#*****************************************************************************
from sage.misc.decorators import options, rename_keyword
from sage.plot.colors import to_mpl_color
from sage.plot.primitive import GraphicPrimitive_xydata
import collections

# TODO: create _allowed_options for 3D point classes to
# improve bad option handling in plot3d?
class Point(GraphicPrimitive_xydata):
    """
    Primitive class for the point graphics type.  See point?, point2d?
    or point3d? for information about actually plotting points.

    INPUT:

    - xdata - list of x values for points in Point object

    - ydata - list of y values for points in Point object

    - options - dict of valid plot options to pass to constructor

    EXAMPLES:

    Note this should normally be used indirectly via ``point`` and friends::

        sage: from sage.plot.point import Point
        sage: P = Point([1,2],[2,3],{'alpha':.5})
        sage: P
        Point set defined by 2 point(s)
        sage: P.options()['alpha']
        0.500000000000000
        sage: P.xdata
        [1, 2]

    TESTS:

    We test creating a point::

        sage: point((3,3))
        Graphics object consisting of 1 graphics primitive
    """
    def __init__(self, xdata, ydata, options):
        """
        Initializes base class Point.

        EXAMPLES::

            sage: P = point((3,4))
            sage: P[0].xdata
            [3.0]
            sage: P[0].options()['alpha']
            1
        """
        self.xdata = xdata
        self.ydata = ydata
        GraphicPrimitive_xydata.__init__(self, options)

    def _allowed_options(self):
        """
        Return the allowed options for the Point class.

        EXAMPLES::

            sage: P = point((3,4))
            sage: P[0]._allowed_options()['size']
            'How big the point is (i.e., area in points^2=(1/72 inch)^2).'
        """
        return {'alpha':'How transparent the point is.',
                'faceted': 'If True color the edge of the point. (only for 2D plots)',
                'hue':'The color given as a hue.',
                'legend_color':'The color of the legend text',
                'legend_label':'The label for this item in the legend.',
                'marker':'the marker symbol for 2D plots only (see documentation of plot() for details)',
                'markeredgecolor':'the color of the marker edge (only for 2D plots)',
                'rgbcolor':'The color as an RGB tuple.',
                'size': 'How big the point is (i.e., area in points^2=(1/72 inch)^2).',
                'zorder':'The layer level in which to draw'}

    def _plot3d_options(self, options=None):
        """
        Translate 2D plot options into 3D plot options.

        EXAMPLES::

            sage: A=point((1,1),size=22)
            sage: a=A[0];a
            Point set defined by 1 point(s)
            sage: b=a.plot3d()
            sage: b.size
            22
            sage: b=a.plot3d(size=3)
            sage: b.size
            3
        """
        if options is None:
            options = dict(self.options())
        options_3d = {}
        if 'size' in options:
            options_3d['size'] = options['size']
            del options['size']
        if options.pop('faceted', False):
            raise NotImplementedError("3D points can not be faceted.")
        for o in ('marker', 'markeredgecolor'): # remove 2D options
            if o in options:
                del options[o]

        options_3d.update(GraphicPrimitive_xydata._plot3d_options(self, options))
        return options_3d

    def plot3d(self, z=0, **kwds):
        """
        Plots a two-dimensional point in 3-D, with default height zero.

        INPUT:

        -  ``z`` - optional 3D height above `xy`-plane.  May be a list
           if self is a list of points.

        EXAMPLES:

        One point::

            sage: A=point((1,1))
            sage: a=A[0];a
            Point set defined by 1 point(s)
            sage: b=a.plot3d()

        One point with a height::

            sage: A=point((1,1))
            sage: a=A[0];a
            Point set defined by 1 point(s)
            sage: b=a.plot3d(z=3)
            sage: b.loc[2]
            3.0

        Multiple points::

            sage: P=point([(0,0), (1,1)])
            sage: p=P[0]; p
            Point set defined by 2 point(s)
            sage: q=p.plot3d(size=22)

        Multiple points with different heights::

            sage: P=point([(0,0), (1,1)])
            sage: p=P[0]
            sage: q=p.plot3d(z=[2,3])
            sage: q.all[0].loc[2]
            2.0
            sage: q.all[1].loc[2]
            3.0

        Note that keywords passed must be valid point3d options::

            sage: A=point((1,1),size=22)
            sage: a=A[0];a
            Point set defined by 1 point(s)
            sage: b=a.plot3d()
            sage: b.size
            22
            sage: b=a.plot3d(pointsize=23) # only 2D valid option
            sage: b.size
            22
            sage: b=a.plot3d(size=23) # correct keyword
            sage: b.size
            23

        TESTS:

        Heights passed as a list should have same length as
        number of points::

            sage: P=point([(0,0), (1,1), (2,3)])
            sage: p=P[0]
            sage: q=p.plot3d(z=2)
            sage: q.all[1].loc[2]
            2.0
            sage: q=p.plot3d(z=[2,-2])
            Traceback (most recent call last):
            ...
            ValueError: Incorrect number of heights given
        """
        from sage.plot.plot3d.base import Graphics3dGroup
        from sage.plot.plot3d.shapes2 import point3d
        options = self._plot3d_options()
        options.update(kwds)
        zdata=[]
        if isinstance(z, list):
            zdata=z
        else:
            zdata=[z]*len(self.xdata)
        if len(zdata)==len(self.xdata):
            all = [point3d([(x, y, z) for x, y, z in zip(self.xdata, self.ydata, zdata)], **options)]
            if len(all) == 1:
                return all[0]
            else:
                return Graphics3dGroup(all)
        else:
            raise ValueError('Incorrect number of heights given')

    def _repr_(self):
        """
        String representation of Point primitive.

        EXAMPLES::

            sage: P=point([(0,0), (1,1)])
            sage: p=P[0]; p
            Point set defined by 2 point(s)
        """
        return "Point set defined by %s point(s)"%len(self.xdata)

    def __getitem__(self, i):
        """
        Returns tuple of coordinates of point.

        EXAMPLES::

            sage: P=point([(0,0), (1,1), (2,3)])
            sage: p=P[0]; p
            Point set defined by 3 point(s)
            sage: p[1]
            (1.0, 1.0)
        """
        return self.xdata[i], self.ydata[i]

    def _render_on_subplot(self,subplot):
        r"""
        TESTS:

        We check to make sure that :trac:`2076` is fixed by verifying all
        the points are red::

            sage: point(((1,1), (2,2), (3,3)), rgbcolor=hue(1), size=30)
            Graphics object consisting of 1 graphics primitive
        """
        options = self.options()

        #Convert the color to a hex string so that the scatter
        #method does not interpret it as a list of 3 floating
        #point color specifications when there are
        #three points. This is mentioned in the matplotlib 0.98
        #documentation and fixes \#2076
        from matplotlib.colors import rgb2hex
        c = rgb2hex(to_mpl_color(options['rgbcolor']))

        a = float(options['alpha'])
        z = int(options.pop('zorder', 0))
        s = int(options['size'])
        faceted = options['faceted'] #faceted=True colors the edge of point
        markeredgecolor = options['markeredgecolor']

        scatteroptions={}
        if not faceted and markeredgecolor is None:
            scatteroptions['edgecolors'] = 'none'
        elif markeredgecolor is not None:
            scatteroptions['edgecolors'] = to_mpl_color(
                                              options.pop('markeredgecolor'))
        scatteroptions['marker'] = options.pop('marker')

        subplot.scatter(self.xdata, self.ydata, s=s, c=c, alpha=a, zorder=z,
                        label=options['legend_label'], **scatteroptions)

def point(points, **kwds):
    """
    Returns either a 2-dimensional or 3-dimensional point or sum of points.

    INPUT:

    -  ``points`` - either a single point (as a tuple), a list of
       points, a single complex number, or a list of complex numbers.

    For information regarding additional arguments, see either point2d?
    or point3d?.

    .. SEEALSO::

        :func:`sage.plot.point.point2d`, :func:`sage.plot.plot3d.shapes2.point3d`

    EXAMPLES::

        sage: point((1,2))
        Graphics object consisting of 1 graphics primitive

    ::

        sage: point((1,2,3))
        Graphics3d Object

    ::

        sage: point([(0,0), (1,1)])
        Graphics object consisting of 1 graphics primitive

    ::

        sage: point([(0,0,1), (1,1,1)])
        Graphics3d Object

    Extra options will get passed on to show(), as long as they are valid::

        sage: point([(cos(theta), sin(theta)) for theta in srange(0, 2*pi, pi/8)], frame=True)
        Graphics object consisting of 1 graphics primitive
        sage: point([(cos(theta), sin(theta)) for theta in srange(0, 2*pi, pi/8)]).show(frame=True) # These are equivalent

    TESTS:

    One can now use iterators (:trac:`13890`)::

        sage: point(iter([(1,1,1)]))
        Graphics3d Object
        sage: point(iter([(1,2),(3,5)]))
        Graphics object consisting of 1 graphics primitive
    """
    if isinstance(points, collections.Iterator):
        points = list(points)

    try:
        return point2d(points, **kwds)
    except (ValueError, TypeError):
        from sage.plot.plot3d.shapes2 import point3d
        return point3d(points, **kwds)

@rename_keyword(color='rgbcolor', pointsize='size')
@options(alpha=1, aspect_ratio='automatic', faceted=False,
        legend_color=None, legend_label=None, marker='o',
        markeredgecolor=None, rgbcolor=(0,0,1), size=10)
def point2d(points, **options):
    r"""
    A point of size ``size`` defined by point = `(x,y)`.

    INPUT:

    -  ``points`` - either a single point (as a tuple), a list of
       points, a single complex number, or a list of complex numbers.
    - ``alpha`` -- How transparent the point is.
    - ``faceted`` -- If True color the edge of the point. (only for 2D plots)
    - ``hue`` -- The color given as a hue.
    - ``legend_color`` -- The color of the legend text
    - ``legend_label`` -- The label for this item in the legend.
    - ``marker`` -- the marker symbol for 2D plots only (see documentation of
      :func:`plot` for details)
    - ``markeredgecolor`` -- the color of the marker edge (only for 2D plots)
    - ``rgbcolor`` -- The color as an RGB tuple.
    - ``size`` -- How big the point is (i.e., area in points^2=(1/72 inch)^2).
    - ``zorder`` -- The layer level in which to draw

    EXAMPLES:

    A purple point from a single tuple or coordinates::

        sage: point((0.5, 0.5), rgbcolor=hue(0.75))
        Graphics object consisting of 1 graphics primitive

    Points with customized markers and edge colors::

        sage: r = [(random(), random()) for _ in range(10)]
        sage: point(r, marker='d', markeredgecolor='red', size=20)
        Graphics object consisting of 1 graphics primitive

    Passing an empty list returns an empty plot::

        sage: point([])
        Graphics object consisting of 0 graphics primitives
        sage: import numpy; point(numpy.array([]))
        Graphics object consisting of 0 graphics primitives

    If you need a 2D point to live in 3-space later, this is possible::

        sage: A=point((1,1))
        sage: a=A[0];a
        Point set defined by 1 point(s)
        sage: b=a.plot3d(z=3)

    This is also true with multiple points::

        sage: P=point([(0,0), (1,1)])
        sage: p=P[0]
        sage: q=p.plot3d(z=[2,3])

    Here are some random larger red points, given as a list of tuples::

        sage: point(((0.5, 0.5), (1, 2), (0.5, 0.9), (-1, -1)), rgbcolor=hue(1), size=30)
        Graphics object consisting of 1 graphics primitive

    And an example with a legend::

        sage: point((0,0), rgbcolor='black', pointsize=40, legend_label='origin')
        Graphics object consisting of 1 graphics primitive

    The legend can be colored::

        sage: P = points([(0,0),(1,0)], pointsize=40, legend_label='origin', legend_color='red')
        sage: P + plot(x^2,(x,0,1), legend_label='plot', legend_color='green')
        Graphics object consisting of 2 graphics primitives

    Extra options will get passed on to show(), as long as they are valid::

        sage: point([(cos(theta), sin(theta)) for theta in srange(0, 2*pi, pi/8)], frame=True)
        Graphics object consisting of 1 graphics primitive
        sage: point([(cos(theta), sin(theta)) for theta in srange(0, 2*pi, pi/8)]).show(frame=True) # These are equivalent

    For plotting data, we can use a logarithmic scale, as long as we are sure
    not to include any nonpositive points in the logarithmic direction::

        sage: point([(1,2),(2,4),(3,4),(4,8),(4.5,32)],scale='semilogy',base=2)
        Graphics object consisting of 1 graphics primitive

    Since Sage Version 4.4 (:trac:`8599`), the size of a 2d point can be
    given by the argument ``size`` instead of ``pointsize``. The argument
    ``pointsize`` is still supported::

        sage: point((3,4), size=100)
        Graphics object consisting of 1 graphics primitive

    ::

        sage: point((3,4), pointsize=100)
        Graphics object consisting of 1 graphics primitive

    We can plot a single complex number::

        sage: point(CC(1+I), pointsize=100)
        Graphics object consisting of 1 graphics primitive

    We can also plot a list of complex numbers::

        sage: point([CC(I), CC(I+1), CC(2+2*I)], pointsize=100)
        Graphics object consisting of 1 graphics primitive

    """
    from sage.plot.plot import xydata_from_point_list
    from sage.plot.all import Graphics
    from sage.rings.all import CC, CDF
    if points in CC or points in CDF:
        pass
    else:
        try:
            l = len(points)
        except TypeError:
            # argument is an iterator
            points = list(points)
            l = len(points)

        if l == 0:
            return Graphics()

    xdata, ydata = xydata_from_point_list(points)
    g = Graphics()
    g._set_extra_kwds(Graphics._extract_kwds_for_show(options))
    g.add_primitive(Point(xdata, ydata, options))
    if options['legend_label']:
        g.legend(True)
        g._legend_colors = [options['legend_color']]
    return g

points = point


#!/usr/bin/env python
# coding: utf-8

# # This is a TF Estimator end-to-end baseline solution
#
# **For local run**
#
# Tested with
#
# ```
# numpy==1.13.3
# scipy==0.19.1
# tensorflow-gpu==1.4.0
# tqdm
# ```
#
#
# I want to show usage of Estimators with custom python datagenerators.
#
#
# Detailed documentation you can find at https://www.tensorflow.org/api_docs/python/tf/estimator/Estimator
#
# I also recommend to read source code  https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/estimator/estimator.py

# Suppose we have following project structure:
# ```
# .
# ├── data
# │   ├── test            # extracted
# │   │   └── audio          # all test
# │   ├── test.7z         # downloaded
# │   ├── train           # extracted
# │   │   ├── audio          # folder with all train command/file.wav
# │   │   ├── LICENSE
# │   │   ├── README.md
# │   │   ├── testing_list.txt
# │   │   └── validation_list.txt
# │   └── train.7z         # downloaded
# ├── kernel.ipynb      # this ipynb
# └── model-k           # folder for model, checkpoints, logs and submission.csv
# ```

# In[1]:

DATADIR = './data' # unzipped train and test data
OUTDIR = './model-k' # just a random name
# Data Loading
import os
import re
from glob import glob

POSSIBLE_LABELS = 'yes no up down left right on off stop go silence unknown'.split()
id2name = {i: name for i, name in enumerate(POSSIBLE_LABELS)}
name2id = {name: i for i, name in id2name.items()}

def load_data(data_dir):
    """ Return 2 lists of tuples:
    [(class_id, user_id, path), ...] for train
    [(class_id, user_id, path), ...] for validation
    """
    # Just a simple regexp for paths with three groups:
    # prefix, label, user_id
    pattern = re.compile("(.+\/)?(\w+)\/([^_]+)_.+wav")
    all_files = glob(os.path.join(data_dir, 'train/audio/*/*wav'))

    with open(os.path.join(data_dir, 'train/validation_list.txt'), 'r') as fin:
        validation_files = fin.readlines()
    valset = set()
    for entry in validation_files:
        r = re.match(pattern, entry)
        if r:
            valset.add(r.group(3))

    possible = set(POSSIBLE_LABELS)
    train, val = [], []
    for entry in all_files:
        r = re.match(pattern, entry)
        if r:
            label, uid = r.group(2), r.group(3)
            if label == '_background_noise_':
                label = 'silence'
            if label not in possible:
                label = 'unknown'

            label_id = name2id[label]

            sample = (label_id, uid, entry)
            if uid in valset:
                val.append(sample)
            else:
                train.append(sample)

    print('There are {} train and {} val samples'.format(len(train), len(val)))
    return train, val

trainset, valset = load_data(DATADIR)

# Let me introduce pythonic datagenerator.
# It is just a python/numpy/... function **without tf** that yields dicts such that
# ```
# {
#   'x': np.array(...),
#   'str_key': np.string_(...),
#   'label': np.int32(...),
# }
# ```
#
# Be sure, every value in this dict has `.dtype` method.

# In[2]:

import numpy as np
from scipy.io import wavfile

def data_generator(data, params, mode='train'):
    def generator():
        if mode == 'train':
            np.random.shuffle(data)
        # Feel free to add any augmentation
        for (label_id, uid, fname) in data:
            try:
                _, wav = wavfile.read(fname)
                wav = wav.astype(np.float32) / np.iinfo(np.int16).max

                L = 16000  # be aware, some files are shorter than 1 sec!
                if len(wav) < L:
                    continue
                # let's generate more silence!
                samples_per_file = 1 if label_id != name2id['silence'] else 20
                for _ in range(samples_per_file):
                    if len(wav) > L:
                        beg = np.random.randint(0, len(wav) - L)
                    else:
                        beg = 0
                    yield dict(
                        target=np.int32(label_id),
                        wav=wav[beg: beg + L],
                    )
            except Exception as err:
                print(err, label_id, uid, fname)

    return generator

#
# Suppose, we have spectrograms and want to write feature extractor that produces logits.
#
#
# Let's write some simple net, treat sound as a picture.
#
#
# **Spectrograms** (input x) have shape `(batch_size, time_frames, freq_bins, 2)`.
#
# **Logits** is a tensor with shape `(batch_size, num_classes)`.

# In[3]:

import tensorflow as tf
from tensorflow.contrib import layers

def baseline(x, params, is_training):
    x = layers.batch_norm(x, is_training=is_training)
    for i in range(4):
        x = layers.conv2d(
            x, 16 * (2 ** i), 3, 1,
            activation_fn=tf.nn.elu,
            normalizer_fn=layers.batch_norm if params.use_batch_norm else None,
            normalizer_params={'is_training': is_training}
        )
        x = layers.max_pool2d(x, 2, 2)

    # just take two kind of pooling and then mix them, why not :)
    mpool = tf.reduce_max(x, axis=[1, 2], keep_dims=True)
    apool = tf.reduce_mean(x, axis=[1, 2], keep_dims=True)

    x = 0.5 * (mpool + apool)
    # we can use conv2d 1x1 instead of dense
    x = layers.conv2d(x, 128, 1, 1, activation_fn=tf.nn.elu)
    x = tf.nn.dropout(x, keep_prob=params.keep_prob if is_training else 1.0)

    # again conv2d 1x1 instead of dense layer
    logits = layers.conv2d(x, params.num_classes, 1, 1, activation_fn=None)
    return tf.squeeze(logits, [1, 2])

# We need to write a model handler for three regimes:
# - train
# - eval
# - predict
#
# Loss function, train_op, additional metrics and summaries should be defined.
#
# Also, we need to convert sound waveform into spectrograms (we could do it with numpy/scipy/librosa in data generator, but TF has new signal processing API)

# In[4]:

from tensorflow.contrib import signal

# features is a dict with keys: tensors from our datagenerator
# labels also were in features, but excluded in generator_input_fn by target_key

def model_handler(features, labels, mode, params, config):
    # Im really like to use make_template instead of variable_scopes and re-usage
    extractor = tf.make_template(
        'extractor', baseline,
        create_scope_now_=True,
    )
    # wav is a waveform signal with shape (16000, )
    wav = features['wav']
    # we want to compute spectograms by means of short time fourier transform:
    specgram = signal.stft(
        wav,
        400,  # 16000 [samples per second] * 0.025 [s] -- default stft window frame
        160,  # 16000 * 0.010 -- default stride
    )
    # specgram is a complex tensor, so split it into abs and phase parts:
    phase = tf.angle(specgram) / np.pi
    # log(1 + abs) is a default transformation for energy units
    amp = tf.log1p(tf.abs(specgram))

    x = tf.stack([amp, phase], axis=3) # shape is [bs, time, freq_bins, 2]
    x = tf.to_float(x)  # we want to have float32, not float64

    logits = extractor(x, params, mode == tf.estimator.ModeKeys.TRAIN)

    if mode == tf.estimator.ModeKeys.TRAIN:
        loss = tf.reduce_mean(
            tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits))
        # some lr tuner, you could use move interesting functions
        def learning_rate_decay_fn(learning_rate, global_step):
            return tf.train.exponential_decay(
                learning_rate, global_step, decay_steps=10000, decay_rate=0.99)

        train_op = tf.contrib.layers.optimize_loss(
            loss=loss,
            global_step=tf.contrib.framework.get_global_step(),
            learning_rate=params.learning_rate,
            optimizer=lambda lr: tf.train.MomentumOptimizer(lr, 0.9, use_nesterov=True),
            learning_rate_decay_fn=learning_rate_decay_fn,
            clip_gradients=params.clip_gradients,
            variables=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES))

        specs = dict(
            mode=mode,
            loss=loss,
            train_op=train_op,
        )

    if mode == tf.estimator.ModeKeys.EVAL:
        prediction = tf.argmax(logits, axis=-1)
        acc, acc_op = tf.metrics.mean_per_class_accuracy(
            labels, prediction, params.num_classes)
        loss = tf.reduce_mean(
            tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels, logits=logits))
        specs = dict(
            mode=mode,
            loss=loss,
            eval_metric_ops=dict(
                acc=(acc, acc_op),
            )
        )

    if mode == tf.estimator.ModeKeys.PREDICT:
        predictions = {
            'label': tf.argmax(logits, axis=-1),  # for probability just take tf.nn.softmax()
            'sample': features['sample'], # it's a hack for simplicity
        }
        specs = dict(
            mode=mode,
            predictions=predictions,
        )
    return tf.estimator.EstimatorSpec(**specs)

def create_model(config=None, hparams=None):
    return tf.estimator.Estimator(
        model_fn=model_handler,
        config=config,
        params=hparams,
    )

# Define some params. Move model hyperparams (optimizer, extractor, num of layers, activation fn, ...) here

# In[5]:

params=dict(
    seed=2018,
    batch_size=64,
    keep_prob=0.5,
    learning_rate=1e-3,
    clip_gradients=15.0,
    use_batch_norm=True,
    num_classes=len(POSSIBLE_LABELS),
)

hparams = tf.contrib.training.HParams(**params)
os.makedirs(os.path.join(OUTDIR, 'eval'), exist_ok=True)
model_dir = OUTDIR

run_config = tf.contrib.learn.RunConfig(model_dir=model_dir)

# **Let's run training!**

# In[6]:

# it's a magic function :)
from tensorflow.contrib.learn.python.learn.learn_io.generator_io import generator_input_fn

train_input_fn = generator_input_fn(
    x=data_generator(trainset, hparams, 'train'),
    target_key='target',  # you could leave target_key in features, so labels in model_handler will be empty
    batch_size=hparams.batch_size, shuffle=True, num_epochs=None,
    queue_capacity=3 * hparams.batch_size + 10, num_threads=1,
)

val_input_fn = generator_input_fn(
    x=data_generator(valset, hparams, 'val'),
    target_key='target',
    batch_size=hparams.batch_size, shuffle=True, num_epochs=None,
    queue_capacity=3 * hparams.batch_size + 10, num_threads=1,
)

def _create_my_experiment(run_config, hparams):
    exp = tf.contrib.learn.Experiment(
        estimator=create_model(config=run_config, hparams=hparams),
        train_input_fn=train_input_fn,
        eval_input_fn=val_input_fn,
        train_steps=10000, # just randomly selected params
        eval_steps=200,  # read source code for steps-epochs ariphmetics
        train_steps_per_iteration=1000,
    )
    return exp

tf.contrib.learn.learn_runner.run(
    experiment_fn=_create_my_experiment,
    run_config=run_config,
    schedule="continuous_train_and_eval",
    hparams=hparams)

#
# While it trains (~10-20min on i5 + 1080), you could start tensorboard on model_dir and see live chart like this
#
# ![Tensorboard](https://pp.userapi.com/c841329/v841329524/3db60/fdNDyRMJHMQ.jpg)
#
#
# Now we want to predict testset and make submission file.
#
# 1. Create datagenerator and input_function
# 2. Load model
# 3. Iterate over predictions and store results

# In[7]:

from tqdm import tqdm
# now we want to predict!
paths = glob(os.path.join(DATADIR, 'test/audio/*wav'))

def test_data_generator(data):
    def generator():
        for path in data:
            _, wav = wavfile.read(path)
            wav = wav.astype(np.float32) / np.iinfo(np.int16).max
            fname = os.path.basename(path)
            yield dict(
                sample=np.string_(fname),
                wav=wav,
            )

    return generator

test_input_fn = generator_input_fn(
    x=test_data_generator(paths),
    batch_size=hparams.batch_size,
    shuffle=False,
    num_epochs=1,
    queue_capacity= 10 * hparams.batch_size,
    num_threads=1,
)

model = create_model(config=run_config, hparams=hparams)
it = model.predict(input_fn=test_input_fn)

# last batch will contain padding, so remove duplicates
submission = dict()
for t in tqdm(it):
    fname, label = t['sample'].decode(), id2name[t['label']]
    submission[fname] = label

with open(os.path.join(model_dir, 'submission.csv'), 'w') as fout:
    fout.write('fname,label\n')
    for fname, label in submission.items():
        fout.write('{},{}\n'.format(fname, label))

# ## About tf.Estimators
#
# **Pros**:
# - no need to control Session
# - datagenerator feeds model via queues without explicit queue coding :)
# - you could naturaly export models into production
#
# **Cons**:
# - it's very hard to debug computational graph (use `tf.add_check_numerics()` and `tf.Print` in case of problems)
# - boilerplate code
# - need to read source code for making interesting things
#
#
# **Conclusion**:
# Estimator is a nice abstraction with some boilerplate code :)
#
#
# ## About Speech Recognition Challenge:
#
# You could start from this end-to-end ipynb, improving several functions for much better results.
#
#
#
# May the gradient flow be with you.

# In[8]:


#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  ALT_library.py
#
#  Copyright 2020: Profesores de Algorítmica (UPV, GII/ETSINF)

import collections
import numpy as np

class Trie:
    """
    Clase que implementa un Trie para almacenar el vocabulario
    """

    def __init__(self, vocabulary):
        """Método constructor de la clase Trie.

        Llama al método build_trie para crear el trie a partir de un vocabulario de términos.

        Args:
            vocabulary (list of str): Vocabulario de términos.
                Ejemplo: ["aa","ac","bb","c","cab","cac", "aá", "aà"].
                Los términos deben estar ordenados lexicográficamente (sort() o sorted())

        Raises:
            Exception: si el vocabulario no es una lista de cadenas ordenada lexicográficamente.
        """

        if not (isinstance(vocabulary,list) and
                all(isinstance(w, str) and len(w) > 0 for w in vocabulary) and
                all(w1 < w2 for w1, w2 in zip(vocabulary, vocabulary[1:]))):
            raise Exception("vocabulario incorrecto")

        self.vocabulary = vocabulary # nos quedamos una referencia, esto podría dar problemas
                                     # si más adelante alguien cambia dicha lista, ¡cuidado!
                                     # una alternativa sería hacer copia local pero ocupa más espacio
                                     # si varios objetos van a compartir el vocabulario sin modificarlo
        self.build_trie() # construimos el Trie propiamente dicho

    def build_trie(self):
        """Método para construir el trie.

        Crea los atributos que representan el trie (trasparencias 34-38 del boletín de prácticas):

            label (np.array of unicode symbols): cada posición i del array contiene
                el símbolo/letra/carácter/unicode_rune asociado al nodo i del trie.

            firstchild (np.array of int): cada posición i del array contiene
                el índice del PRIMER hijo del nodo i del trie.
                La idea es que puedes recorrer todos los hijos de i haciendo un bucle
                en range(self.firstchild[i], self.firstchild[i+1])
                Se utiliza para programación dinámica hacia adelante (consultar a dónde vas)

            parent (np.array of int): cada posición i del array contiene
                el padre del nodo i. El nodo raiz y el centinela comparten otro nodo ficticio
                como padre (-1).
                Se utiliza para programación dinámica hacia atrás (consultar de dónde vienes)

            output es un diccionario que a los nodos finales les asocia la cadena que se emite
                la forma de saber si un nodo i es final es comprobar que i es una clave de este
                diccionario.
        """

        label = []     # lista python, después self.label un array numpy
        firstchild = [] # lista python, después self.firstchild un array numpy
        parent = []     # lista python, después self.parent array numpy
        self.output = {} # diccionario asocia a cada nodo final una palabra del vocabulario
        label.append(" ")   # asociado al nodo raíz
        firstchild.append(1) # asociado al nodo raíz
        parent.append(-1)   # asociado al nodo raíz

        # esto básicamente hace un recorrido por niveles del Trie, de
        # ahí que utilice una cola:
        Q = collections.deque()
        # cada elemento guardado en la cola es una tupla con estos campos:
        # first_index, last_index, triedepth, root = Q.popleft()
        # donde:
        # first_index y last_index: ambos referidos al vector self.vocabulary,
        #  el 2º apunta a después como range
        #  ambos delimitan la zona a tratar
        # triedepth se refiere a la prof. de root en el trie, donde
        # root es el estado raíz del subárbol que se va a generar
        Q.append((0, len(self.vocabulary), 1, 0))

        while len(Q) > 0:
            first_index, last_index, triedepth, root = Q.popleft()
            firstchild[root] = len(firstchild) # ver más abajo el (*)
            positions = []
            lastletter = ""
            for i in range(first_index, last_index):
                word = self.vocabulary[i]
                if len(word) >= triedepth:
                    letter = word[triedepth - 1]
                    if letter != lastletter:
                        lastletter = letter
                        label.append(letter)
                        parent.append(root)
                        firstchild.append(0) # se pone a 0 por poner algo, luego se modifica arriba en (*)
                        thischildpos = len(firstchild) - 1
                        if len(word) == triedepth:
                            self.output[thischildpos] = word
                        positions.append((i, thischildpos))
            positions.append((last_index, 0))
            for (frompos, root), (lastpos, dummy) in zip(positions, positions[1:]):
                Q.append((frompos, lastpos, triedepth + 1, root))

        # un nodo centinela al final para poder usar
        # range(self.firstchild[i], self.firstchild[i+1]) en el último nodo
        label.append(" ")
        firstchild.append(len(firstchild))
        parent.append(-1)
        # convertimos todo a np.array y lo guardamos como atributos de la clase:
        self.label = np.array(label, dtype=np.unicode_)
        self.firstchild = np.array(firstchild, dtype=np.int)
        self.parent = np.array(parent, dtype=np.int)

    def get_root(self):
        return 0

    def get_num_states(self):
        return len(self.firstchild)-1 # sentinel is not included

    def get_label(self, node):
        return self.label[node]

    def get_parent(self, node):
        return self.parent[node]

    def iter_children(self, node):
        return range(self.firstchild[node], self.firstchild[node + 1])

    def is_leaf(self, node):
        return self.firstchild[node] == self.firstchild[node+1]

    def is_final(self, node):
        return node in self.output

    def get_output(self, node):
        return self.output.get(node,"")

    def num_children(self, node):
        return self.firstchild[node + 1] - self.firstchild[node]

    def __str__(self):
        lines = ["vocabulary " + repr(self.vocabulary),
                 "pos label parent firstchild nºchild output"]
        num_nodes = len(self.label)-1 # sentinel is not included
        for node in range(num_nodes):
            etiqueta = '"'+str(self.label[node])+'"'
            num_children = self.firstchild[node + 1] - self.firstchild[node]
            out = '' if node not in self.output else '"'+self.output[node]+'"'
            line = f'{node:3} {etiqueta:4}  {self.parent[node]:3}    {self.firstchild[node]:3}        {num_children:2}      {out}'
            lines.append(line)
        return "\n".join(lines)

if __name__ == "__main__":
    vocabulario = ["aa","ac","bb","c","cab","cac"]
    trie = Trie(vocabulario)
    print(trie)
    print()
    vocabulario = ['a', 'ata', 'ato', 'cama', 'casa', 'caso', 'casó', 'caña']
    trie = Trie(vocabulario)
    print(trie)