import numpy as np
import pandas as pd
import sys
from copy import deepcopy
from scipy.sparse import coo_matrix
from .load_hic_file import get_chromosome_lengths, load_HiC
from ..embedding import pairwise_distances, MDS, tSNE, PHATE, SpectralEmbedding, PCA
from .processing_utils import matrix_operation
# from ..analysis import scatter
from ..analysis import kmeans, spectral_clustering, HAC
import matplotlib.pyplot as plt
try:
    import multiprocessing as mp
except:
    mp=None

def add_cell(ch, idx, file, resolution, chromosome_lengths,
             store_full_map, keep_n_strata, format, customized_format,
             header, adjust_resolution, map_filter, sparse, gzip, operations, sep=' ',):
    if ('ch' in ch) and ('chr' not in ch):
        ch=ch.replace("ch", "chr")
    mat, strata = load_HiC(
        file, genome_length=chromosome_lengths,
        format=format, custom_format=customized_format,
        header=header, chromosome=ch, resolution=resolution,
        resolution_adjust=adjust_resolution,
        map_filter=map_filter, sparse=sparse, gzip=gzip,
        keep_n_strata=keep_n_strata, operations=operations, sep=sep)

    contacts=np.sum(mat)/2+ np.trace(mat)/2

    short_range=sum([np.sum(mat[i,i:i+int(2000000/resolution)]) for i in range(len(mat))])
    mitotic=sum([np.sum(mat[i,i+int(2000000/resolution):i+int(12000000/resolution)]) for i in range(len(mat))])

    # if not keep_n_strata:
        # self.strata[ch][idx] = strata
        # for strata_idx, stratum in enumerate(strata):
        #     self.strata[ch][strata_idx][idx, :] = stratum
    if store_full_map:
        return [contacts, short_range, mitotic, strata, mat]
    else:
        return [contacts, short_range, mitotic, strata]
    # else:
    #     raise ValueError('`keep_n_strata` should be an positive intger.')

class scHiCs:
    def __init__(self, list_of_files, reference_genome, resolution,
                 adjust_resolution=True, sparse=False, chromosomes='all',
                 format='customized', keep_n_strata=10, store_full_map=False,
                 operations=None, header=0, customized_format=None,
                 map_filter=0., gzip=False, sep=' ',
                 parallelize=False, n_processes=None, **kwargs):
        """

        Parameters
        ----------
        list_of_files : list
            List of HiC file paths.

        reference_genome : str or dict
            Now supporting 'mm9', 'mm10', 'hg19', 'hg38',
            if using other references,you can simply provide the chromosome name
            and corresponding size (bp) with a dictionary in Python.
            e.g. {'chr1': 150000000, 'chr2': 130000000, 'chr3': 200000000}

        resolution : int
            The resolution to separate genome into bins. If using .hic file format,
            the given resolution must match with the resolutions in .hic file.

        adjust_resolution : bool, optional
            Whether to adjust resolution for input file.
            Sometimes the input file is already in the proper resolution
            (e.g. position 3000000 has already been changed to 6 in 500kb resolution),
            then you can set `adjust_resolution=False`.  The default is True.

        sparse : bool, optional
            Whether to use sparse matrix to store (only effective when max_distance=None). The default is False.

        chromosomes : list or str, optional
            Chromosomes to use,
            eg. ['chr1', 'chr2'], or just 'except Y', 'except XY','all',
            which means chr 1-19 + XY for mouse and chr 1-22 + XY for human.
            The default is 'all'.

        format : str, optional
            HiC files' format.
            e.g., '.hic', 'customized', '.cool'. The default is 'customized'.

        keep_n_strata : int, optional
            Only consider contacts within this genomic distance.
            If `None`, it will store full matrices in numpy matrix or
            scipy sparse format, which will use too much memory sometimes.
            The default is 10.

        store_full_map : bool, optional
            Whether store contact maps. The default is False.

        operations : list, optional
            The methods use for pre-processing or smoothing the maps given in a list.
            The operations will happen in the given order.
            Operations: 'convolution', 'random_walk', 'network_enhancing'.
            For pre-processing and smoothing operations, sometimes you need additional arguments.
            You can check docstrings for pre-processing and smoothing for more information.
            The default is None.

        header : int, optional
            The number of header line(s).
            If `header=0`, HiC files do not have header.
            The default is 0.

        customized_format : int or list, optional
            Format for each line. The default is None.

        map_filter : float, optional
            The threshold to filter some reads by map quality.
            The default is 0..

        gzip : bool, optional
            If the HiC files are zip files.
            If `True`, the HiC files are zip files.
            The default is False.

        parallelize : bool, optional
            If `True`, parallelize file reading process.
            The default is False.

        n_processes : int, optional
            Number of cores to use in parallelization.
            When n_processes=Null,  use number of CPUs -1 for parallelization.
            The default is Null.

        **kwargs :
            Other arguments specify smoothing methods passed to function.
            See `scHiCTools.load.processing_utils.matrix_operation` function.

        Returns
        -------
        None.

        """

        self.resolution = resolution
        self.chromosomes, self.chromosome_lengths = get_chromosome_lengths(reference_genome, chromosomes, resolution)
        self.num_of_cells = len(list_of_files)
        self.sparse = sparse
        self.keep_n_strata = keep_n_strata
        self.contacts=np.array([0]*len(list_of_files))
        self.short_range=np.array([0.0]*len(list_of_files))
        self.mitotic=np.array([0.0]*len(list_of_files))
        self.files=list_of_files
        self.strata = {
            ch: [np.zeros((self.num_of_cells, self.chromosome_lengths[ch] - i)) for i in range(keep_n_strata)]
            for ch in self.chromosomes} if keep_n_strata else None
        self.full_maps = None
        self.similarity_method=None
        self.distance=None

        assert keep_n_strata is not None or store_full_map is True

        if not store_full_map:
            self.full_maps = None
        elif sparse:
            self.full_maps = {ch: [None] * self.num_of_cells for ch in self.chromosomes}
        else:
            self.full_maps = {
                ch: np.zeros((self.num_of_cells, self.chromosome_lengths[ch], self.chromosome_lengths[ch]))
                for ch in self.chromosomes}

        print('Loading HiC data...')

        if parallelize:
            if mp is None:
                raise ImportError('Need `multiprocessing` installed to parallelize data loading process.')
            if n_processes is None:
                n_processes=mp.cpu_count()-1

            # print(n_processes)

            for ch in self.chromosomes:
                pool = mp.Pool(n_processes)

                results = [pool.apply(add_cell, args=(
                    ch,idx, file, self.resolution, self.chromosome_lengths,
                    store_full_map, keep_n_strata, format, customized_format,
                    header, adjust_resolution, map_filter, sparse, gzip,
                    operations)
                    ) for idx, file in enumerate(self.files)]

                if store_full_map:
                    for idx in range(len(self.files)):
                        self.full_maps[ch][idx]=results[idx].pop(4)
                for idx in range(len(self.files)):
                    strata=results[idx].pop(3)
                    for strata_idx, stratum in enumerate(strata):
                        self.strata[ch][strata_idx][idx, :]=stratum

                self.contacts+=[int(i[0]) for i in results]
                self.short_range+=[i[1] for i in results]
                self.mitotic+=[i[2] for i in results]

                pool.close()

                sys.stdout.write('\r')
                sys.stdout.write("Process chromosome: {} ".format(ch))
                sys.stdout.flush()

        else:
            for idx, file in enumerate(self.files):
                # print('Processing {0} out of {1} files: {2}'.format(idx+1,len(list_of_files),file))

                for ch in self.chromosomes:
                    if ('ch' in ch) and ('chr' not in ch):
                        ch=ch.replace("ch", "chr")
                    mat, strata = load_HiC(
                        file, genome_length=self.chromosome_lengths,
                        format=format, custom_format=customized_format,
                        header=header, chromosome=ch, resolution=resolution,
                        resolution_adjust=adjust_resolution,
                        map_filter=map_filter, sparse=sparse, gzip=gzip,
                        keep_n_strata=keep_n_strata, operations=operations, sep=sep,
                        **kwargs)

                    self.contacts[idx]+=np.sum(mat)/2+ np.trace(mat)/2
    # ??
                    self.short_range[idx]+=sum([np.sum(mat[i,i:i+int(2000000/self.resolution)]) for i in range(len(mat))])
                    self.mitotic[idx]+=sum([np.sum(mat[i,i+int(2000000/self.resolution):i+int(12000000/self.resolution)]) for i in range(len(mat))])

                    if store_full_map:
                        self.full_maps[ch][idx] = mat

                    if keep_n_strata:
                        # self.strata[ch][idx] = strata
                        for strata_idx, stratum in enumerate(strata):
                            self.strata[ch][strata_idx][idx, :] = stratum

                sys.stdout.write('\r')
                sys.stdout.write("[%-50s] %d/%d \t" % ('='*int((idx+1)/len(self.files)*50), idx+1,len(self.files)))
                #  File %s loaded  ,file
                sys.stdout.flush()

    def cal_strata(self, n_strata):
        """

        Alter the number of strata kept in a `scHiCs` object.

        Parameters
        ----------
        n_strata : int
            Number of strata to keep.

        Returns
        -------
        dict
            Strata of cells.

        """

        if self.full_maps is None:
            if self.keep_n_strata <= n_strata:
                print(' Only {0} strata are kept!'.format(self.keep_n_strata))
                return deepcopy(self.strata)
            else:
                return deepcopy({ch: self.strata[ch][:n_strata] for ch in self.chromosomes})
        else:
            if self.keep_n_strata is None:
                new_strata = {
                    ch: [np.zeros((self.num_of_cells, self.chromosome_lengths[ch] - i))
                         for i in range(n_strata)] for ch in self.chromosomes}
                for ch in self.chromosomes:
                    for idx in range(self.num_of_cells):
                        fmap = self.full_maps[ch][idx].toarray() if self.sparse else self.full_maps[ch][idx]
                        for i in range(n_strata):
                            new_strata[ch][i][idx, :] = np.diag(fmap[i:, :-i])
                return new_strata
            elif self.keep_n_strata >= n_strata:
                return deepcopy({ch: self.strata[ch][:n_strata] for ch in self.chromosomes})
            else:
                for ch in self.chromosomes:
                    self.strata[ch] += [(np.zeros(self.num_of_cells, self.chromosome_lengths[ch] - i))
                                        for i in range(self.keep_n_strata, n_strata)]
                    for idx in range(self.num_of_cells):
                        fmap = self.full_maps[ch][idx].toarray() if self.sparse else self.full_maps[ch][idx]
                        for i in range(self.keep_n_strata, n_strata):
                            self.strata[ch][i][idx, :] = np.diag(fmap[i:, :-i])
                return deepcopy(self.strata)

    def test_print_chr8_inner(self):
        pair_dis = pairwise_distances(self.strata['chr8'], 'inner_product')
        print(pair_dis)

    def test_selfish(self):
        chrs = ["chr1", "chr2", "chrX", "chr3", "chr4", "chr5", "chr6", "chr7",
                "chr8", "chr9", "chr10", "chr11", "chr12", "chr13", "chr14",
                "chr15", "chr16", "chr17", "chr18", "chr19", "chrY"]
        for chr in chrs:
            pair_dis = pairwise_distances(self.strata[chr], 'selfish')

    def test_fast(self):
        chrs = ["chr1", "chr2", "chrX", "chr3", "chr4", "chr5", "chr6", "chr7",
                "chr8", "chr9", "chr10", "chr11", "chr12", "chr13", "chr14",
                "chr15", "chr16", "chr17", "chr18", "chr19", "chrY"]
        timeS =0.0;
        for chr in chrs:
            pair_dis, time = pairwise_distances(self.strata['chr8'], 'hicrep')
            timeS += time
        print("all chr total: ", timeS)

    def processing(self, operations, **kwargs):
        """

        Apply a smoothing method to contact maps.
        Requre the `scHiCs` object to store the full map of contacts maps.

        Parameters
        ----------
        operations : str
            The methods use for smoothing the maps.
            Avaliable operations: 'convolution', 'random_walk', 'network_enhancing'.

        **kwargs :
            Other arguments specify smoothing methods passed to function.
            See function `scHiCTools.load.processing_utils.matrix_operation`.

        Returns
        -------
        None.

        """

        if self.full_maps is None:
            raise ValueError('No full maps stored. Processing is not doable.')
        if self.sparse:
            for ch in self.chromosomes:
                for i, mat in enumerate(self.full_maps[ch]):
                    self.full_maps[ch][i] = coo_matrix(matrix_operation(mat.toarray(), operations, **kwargs))
        else:
            for ch in self.chromosomes:
                for i, mat in enumerate(self.full_maps[ch]):
                    self.full_maps[ch][i, :, :] = matrix_operation(mat, operations, **kwargs)
        # Update the strata
        if self.keep_n_strata is not None:
            for ch in self.chromosomes:
                for i, mat in enumerate(self.full_maps[ch]):
                    for j in range(self.keep_n_strata):
                        self.strata[ch][j][i, :] = np.diag(mat[j:, :len(mat) - j])

    def plot_contacts(self, hist=True, percent=True,
                      size=1.0, bins=10, color='#1f77b4'):
        """

        Generate two plots:
        Histogram of contacts and
        scatter plot of short-range contacts v.s. contacts at the mitotic band.

        Parameters
        ----------
        hist : bool, optional
            Whether to plot Histogram of contacts.
            If `True`, plot Histogram of contacts.
            The default is True.

        percent : int, optional
            Whether to plot scatter plot of short-range contacts v.s. contacts at the mitotic band.
            If `True`, plot scatter plot of short-range contacts v.s. contacts at the mitotic band.
            The default is True.

        size : float, optional
            The point size of scatter plot.
            The default is 1.0.

        bins : int, optional
            Number of bins in histogram.
            The default is 10.

        color : str, optional
            The color of the plot.
            The default is '#1f77b4'.

        Returns
        -------
        None.

        """

        if hist:
            if percent:
                plt.subplot(1,2,1)

            plt.hist(self.contacts,bins=bins,color=color)
            plt.xlabel("Number of contacts")
            plt.ylabel('Frequency')
            plt.title('Histogram of contacts')

        if percent:
            if hist:
                plt.subplot(1,2,2)

            plt.scatter(self.mitotic*100/self.contacts,self.short_range*100/self.contacts, s=size, c=color)
            plt.xlabel("% Mitotic contacts")
            plt.ylabel("% Short-range contacts")
            plt.title('Short-range contacts v.s. contacts at the mitotic band')

    def select_cells(self, n_contacts=[0,float("inf")],
                     short_range=[0,1],
                     mitotic=[0,1],
                     selected=None):
        """

        Select qualify cells based on minimum number of contacts and
        maxium percent of short range contact.

        Parameters
        ----------
        min_n_contacts : int, optional
            The threshold of minimum number of contacts in each cell.
            The default is 0.

        short_range : list, optional
            The threshold of minimum and maximum proportion of short range contact in every cell.
            The default is [0,1].

        mitotic : list, optional
            The threshold of minimum and maximum proportion of mitotic contact in every cell.
            The default is [0,1].

        selected : list, optional
            A list of cells to be selected. Elements in the list can be either bool or str.
            The default is None.

        Returns
        -------
        list
            Selected files.

        """

        # files=np.array(self.files)
        if selected is None:
            selected=[True]*self.num_of_cells

        if np.all([type(i)==str for i in selected]):
            selected=[True if i in selected else False for i in self.files]
        elif not np.all([type(i)==bool for i in selected]):
            raise ValueError("Elements of `selected` should be bool or str!")
        elif len(selected)!=len(self.files):
            raise ValueError("When elements of `selected` are bool, length of `selected` should be length of files({})!".format(len(self.files)))

        selected=np.all([selected,
                         short_range[0] <= self.short_range/self.contacts,
                         self.short_range/self.contacts <= short_range[1],
                         n_contacts[0] <= self.contacts,
                         self.contacts <= n_contacts[1],
                         mitotic[0] <= self.mitotic/self.contacts,
                         self.mitotic/self.contacts <= mitotic[1]
                         ],axis=0)
        self.num_of_cells=sum(selected)
        self.files=[self.files[i] for i in  range(len(self.files)) if selected[i]]
        self.contacts=self.contacts[selected]
        self.short_range=self.short_range[selected]
        self.mitotic=self.mitotic[selected]
        if self.strata is not None:
            for ch in self.chromosomes:
                for s in range(self.keep_n_strata):
                    self.strata[ch][s]=self.strata[ch][s][selected,:]
        if self.full_maps is not None:
            for ch in self.chromosomes:
                self.full_maps[ch]=self.full_maps[ch][selected]
        if self.distance is not None:
            self.distance=self.distance[:,selected,:][:,:,selected]

        return self.files

    def scHiCluster(self,dim=2,n_clusters=4,cutoff=0.8,n_PCs=10,**kwargs):

        """

        Embedding and clustering single cells using HiCluster.
        Reference:
            Zhou J, Ma J, Chen Y, Cheng C, Bao B, Peng J, et al.
            Robust single-cell Hi-C clustering by convolution- and random-walk–based imputation.
            PNAS. 2019 Jul 9;116(28):14011–8.

        Parameters
        ----------
        dim : int, optional
            Number of dimension of embedding. The default is 2.

        n_clusters : int, optional
            Number of clusters. The default is 4.

        cutoff : float, optional
            The cutoff proportion to convert the real contact
            matrix into binary matrix. The default is 0.8.

        n_PCs : int, optional
            Number of principal components. The default is 10.

        **kwargs :
            Other arguments passed to kmeans.
            See `scHiCTools.analysis.clustering.kmeans` function.

        Returns
        -------
        embeddings : numpy.ndarray
            The embedding of cells using HiCluster.

        label : numpy.ndarray
            An array of cell labels clustered by HiCluster.

        """

        if self.full_maps is None:
            raise ValueError('No full maps stored. scHiCluster is not doable.')

        X=None
        for ch in self.chromosomes:
            sys.stdout.write('\r')
            sys.stdout.write('HiCluster processing chromosome {}. '.format(ch))
            # print('HiCluster processing chromosomes {}'.format(ch))
            A=self.full_maps[ch].copy()
            if len(A.shape)==3:
                n=A.shape[1]*A.shape[2]
                A.shape=(A.shape[0],n)
            A=np.quantile(A,cutoff,axis=1)<np.transpose(A)
            A = PCA(A.T,n_PCs)
            if X is None:
                X=A
            else:
                X=np.append(X, A, axis=1)

        X=PCA(X,n_PCs)
        label=kmeans(X,n_clusters,kwargs.pop('weights',None),kwargs.pop('iteration',1000))

        return X[:,:dim], label

    def learn_embedding(self, similarity_method, embedding_method,
                        dim=2, aggregation='median', n_strata=None, return_distance=False,
                        print_time=False, parallelize=False, n_processes=1,
                        **kwargs):
        """

        Function to find a low-dimensional embedding for cells.

        Parameters
        ----------
        similarity_method : str
            The method used to calculate similarity matrix.
            Now support 'inner_product', 'HiCRep' and 'Selfish'.

        embedding_method : str
            The method used to project cells into lower-dimensional space.
            Now support 'MDS', 'tSNE', 'phate', 'spectral_embedding'.

        dim : int, optional
            Dimension of the embedding space.
            The default is 2.

        aggregation : str, optional
            Method to find the distance matrix based on distance matrices of chromesomes.
            Must be 'mean' or 'median'.
            The default is 'median'.

        n_strata : int, optional
            Number of strata used in calculation.
            The default is None.

        return_distance : bool, optional
            Whether to return the distance matrix of cells.
            If True, return (embeddings, distance_matrix);
            if False, only return embeddings.
            The default is False.

        print_time : bool, optional
            Whether to print process time. The default is False.

        **kwargs :
            Including two arguments for Selfish
            (see funciton `pairwise_distances`):\
            `n_windows`: number of Selfish windows\
            `sigma`: sigma in the Gaussian-like kernel\
            and some arguments specify different embedding method
            (see functions in `scHiCTools.embedding.embedding`).

        Returns
        -------
        embeddings: numpy.ndarray
            The embedding of cells in lower-dimensional space.

        final_distance: numpy.ndarray, optional
            The pairwise distance calculated.

        """

        if self.distance is None or self.similarity_method!=similarity_method:
            self.similarity_method=similarity_method
            distance_matrices = []
            assert embedding_method.lower() in ['mds', 'tsne', 'umap', 'phate', 'spectral_embedding']
            assert n_strata is not None or self.keep_n_strata is not None
            n_strata = n_strata if n_strata is not None else self.keep_n_strata
            new_strata = self.cal_strata(n_strata)
            if print_time:
                time1=0
                time2=0
                for ch in self.chromosomes:
                    print(ch)
                    distance_mat,t1,t2 = pairwise_distances(new_strata[ch], similarity_method, print_time, kwargs.get('sigma',.5), kwargs.get('window_size',10),parallelize, n_processes)
                    time1=time1+t1
                    time2=time2+t2
                    distance_matrices.append(distance_mat)
                print('Sum of time 1:', time1)
                print('Sum of time 2:', time2)
            else:
                for i,ch in enumerate(self.chromosomes):
                    # print(ch)
                    distance_mat = pairwise_distances(new_strata[ch],
                                   similarity_method,
                                   print_time,
                                   kwargs.get('sigma',.5),
                                   kwargs.get('window_size',10))
                    distance_matrices.append(distance_mat)
                    sys.stdout.write('\r')
                    sys.stdout.write("[%-30s] %d/%d \t Calculating chromosome %s. " % ('='*int((i+1)/len(self.chromosomes)*30),i+1,len(self.chromosomes),ch))

            self.distance = np.array(distance_matrices)

        if aggregation == 'mean':
            final_distance = np.mean(self.distance, axis=0)
        elif aggregation == 'median':
            final_distance = np.median(self.distance, axis=0)
        else:
            raise ValueError('Aggregation method {0} not supported. Only "mean" or "median".'.format(aggregation))

        np.fill_diagonal(final_distance, 0)

        embedding_method = embedding_method.lower()
        if embedding_method == 'mds':
            embeddings = MDS(final_distance, dim)
        elif embedding_method == 'tsne':
            embeddings = tSNE(final_distance, dim,
                              kwargs.pop('perp',30),
                              kwargs.pop('iteration',1000),
                              kwargs.pop('momentum', 0.5),
                              kwargs.pop('rate', 200),
                              kwargs.pop('tol',1e-5))
        # elif embedding_method == 'umap':
        #     embeddings = UMAP(final_distance, dim,
        #                       kwargs.pop('n',5),
        #                       kwargs.pop('min_dist',1),
        #                       kwargs.pop('n_epochs',10),
        #                       kwargs.pop('alpha',1),
        #                       kwargs.pop('n_neg_samples',0))
        elif embedding_method == 'phate':
            embeddings = PHATE(final_distance, dim,
                               kwargs.pop('k',5),
                               kwargs.pop('a',1),
                               kwargs.pop('gamma',1),
                               kwargs.pop('t_max',100),
                               kwargs.pop('momentum',.1),
                               kwargs.pop('iteration',1000))
        elif embedding_method == 'spectral_embedding':
            graph=np.exp(-np.square(final_distance)/np.mean(final_distance**2))
            graph = graph-np.diag(graph.diagonal())
            embeddings = SpectralEmbedding(graph, dim)
        else:
            raise ValueError('Embedding method {0} not supported. '.format(embedding_method))

        if return_distance:
            return embeddings, final_distance
        else:
            return embeddings

    def clustering(self,
                   n_clusters,
                   clustering_method,
                   similarity_method,
                   aggregation='median',
                   n_strata=None,
                   print_time=False,
                   **kwargs):
        """

        Parameters
        ----------
        n_clusters : int
            Number of clusters.

        clustering_method : str
            Clustering method in 'kmeans', 'spectral_clustering' or 'HAC'(hierarchical agglomerative clustering).

        similarity_method : str
            Reproducibility measure.
            Value in ‘InnerProduct’, ‘HiCRep’ or ‘Selfish’.

        aggregation : str, optional
             Method to aggregate different chromosomes.
             Value is either 'mean' or 'median'.
             The default is 'median'.

        n_strata : int or None, optional
            Only consider contacts within this genomic distance.
            If it is None, it will use the all strata kept from previous loading process.
            The default is None.

        print_time : bool, optional
            Whether to print the processing time. The default is False.

        **kwargs :
            Other arguments pass to function `scHiCTools.embedding.reproducibility.pairwise_distances `,
            and the clustering function in `scHiCTools.analysis.clustering`.

        Returns
        -------
        label : numpy.ndarray
            An array of cell labels clustered.

        """
        if self.distance is None or self.similarity_method!=similarity_method:
            self.similarity_method=similarity_method
            distance_matrices = []
            assert n_strata is not None or self.keep_n_strata is not None
            n_strata = n_strata if n_strata is not None else self.keep_n_strata
            new_strata = self.cal_strata(n_strata)

            for i,ch in enumerate(self.chromosomes):
                # print(ch)
                distance_mat = pairwise_distances(new_strata[ch],
                                              similarity_method,
                                              print_time,
                                              kwargs.get('sigma',.5),
                                              kwargs.get('window_size',10))
                distance_matrices.append(distance_mat)
                sys.stdout.write('\r')
                sys.stdout.write("[%-30s] %d/%d \t Calculating chromosome %s. " % ('='*int((i+1)/len(self.chromosomes)*30), i+1, len(self.chromosomes), ch))
            self.distance = np.array(distance_matrices)

        if aggregation == 'mean':
            final_distance = np.mean(self.distance, axis=0)
        elif aggregation == 'median':
            final_distance = np.median(self.distance, axis=0)
        else:
            raise ValueError('Aggregation method {0} not supported. Only "mean" or "median".'.format(aggregation))

        np.fill_diagonal(final_distance, 0)

        clustering_method=clustering_method.lower()
        if clustering_method=='kmeans':
            embeddings = MDS(final_distance, n_clusters)
            label=kmeans(embeddings,
                         k=n_clusters,
                         **kwargs)
        elif clustering_method=='spectral_clustering':
            label=spectral_clustering(final_distance,
                                      data_type='distance_matrix',
                                      n_clusters=n_clusters,
                                      **kwargs)
        elif clustering_method=='hac':
            label=HAC(final_distance,
                      'distance_matrix',
                      n_clusters,
                      kwargs.pop('method','centroid'))
        else:
            raise ValueError('Embedding method {0} not supported. '.format(clustering_method))

        return label


import numpy as np
import tensorflow as tf
import collections
from konlpy.tag import Okt as Twitter
import argparse
import re
import math
import random
import time
'''
    Step 1 : Parse Arguments.
'''
parser = argparse.ArgumentParser()
parser.add_argument("input", type=str, help="input text file for training: one sentence per line")
# 수정
parser.add_argument("--embedding_size", type=int, help="embedding vector size (default=150)", default=300)
parser.add_argument("--window_size", type=int, help="window size (default=5)", default=5)
parser.add_argument("--min_count", type=int, help="minimal number of word occurences (default=5)", default=5)
parser.add_argument("--num_sampled", type=int, help="number of negatives sampled (default=50)", default=50)
parser.add_argument("--learning_rate", type=float, help="learning rate (default=1.0)", default=1.0)
parser.add_argument("--sampling_rate", type=int, help="rate for subsampling frequent words (default=0.0001)", default=0.0001)

parser.add_argument("--epochs", type=int, help="number of epochs (default=4)", default=4)

# batch_size를 어떻게 해야할까.
parser.add_argument("--batch_size", type=int, help="batch size (default=150)", default=150)

args = parser.parse_args()

'''
    Step 2 : Pre-process Data.
'''

def build_dataset(train_text, min_count, sampling_rate):
    words = list()
    with open(train_text, 'r') as f:
        lines = f.readlines()
        for line in lines:
            #  line -> line.split('\t')[1]
            sentence = re.sub(r"[^ㄱ-힣a-zA-Z0-9]+", ' ', line).strip().split()
            # sentence = re.sub(r"[^ㄱ-힣a-zA-Z0-9]+", ' ', line.split('\t')[1]).strip().split()
            if sentence:
                words.append(sentence)

    word_counter = [['UNK', -1]]
    word_counter.extend(collections.Counter([word for sentence in words for word in sentence]).most_common())
    word_counter = [item for item in word_counter if item[1] >= min_count or item[0] == 'UNK']

    word_dict = dict()
    for word, count in word_counter:
        word_dict[word] = len(word_dict)
    word_reverse_dict = dict(zip(word_dict.values(), word_dict.keys()))

    word_to_pos_li = dict()
    pos_list = list()
    twitter = Twitter()
    for w in word_dict:
        w_pos_li = list()
        for pos in twitter.pos(w, norm=True):
            w_pos_li.append(pos)

        word_to_pos_li[word_dict[w]] = w_pos_li
        pos_list += w_pos_li

    pos_counter = collections.Counter(pos_list).most_common()

    pos_dict = dict()
    for pos, _ in pos_counter:
        pos_dict[pos] = len(pos_dict)

    pos_reverse_dict = dict(zip(pos_dict.values(), pos_dict.keys()))

    word_to_pos_dict = dict()

    for word_id, pos_li in word_to_pos_li.items():
        pos_id_li = list()
        for pos in pos_li:
            pos_id_li.append(pos_dict[pos])
        word_to_pos_dict[word_id] = pos_id_li

    data = list()
    unk_count = 0
    for sentence in words:
        s = list()
        for word in sentence:
            if word in word_dict:
                index = word_dict[word]
            else:
                index = word_dict['UNK']
                unk_count += 1
            s.append(index)
        data.append(s)
    word_counter[0][1] = max(1, unk_count)

    data = sub_sampling(data, word_counter, word_dict, sampling_rate)

    return data, word_dict, word_reverse_dict, pos_dict, pos_reverse_dict, word_to_pos_dict

# Sub-sampling frequent words according to sampling_rate
def sub_sampling(data, word_counter, word_dict, sampling_rate):
    total_words = sum([len(sentence) for sentence in data])
    prob_dict = dict()
    for word, count in word_counter:
        f = count / total_words
        p = max(0, 1 - math.sqrt(sampling_rate / f))
        prob_dict[word_dict[word]] = p

    new_data = list()
    for sentence in data:
        s = list()
        for word in sentence:
            prob = prob_dict[word]
            if random.random() > prob:
                s.append(word)
        new_data.append(s)

    return new_data

data, word_dict, word_reverse_dict, pos_dict, pos_reverse_dict, word_to_pos_dict \
        = build_dataset(args.input, args.min_count, args.sampling_rate)

vocabulary_size = len(word_dict)
pos_size = len(pos_dict)
num_sentences = len(data)

print("number of sentences :", num_sentences)
print("vocabulary size :", vocabulary_size)
print("pos size :", pos_size)

pos_li = []
for key in sorted(pos_reverse_dict):
    pos_li.append(pos_reverse_dict[key])

'''
    Step 3 : Function to generate a training batch
'''

window_size = args.window_size
batch_size = args.batch_size

def generate_input_output_list(data, window_size):
    input_li = list()
    output_li = list()
    for sentence in data:
        for i in range(len(sentence)):
            for j in range(max(0, i - window_size), min(len(sentence), i + window_size + 1)):
                if i != j:
                    if sentence[i]!=word_dict['UNK'] and sentence[j]!=word_dict['UNK']:
                        input_li.append(sentence[i])
                        output_li.append(sentence[j])
    return input_li, output_li

input_li, output_li = generate_input_output_list(data, window_size)
input_li_size = len(input_li)

def generate_batch(iter, batch_size):
    index = (iter % (input_li_size//batch_size)) * batch_size
    batch_input = input_li[index:index+batch_size]
    batch_output_li = output_li[index:index+batch_size]
    batch_output = [[i] for i in batch_output_li]

    return np.array(batch_input), np.array(batch_output)

'''
    Step 4 : Build a model.
'''

embedding_size = args.embedding_size
num_sampled = args.num_sampled
learning_rate = args.learning_rate

valid_size = 20     # Random set of words to evaluate similarity on.
valid_window = 200  # Only pick dev samples in the head of the distribution.
valid_examples = np.random.choice(valid_window, valid_size, replace=False)

graph = tf.Graph()

with graph.as_default():
    # Input data
    train_inputs = tf.placeholder(tf.int32, shape=[batch_size])
    train_labels = tf.placeholder(tf.int32, shape=[batch_size, 1])
    words_matrix = [tf.placeholder(tf.int32, shape=None) for _ in range(batch_size)]
    vocabulary_matrix = [tf.placeholder(tf.int32, shape=(None)) for _ in range(vocabulary_size)]
    valid_dataset = tf.constant(valid_examples, dtype=tf.int32)

    with tf.device('/cpu:0'):
        pos_embeddings = tf.Variable(tf.random_uniform([pos_size, embedding_size], -1.0, 1.0), name='pos_embeddings')

        word_vec_list = []
        for i in range(batch_size):
            word_vec = tf.reduce_sum(tf.nn.embedding_lookup(pos_embeddings, words_matrix[i]), 0)
            word_vec_list.append(word_vec)
        word_embeddings = tf.stack(word_vec_list)

        nce_weights = tf.Variable(
            tf.truncated_normal([vocabulary_size, embedding_size], stddev=1.0 / math.sqrt(embedding_size)), name='nce_weights'
        )
        nce_biases = tf.Variable(tf.zeros([vocabulary_size]), name='nce_biases')

    loss = tf.reduce_mean(
        tf.nn.nce_loss(weights=nce_weights,
                       biases=nce_biases,
                       labels=train_labels,
                       inputs=word_embeddings,
                       num_sampled=num_sampled,
                       num_classes=vocabulary_size))

    optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss)
    init = tf.global_variables_initializer()

    # Compute the cosine similarity between minibatch exaples and all embeddings.
    norm = tf.sqrt(tf.reduce_sum(tf.square(pos_embeddings), 1, keep_dims=True))
    normalized_embeddings = pos_embeddings / norm
    valid_embeddings = tf.nn.embedding_lookup(normalized_embeddings, valid_dataset)
    similarity = tf.matmul(valid_embeddings, normalized_embeddings, transpose_b=True)

# Function to save vectors.
def save_model(pos_list, embeddings, file_name):
    with open(file_name, 'w') as f:
        f.write(str(len(pos_list)))
        f.write(" ")
        f.write(str(embedding_size))
        f.write("\n")
        for i in range(len(pos_list)):
            pos = pos_list[i]
            f.write(str(pos).replace("', '", "','") + " ")
            f.write(' '.join(map(str, embeddings[i])))
            f.write("\n")

'''
    Step 5 : Train a model.
'''

num_iterations = input_li_size // batch_size
print("number of iterations for each epoch :", num_iterations)
epochs = args.epochs
num_steps = num_iterations * epochs + 1

with tf.Session(graph=graph) as session:
    init.run()
    print("Initialized - Tensorflow")

    average_loss = 0
    for step in range(num_steps):
        batch_inputs, batch_labels = generate_batch(step, batch_size)

        word_list = []
        for word in batch_inputs:
            word_list.append(word_to_pos_dict[word])

        feed_dict = {}
        for i in range(batch_size):
            feed_dict[words_matrix[i]] = word_list[i]
        feed_dict[train_inputs] = batch_inputs
        feed_dict[train_labels] = batch_labels

        _, loss_val = session.run([optimizer, loss], feed_dict=feed_dict)
        average_loss += loss_val

        if step % 2000 == 0:
            if step > 0:
                average_loss /= 2000
            print("Average loss at step ", step, ": ", average_loss)
            average_loss = 0

        if step % 20000 == 0:
            pos_embed = pos_embeddings.eval()

            # Print nearest words
            # sim = similarity.eval()
            # for i in range(valid_size):
            #     valid_pos = pos_reverse_dict[valid_examples[i]]
            #     top_k = 8
            #     nearest = (-sim[i, :]).argsort()[1:top_k + 1]
            #     log_str = 'Nearest to %s:' % str(valid_pos)
            #     for k in range(top_k):
            #         close_word = pos_reverse_dict[nearest[k]]
            #         log_str = '%s %s,' % (log_str, str(close_word))
            #     print(log_str)

    # Save vectors
    save_model(pos_li, pos_embeddings.eval(), "청원문.vec")


"""Classification algorithm for the kids database.

Requirements:
    - A csv file containing the extracted landmark points from the kids
    videos.
    - A trained lstm.

Example:
    python cnn_lstm_kids.py --source ./extracted.csv --model ./lstm.hdf5
    --output ./results.csv
"""
import cv2
from fee.classification import Expression as Exp
from fee.io import open_landmarks_files
from fee.io import read_landmarks_file
from keras.models import load_model
import numpy as np
import argparse

# Little hack that I don't understand, in order to reuse an existing model.
from keras import backend as K
K.set_learning_phase(1)  # Set learning phase

# Parsing the command arguments.
parser = argparse.ArgumentParser()
parser.add_argument("-s", "--source",
                    help="Path of the kids csv files.")
parser.add_argument("-m", "--model",
                    help="Path of the trained lstm model.")
parser.add_argument("-o", "--output",
                    help="Path of the output file.")
args = parser.parse_args()

# The global paths variables
source_path = args.source   # The extracted landmarks csv file path
model_path  = args.model    # The pre-trained model path
output_path = args.output   # The output csv file path

def group_line_by_file(file):
    """Parse an extracted landmarks file.

    Read each line of a csv file and return a list of all datas from a same
    video.

    Parameters:
        file: the landmarks csv file, opened in read mode.

    Returns:
        current_file: the current returned file path
        exp: the current returned file main expression
        datas: a list of the current returned file landmarks
        (see fee.io.read_landmarks_file)
    """
    current_file = None
    datas = []
    for data in read_landmarks_file(file):
        filepath, frame_id, exp, bounds, points = data
        if current_file != filepath:
            if current_file is not None:
                yield current_file, exp, datas
            datas = []
            current_file = filepath
        datas.append(data)
    yield current_file, exp, datas

def extract_face_frame(frame, bounds, size, offset):
    """Crop and resize a frame.

    Parameters:
        frame: A frame.
        bounds: The bounding box to crop.
        size: The size to apply to the cropped part.
        offset: The offset to add to the bounding box.

    Returns:
        A frame.
    """
    top_margin = bounds["width"] - bounds["height"]
    frame = frame[bounds["top"] - offset - top_margin:
                  bounds["top"] + bounds["height"] + offset,
                  bounds["left"] - offset:
                  bounds["left"] + bounds["width"] + offset]
    frame = cv2.resize(frame, (size))
    return frame

def format_face_frame(frame):
    """Format for the face_classification model.

    Parameters:
        frame: A frame.

    Returns:
        A frame.
    """
    frame = frame.astype('float32')
    frame = frame / 255.0
    frame = frame - 0.5
    frame = frame * 2.0
    frame = np.expand_dims(frame, -1)
    return frame

# The number of frame for each sequence to classify. Must be equivalent of the
# number of frame used for the trained model.
frame_count = 20

model  = load_model(model_path)
labels = [Exp.ANGER, Exp.DISGUST, Exp.FEAR, Exp.HAPPINESS,
          Exp.SADNESS, Exp.SURPRISE, Exp.NEUTRAL]

# Write the header of the output csv file.
output_file = open(output_path, 'w')
output_file.write('filepath,'+(',').join([l.to_str() for l in labels])+'\n')

# For each extracted landmarks csv file, we do the prediction.
for f in open_landmarks_files(source_path):
    for filepath, exp, datas in group_line_by_file(f):
        exp = Exp.from_str(exp[0])
        # Retrieve the faces frame by frame
        cap = cv2.VideoCapture(filepath, cv2.CAP_FFMPEG)
        frameid = 0
        frames = []
        while(cap.isOpened()):
            ret, frame = cap.read()
            if ret:
                filepath, frame_id, exp, bounds, points = datas[frameid]
                # The dlib algorithm might not succeeded in finding a face.
                # In such case, the bounding box is equal to None. So we'll do
                # the face extraction on the frame which actually have a
                # bounding box.
                if bounds is not None:
                    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                    frame = extract_face_frame(frame, bounds, (64, 64), 15)
                    frame = format_face_frame(frame)
                    frames.append(frame)
                frameid += 1
            else:
                cap.release()
        # Predict the expression. We predict for n sequences of <frame_count>
        # frames. Each sequence starts with a slide of 5 frame from the
        # previous one.
        i = 0
        X = []
        # We fill the predict set.
        while i + frame_count < len(frames):
            X.append(np.array(frames[i:i+frame_count]))
            i += 5
        # There might be no sequence to predict, because the video is not long
        # enough (less than <frame_count> frames) or because the dlib algorithm
        # couldn't find any face.
        if len(X) > 0:
            X = np.array(X)
            print(filepath)
            prediction = model.predict(X)
            argmax = np.argmax(prediction, axis=1)
            # Uncomment the following line to remove the NEUTRAL prediction.
            # argmax = np.argmax(prediction[:, :-1], axis=1)
            unique, counts = np.unique(argmax, return_counts=True)
            counts = [c / len(argmax) for c in counts]
            results = np.zeros(len(labels))
            for i in range(0, len(unique)):
                results[unique[i]] = counts[i]
            output_file.write(filepath+',')
            output_file.write((',').join([str(r) for r in results])+'\n')


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

# Convert FEATURES with trained models.
# By Wen-Chin Huang 2019.06

import json
import os

import tensorflow as tf
import numpy as np

from importlib import import_module

import argparse
import logging

import sys
from preprocessing.vcc2018.feature_reader import Whole_feature_reader
from util.normalizer import MinMaxScaler, StandardScaler
from util.misc import read_hdf5, read_txt, load, get_default_logdir_output

import glob

def main():

    parser = argparse.ArgumentParser(
        description="Conversion.")
    parser.add_argument("--logdir", required=True, type=str,
        help="path of log directory")
    parser.add_argument("--checkpoint", default=None, type=str,
        help="path of checkpoint")

    parser.add_argument("--src", default=None, required=True, type=str,
        help="source speaker")
    parser.add_argument("--trg", default=None, required=True, type=str,
        help="target speaker")
    parser.add_argument("--type", default='test', type=str,
        help="test or valid (default is test)")

    parser.add_argument("--input_feat", required=True,
        type=str, help="input feature type")
    parser.add_argument("--output_feat", required=True,
        type=str, help="output feature type")
    parser.add_argument("--mcd", action='store_true',
        help="calculate mcd or not")
    parser.add_argument("--syn", action='store_true',
        help="synthesize voice or not")
    args = parser.parse_args()

    # make exp directory
    output_dir = get_default_logdir_output(args)
    tf.gfile.MakeDirs(output_dir)

    # set log level
    fmt = '%(asctime)s %(message)s'
    datefmt = '%m/%d/%Y %I:%M:%S'
    logFormatter = logging.Formatter(fmt, datefmt=datefmt)
    logging.basicConfig(
        level=logging.INFO,
        filename=os.path.join(output_dir, 'exp.log'),
        format=fmt,
        datefmt=datefmt,
        )
    consoleHandler = logging.StreamHandler()
    consoleHandler.setFormatter(logFormatter)
    logging.getLogger().addHandler(consoleHandler)
    logging.info('====================')
    logging.info('Conversion start')
    logging.info(args)

    # Load architecture
    # arch = tf.gfile.Glob(os.path.join(args.logdir, 'architecture*.json'))[0]  # should only be 1 file
    arch = glob.glob(os.path.join(args.logdir, 'architecture*.json'))[0]  # should only be 1 file
    with open(arch) as fp:
        arch = json.load(fp)

    # Load the model module
    module = import_module(arch['model_module'], package=None)
    MODEL = getattr(module, arch['model'])

    input_feat = args.input_feat
    input_feat_dim = arch['feat_param']['dim'][input_feat]
    output_feat = args.output_feat

    # read speakers
    spk_list = read_txt(arch['spklist'])

    # Load statistics, normalize and NCHW
    normalizers = {}
    for k in arch['normalizer']:
        normalizers[k] = {}
        for norm_type in arch['normalizer'][k]['type']:
            if norm_type == 'minmax':
                normalizer = MinMaxScaler(
                    xmax=read_hdf5(arch['stats'], '/max/' + k),
                    xmin=read_hdf5(arch['stats'], '/min/' + k),
                )
            elif norm_type == 'meanvar':
                normalizer = StandardScaler(
                    mu=read_hdf5(arch['stats'], '/mean/' + k),
                    std=read_hdf5(arch['stats'], '/scale/' + k),
                )

            normalizers[k][norm_type] = normalizer

    # Define placeholders
    x_pl = tf.placeholder(tf.float32, [None, input_feat_dim])

    yh_pl = tf.placeholder(dtype=tf.int64, shape=[1,])
    yh = yh_pl * tf.ones(shape=[tf.shape(x_pl)[0],], dtype=tf.int64)
    yh = tf.expand_dims(yh, 0)

    # Define model
    model = MODEL(arch, normalizers)
    z, _ = model.encode(x_pl, input_feat)
    xh = model.decode(z, yh, output_feat)

    # make directories for output
    tf.gfile.MakeDirs(os.path.join(output_dir, 'latent'))
    tf.gfile.MakeDirs(os.path.join(output_dir, 'converted-{}'.format(output_feat)))

    # Define session
    with tf.Session() as sess:

        # define saver
        saver = tf.train.Saver()

        # load checkpoint
        if args.checkpoint is None:
            load(saver, sess, args.logdir,)
        else:
            _, ckpt = os.path.split(args.checkpoint)
            load(saver, sess, args.logdir, ckpt=ckpt)

        # get feature list, either validation set or test set
        if args.type == 'test':
            # files = tf.gfile.Glob(arch['conversion']['test_file_pattern'].format(args.src))
            files = glob.glob(arch['conversion']['test_file_pattern'].format(args.src))
        elif args.type == 'valid':
            files = []
            for p in arch['training']['valid_file_pattern']:
                # files.extend(tf.gfile.Glob(p.replace('*', args.src)))
                files.extend(glob.glob(p.replace('*', args.src)))
        files = sorted(files)

        # conversion
        for f in files:
            basename = os.path.split(f)[-1]
            path_to_latent = os.path.join(output_dir, 'latent', '{}-{}-{}'.format(args.src, args.trg, basename))
            path_to_cvt = os.path.join(output_dir, 'converted-{}'.format(output_feat), '{}-{}-{}'.format(args.src, args.trg, basename))
            logging.info(basename)

            # load source features
            src_data = Whole_feature_reader(f, arch['feat_param'])

            #
            latent, cvt = sess.run([z, xh],
                               feed_dict={yh_pl : np.asarray([spk_list.index(args.trg)]),
                                          x_pl : src_data[input_feat] }
                              )
            # save bin
            with open(path_to_latent, 'wb') as fp:
                fp.write(latent.tostring())
            with open(path_to_cvt, 'wb') as fp:
                fp.write(cvt.tostring())

    # optionally calculate MCD
    if args.mcd:
        cmd = "python ./mcd_calculate.py" + \
                    " --type " + args.type + \
                    " --logdir " + output_dir + \
                    " --input_feat " + input_feat + \
                    " --output_feat " + output_feat
        print(cmd)
        os.system(cmd)

    # optionally synthesize waveform
    if args.syn:
        cmd = "python ./synthesize.py" + \
                    " --type " + args.type + \
                    " --logdir " + output_dir + \
                    " --input_feat " + input_feat + \
                    " --output_feat " + output_feat
        print(cmd)
        os.system(cmd)

if __name__ == '__main__':
    main()


from mido import MidiFile, MidiTrack
import lib
import numpy as np
import numpy.random as rd
from pydub import AudioSegment
from pydub.playback import play
from importlib import reload
reload(lib)

b = 480
c = int(b/2)
cc = int(c/2)
t = int(b/3)
w = int(2*b)

ppp = [7, 22]
pp = [23, 38]
p = [39, 54]
mp = [55, 70]
mf = [71, 86]
f = [87, 102]
ff = [103, 111]
fff = [112, 127]

mf = mp

def left_motif1(nuance, vel_init=None, second_nuance=None):
    """Make the motif repeat with distinct initial values and min/max"""
    notes = [['e3', b], ['g3', b], ['e3', b], ['g3', b]]
    notes += [['d3', b], ['g3', b], ['d3', b], ['g3', b]]
    notes += [['d3', b], ['f3+', b], ['d3', b], ['f3+', b]]
    notes += [['d3', b], ['f3+', b], ['d3', b], ['f3+', b]]

    if not vel_init:
        vel_init = rd.randint(nuance[0], nuance[1])

    if not second_nuance:
        vel = lib.velocity(vel_init, nuance, [[3, 0], [4, 1], [4, 0], [4, -1]],
                           offset=10)
    else:
        vel = lib.velocity(vel_init, nuance, [[3, 0], [4, 1]],
                           offset=10)
        vel += lib.velocity(vel[-1], second_nuance, [[3, 0], [4, -1]],
                           offset=10)
    lhand1 = [notes[i] + [c_vel] for i, c_vel in enumerate(vel)]

    return lhand1, vel

def left_motif2(nuance, vel_init=None, second_nuance=None):
    """Make the motif repeat with distinct initial values and min/max"""
    notes = 4 * [['e4', c], ['b3', c]]
    notes += 4 * [['d4', c], ['b3', c]]
    notes += 4 * [['d4', c], ['b3', c]]
    notes += 4 * [['d4', c], ['a3', c]]

    if not vel_init:
        vel_init = rd.randint(nuance[0], nuance[1])

    if not second_nuance:
        vel = lib.velocity(vel_init, nuance, [[7, 0], [8, 1], [8, 0], [8, -1]],
                           offset=15)
    else:
        vel = lib.velocity(vel_init, nuance, [[7, 0], [8, 1]],
                           offset=15)
        vel += lib.velocity(vel[-1], second_nuance, [[7, 0], [8, -1]],
                           offset=15)

    lhand2 = [notes[i] + [c_vel] for i, c_vel in enumerate(vel)]

    return lhand2, vel

def right_motif1(nuance, vel_init=None, second=False):
    if not second:
        notes = [['s', c], ['g4', cc], ['f4+', cc], ['g4', c], ['b4', cc], ['c5', cc], ['b4', w],
                 ['s', c], ['f4+', cc], ['g4', cc], ['f4+', c], ['g4', cc], ['a4', cc], ['g4', w],
                 ['s', c], ['f4+', cc], ['e4', cc], ['f4+', c], ['b4', cc], ['c5', cc], ['b4', w],
                 ['s', c], ['f4+', cc], ['e4', cc], ['f4+', w+b]]
    else:
        notes = [['s', c], ['g5', cc], ['f5+', cc], ['g5', c], ['b5', cc], ['c6', cc], ['b5', w],
                 ['s', c], ['f5+', cc], ['g5', cc], ['f5+', c], ['g5', cc], ['a5', cc], ['g5', w],
                 ['s', c], ['f5+', cc], ['e5', cc], ['f5+', c], ['b5', cc], ['c6', cc], ['b5', w],
                 ['s', c], ['f5+', cc], ['e5', cc], ['f5+', w+b]]

    if not vel_init:
        vel_init = rd.randint(nuance[0], nuance[1])
    vel = lib.velocity(vel_init, nuance, [[7, 0], [3, 1], [4, -1], [3, 1],
                                          [4, -1], [3, -1]],
                       acc=[[0, 0], [7, 0], [14, 0], [21, 0]])

    rhand = [notes[i] + [c_vel] for i, c_vel in enumerate(vel)]
    return rhand, vel

def right_motif2(nuance1, nuance2, vel_init=None ,second=False):
    if not second:
        notes = [['e5', b+c], ['b4', c+w], ['d5', b+c], ['b4', w+c],
                       ['f5+', b+c], ['b4', c+w], ['f5+', b+c], ['a4', w+c]]
        notes += [['b4 g5', b+c], ['g4 e5', c+w], ['b4 g5', b+c], ['g4 d5', w+c],
                        ['b4 f5+', b+c], ['f4+ d5', c+w], ['a4 f5+', b+c], ['f4+ d5', w+c]]
    else:
        notes = [['e6', b+c], ['b5', c+w], ['d6', b+c], ['b5', w+c],
                       ['f6+', b+c], ['b5', c+w], ['f6+', b+c], ['a5', w+c]]
        notes += [['b5 g6', b+c], ['g5 e6', c+w], ['b5 g6', b+c], ['g5 d6', w+c],
                        ['b5 f6+', b+c], ['f5+ d6', c+w], ['a5 f6+', b+c], ['f5+ d6', w+c]]
    if not vel_init:
        vel_init = rd.randint(nuance1[0], nuance1[1])
    vel = lib.velocity(vel_init, nuance1, [[1, 0], [2, 1], [2, 0], [2, -1]])
    vel_init = rd.randint(nuance2[0], nuance2[1])
    vel += lib.velocity(vel_init, nuance2, [[1, 0], [2, 1], [2, 0], [2, -1]])

    rhand = [notes[i] + [c_vel] for i, c_vel in enumerate(vel)]
    return rhand, vel

def right_motif3(nuance, vel_init=None, rep=False, first=True, second_nuance=None):
    if first:
        notes = 4*['b4', 'e5', 'b5'] + ['b4', 'e5', 'c5', 'e5']
        notes += 4*['b4', 'd5', 'b5'] + ['b4', 'd5', 'a4', 'd5']
        notes += 4*['f4+', 'b4', 'f5+'] + ['f4+', 'b4', 'g4', 'b4']
        if rep:
            notes += 4*['a4', 'd5', 'a5'] + ['a4', 'd5', 'a5']
        else:
            notes += 4*['a4', 'd5', 'a5'] + ['a4', 'd5', 'g4', 'd5']
    else:
        notes = 4*['b5', 'e6', 'b6'] + ['b5', 'e6', 'c6', 'e6']
        notes += 4*['b5', 'd6', 'b6'] + ['b5', 'd6', 'a5', 'd6']
        notes += 4*['f5+', 'b5', 'f6+'] + ['f5+', 'b5', 'g5', 'b5']
        if not rep:
            notes += 4*['a5', 'd6', 'a6'] + ['a5', 'd6', 'g5', 'd6']
        else:
            notes += 4*['a5', 'd6', 'a6'] + ['a5', 'd6', 'a6', 'g6']

    if not vel_init:
        vel_init = rd.randint(nuance[0], nuance[1])

    acc = np.array([0, 12, 14])
    if rep:
        acc_t = acc.tolist() + (acc + 16).tolist() + (acc + 32).tolist() + (acc + 48).tolist()
    else:
        acc_t = acc.tolist() + (acc + 16).tolist() + (acc + 32).tolist() + [48]

    if second_nuance:
        if rep:
            acc_t = acc.tolist() + (acc + 16).tolist()
        else:
            acc_t = acc.tolist() + [16]

    acc = [[a , int(rd.randint(nuance[1], nuance[1]+10))] for a in acc_t]

    if first:
        if rep:
            vel = lib.velocity(vel_init, nuance, [[15, 0], [16, 1], [16, 0],
                                              [15, -1]], acc=acc)
        else:
            vel = lib.velocity(vel_init, nuance, [[15, 0], [16, 1], [16, 0],
                                                  [16, -1]], acc=acc)
    else:
        vel = lib.velocity(vel_init, nuance, [[15, 0], [16, 1]],
                           acc=acc)
        vel += lib.velocity(vel[-1], second_nuance, [[15, 0], [16, -1]], acc=acc)

    rhand = [[notes[i], cc] + [c_vel] for i, c_vel in enumerate(vel)]
    if rep and first:
        rhand[-1][1] = c
    return rhand, vel

def left(mid):
    lh1, vl1 = left_motif1(p)
    lh2, vl2 = left_motif2(p)
    for i in range(2):
        lhand1, vl1 = left_motif1(mf)
        lhand2, vl2 = left_motif2(mf)
        lh1 += lhand1
        lh2 += lhand2
    lhand1, vl1 = left_motif1(p)
    lhand2, vl2 = left_motif2(p)
    lh1 += lhand1
    lh2 += lhand2
    lhand1, vl1 = left_motif1(mf)
    lhand2, vl2 = left_motif2(mf)
    lh1 += lhand1
    lh2 += lhand2
    for i in range(2):
        lhand1, vl1 = left_motif1(mf, vel_init=vl1[-1])
        lhand2, vl2 = left_motif2(mf, vel_init=vl2[-1])
        lh1 += lhand1
        lh2 += lhand2
    lhand1, vl1 = left_motif1(p)
    lhand2, vl2 = left_motif2(p)
    lh1 += lhand1
    lh2 += lhand2
    lhand1, vl1 = left_motif1(p, vl1[-1])
    lhand2, vl2 = left_motif2(p, vl2[-1])
    lh1 += lhand1
    lh2 += lhand2
    lhand1, vl1 = left_motif1(p, vl1[-1])
    lhand2, vl2 = left_motif2(p, vl2[-1])
    lh1 += lhand1
    lh2 += lhand2
    lhand1, vl1 = left_motif1(mf)
    lhand2, vl2 = left_motif2(mf)
    lh1 += lhand1
    lh2 += lhand2
    lhand1, vl1 = left_motif1(p, vel_init=vl1[-1], second_nuance=mf)
    lhand2, vl2 = left_motif2(p, vel_init=vl2[-1], second_nuance=mf)
    lh1 += lhand1
    lh2 += lhand2
    lhand1, vl1 = left_motif1(p, vel_init=vl1[-1], second_nuance=mf)
    lhand2, vl2 = left_motif2(p, vel_init=vl2[-1], second_nuance=mf)
    lhand1[-2][1] = w
    lh1 += lhand1[:-1] + [["e3 b3 e4", w, rd.randint(pp[0], pp[1])]]
    lhand2[-2][1] = b
    lh2 += lhand2[:-1]
    trkl1 = lib.notes2trk(lh1)
    trkl2 = lib.notes2trk(lh2)
    mid.tracks.append(trkl1)
    mid.tracks.append(trkl2)
    return mid

def right(mid):
    rh = [['s', 8*w]]
    rh_t, vel = right_motif1(mf)
    rh += rh_t
    rh_t, vel = right_motif1(mf, vel_init=vel[-1])
    rh += rh_t
    rh_t, vel = right_motif2(p, mf, vel_init=vel[-1])
    rh += rh_t
    rh_t, vel = right_motif3(mf, vel_init=vel[-1], first=True, rep=False)
    rh += rh_t
    rh_t, vel = right_motif3(mf, vel_init=vel[-1], first=True, rep=True)
    rh += rh_t
    rh_t, vel = right_motif1(p, vel_init=vel[-1], second=True)
    rh += rh_t
    rh_t, vel = right_motif1(p, vel_init=vel[-1], second=True)
    rh += rh_t
    rh_t, vel = right_motif2(p, mp, vel_init=vel[-1], second=True)
    rh += rh_t
    rh_t, vel = right_motif3(p, vel_init=vel[-1], first=False, rep=False,
                             second_nuance=mf)
    rh += rh_t
    rh_t, vel = right_motif3(p, vel_init=vel[-1], first=False, rep=True,
                             second_nuance=mf)
    rh += rh_t + [["g5 b5 e6", w, rd.randint(pp[0], pp[1])]]
    trkr = lib.notes2trk(rh)
    mid.tracks.append(trkr)
    return mid

def tempo(mid):
    beats = [rd.randint(600000, 700000)]
    beats += lib.rand_ev(600000, 700000, 6*16+11, 10000, seed=beats[-1])
    beats += lib.rand_ev(650000, 700000, 4, 10000, seed=beats[-1], dec=False)
    beats += [rd.randint(600000, 700000)]
    beats += lib.rand_ev(600000, 700000, 5*16+11, 10000, seed=beats[-1])
    beats += lib.rand_ev(1800000, 2500000, 4, 100000, seed=beats[-1], dec=False)
    trk_tempo = lib.tempo(beats)
    mid.tracks.append(trk_tempo)
    return mid

def song(n="Comptine1", ticks_per_beat=b):
    """Create the Audiosegment"""
    mid = MidiFile()
    mid.ticks_per_beat = ticks_per_beat
    mid = left(mid)
    mid = right(mid)
    mid = tempo(mid)
    mid.save("I" + n + ".mid")
    return lib.mid2aud("I" + n) + 10

silence = AudioSegment.silent(duration=2000)
out = silence

for i in range(5):
    out += song(n) + silence

out.export("Comptine.wav")


import face_recognition
import cv2
import numpy as np
from livenessmodel import get_liveness_model
from common import get_users
font = cv2.FONT_HERSHEY_DUPLEX

# Get the liveness network
model = get_liveness_model()

# load weights into new model
model.load_weights("model/model.h5")
print("Loaded model from disk")

# Read the users data and create face encodings
known_names, known_encods = get_users()

video_capture = cv2.VideoCapture(0)
video_capture.set(3, 640)
video_capture.set(4, 480)

# Initialize some variables
face_locations = []
face_encodings = []
face_names = []
process_this_frame = True
input_vid = []

while True:
    # Grab a single frame of video
    if len(input_vid) < 24:

        ret, frame = video_capture.read()

        liveimg = cv2.resize(frame, (100,100))
        liveimg = cv2.cvtColor(liveimg, cv2.COLOR_BGR2GRAY)
        input_vid.append(liveimg)
    else:
        ret, frame = video_capture.read()

        liveimg = cv2.resize(frame, (100,100))
        liveimg = cv2.cvtColor(liveimg, cv2.COLOR_BGR2GRAY)
        input_vid.append(liveimg)
        inp = np.array([input_vid[-24:]])
        inp = inp/255
        inp = inp.reshape(1,24,100,100,1)
        pred = model.predict(inp)
        input_vid = input_vid[-25:]

        if pred[0][0]> .95:

            # Resize frame of video to 1/4 size for faster face recognition processing
            small_frame = cv2.resize(frame, (0, 0), fx=0.25, fy=0.25)

            # Only process every other frame of video to save time
            if process_this_frame:
                # Find all the faces and face encodings in the current frame of video
                face_locations = face_recognition.face_locations(small_frame)
                face_encodings = face_recognition.face_encodings(small_frame, face_locations)
                name = "Unknown"
                face_names = []
                for face_encoding in face_encodings:
                    for ii in range(len(known_encods)):
                        # See if the face is a match for the known face(s)
                        match = face_recognition.compare_faces([known_encods[ii]], face_encoding)

                        if match[0]:
                            name = known_names[ii]

                    face_names.append(name)

            process_this_frame = not process_this_frame

            unlock = False
            for n in face_names:

                if n != 'Unknown':
                    unlock=True

            # Display the results
            for (top, right, bottom, left), name in zip(face_locations, face_names):
                # Scale back up face locations since the frame we detected in was scaled to 1/4 size
                top *= 4
                right *= 4
                bottom *= 4
                left *= 4

                # Draw a box around the face
                cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)

                # Draw a label with a name below the face
                cv2.rectangle(frame, (left, bottom - 35), (right, bottom), (0, 0, 255), cv2.FILLED)

                cv2.putText(frame, name, (left + 6, bottom - 6), font, 1.0, (255, 255, 255), 1)
                if unlock:
                    cv2.putText(frame, 'UNLOCK', (frame.shape[1]//2, frame.shape[0]//2), font, 1.0, (255, 255, 255), 1)
                else:
                    cv2.putText(frame, 'LOCKED!', (frame.shape[1]//2, frame.shape[0]//2), font, 1.0, (255, 255, 255), 1)
        else:
            cv2.putText(frame, 'WARNING!', (frame.shape[1]//2, frame.shape[0]//2), font, 1.0, (255, 255, 255), 1)
        # Display the liveness score in top left corner
        cv2.putText(frame, str(pred[0][0]), (20, 20), font, 1.0, (255, 255, 0), 1)
        # Display the resulting image
        cv2.imshow('Video', frame)

        # Hit 'q' on the keyboard to quit!
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

# Release handle to the webcam
video_capture.release()
cv2.destroyAllWindows()


import tensorflow as tf
import numpy as np
import pickle
#from reward_model.utils import LimitedRunningStat, RunningStat
from utils import DynamicRunningStat, LimitedRunningStat, RunningStat
import random
from math import sqrt
from utils import *

eps = 1e-12

class RewardModel:

    def __init__(self, actions_size, policy, sess = None, gamma=0.99, lr=1e-5, batch_size=32, num_itr=20,
                 use_vairl=False, mutual_information=0.5, alpha=0.0005, with_action=False, name='reward_model',
                 entropy_weight = 0.5, with_value=True, fixed_reward_model=False,
                 vs=None, **kwargs):

        # Initialize some model attributes
        # RunningStat to normalize reward from the model
        if not fixed_reward_model:
            self.r_norm = DynamicRunningStat()
        else:
            self.r_norm = RunningStat(1)

        # Discount factor
        self.gamma = gamma
        # Policy agent needed to compute the discriminator
        self.policy = policy
        # Demonstrations buffer
        self.expert_traj = None
        self.validation_traj = None
        # Num of actions available in the environment
        self.actions_size = actions_size
        # If is state-only or state-action discriminator
        self.with_action = with_action

        # TF parameters
        self.sess = sess
        self.lr = lr
        self.batch_size = batch_size
        self.num_itr = num_itr
        self.entropy_weight = entropy_weight

        # Use Variation Bottleneck Autoencoder
        self.use_vairl = use_vairl
        self.mutual_information = mutual_information
        self.alpha = alpha
        self.name = name

        # Buffer of policy experience with which train the reward model
        self.buffer = dict()
        self.create_buffer()

        with tf.compat.v1.variable_scope(name) as vs:
            with tf.compat.v1.variable_scope('irl'):

                # Input spec for both reward and value function

                # Current state (DeepCrawl spec)
                self.global_state = tf.compat.v1.placeholder(tf.float32, [None, 10, 10, 52], name='global_state')
                self.local_state = tf.compat.v1.placeholder(tf.float32, [None, 5, 5, 52], name='local_state')
                self.local_two_state = tf.compat.v1.placeholder(tf.float32, [None, 3, 3, 52], name='local_two_state')
                self.agent_stats = tf.compat.v1.placeholder(tf.int32, [None, 16], name='agent_stats')
                self.target_stats = tf.compat.v1.placeholder(tf.int32, [None, 15], name='target_stats')
                if self.with_action:
                    self.acts = tf.compat.v1.placeholder(tf.int32, [None, 1], name='acts')

                # Next state (DeepCrawl spec) - for discriminator
                self.global_state_n = tf.compat.v1.placeholder(tf.float32, [None, 10, 10, 52], name='global_state_n')
                self.local_state_n = tf.compat.v1.placeholder(tf.float32, [None, 5, 5, 52], name='local_state_n')
                self.local_two_state_n = tf.compat.v1.placeholder(tf.float32, [None, 3, 3, 52], name='local_two_state_n')
                self.agent_stats_n = tf.compat.v1.placeholder(tf.int32, [None, 16], name='agent_stats_n')
                self.target_stats_n = tf.compat.v1.placeholder(tf.int32, [None, 15], name='target_stats_n')

                # Probability distribution and labels - whether or not this state belongs to expert buffer
                self.probs = tf.compat.v1.placeholder(tf.float32, [None, 1], name='probs')
                self.labels = tf.compat.v1.placeholder(tf.float32, [None, 1], name='labels')

                # For V-AIRL
                self.use_noise = tf.compat.v1.placeholder(
                    shape=[1], dtype=tf.float32, name="noise"
                )

                self.z_sigma_g = None
                self.z_sigma_h = None
                if self.use_vairl:
                    self.z_sigma_g = tf.compat.v1.get_variable(
                        'z_sigma_g',
                        100,
                        dtype=tf.float32,
                        initializer=tf.compat.v1.ones_initializer(),
                    )
                    self.z_sigma_g_sq = self.z_sigma_g * self.z_sigma_g
                    self.z_log_sigma_g_sq = tf.compat.v1.log(self.z_sigma_g_sq + eps)

                    self.z_sigma_h = tf.compat.v1.get_variable(
                        "z_sigma_h",
                        100,
                        dtype=tf.float32,
                        initializer=tf.compat.v1.ones_initializer(),
                    )
                    self.z_sigma_h_sq = self.z_sigma_h * self.z_sigma_h
                    self.z_log_sigma_h_sq = tf.compat.v1.log(self.z_sigma_h_sq + eps)

                # Reward Funvtion
                with tf.compat.v1.variable_scope('reward'):
                    self.reward, self.z_g = self.conv_net(self.global_state, self.local_state, self.local_two_state,
                                                          self.agent_stats, self.target_stats, with_action = self.with_action,
                                                          z_sigma=self.z_sigma_g, use_noise=self.use_noise)

                # Value Function
                if with_value:
                    with tf.compat.v1.variable_scope('value'):
                        self.value, self.z_h = self.conv_net(self.global_state, self.local_state, self.local_two_state,
                                                             self.agent_stats, self.target_stats,
                                                             z_sigma=self.z_sigma_h, use_noise=self.use_noise, with_action=False)
                    with tf.compat.v1.variable_scope('value', reuse=True):
                        self.value_n, self.z_1_h = self.conv_net(self.global_state_n, self.local_state_n,
                                                                 self.local_two_state_n,
                                                                 self.agent_stats_n, self.target_stats_n,
                                                                 z_sigma=self.z_sigma_h, use_noise=self.use_noise, with_action=False)

                    self.f = self.reward + self.gamma * self.value_n - self.value
                else:
                    self.f = self.reward

                # Discriminator
                self.discriminator = tf.math.divide(tf.math.exp(self.f), tf.math.add(tf.math.exp(self.f), self.probs))

                # Loss Function
                self.loss = -tf.reduce_mean((self.labels * tf.math.log(self.discriminator + eps)) + (
                            (1 - self.labels) * tf.math.log(1 - self.discriminator + eps)))

                # Loss function modification for V-AIRL
                if self.use_vairl:
                    # Define beta
                    self.beta = tf.compat.v1.get_variable(
                        "airl_beta",
                        [],
                        trainable=False,
                        dtype=tf.float32,
                        initializer=tf.compat.v1.ones_initializer(),
                    )

                    # Number of batch element
                    self.batch = tf.compat.v1.shape(self.z_g)[0]
                    self.batch_index = tf.dtypes.cast(self.batch / 2, tf.int32)

                    self.kl_loss = tf.reduce_mean(
                        -tf.reduce_sum(
                            1
                            + self.z_log_sigma_g_sq
                            - 0.5 * tf.square(
                                self.z_g[0:self.batch_index, :] * self.z_h[0:self.batch_index, :] * self.z_1_h[
                                                                                0:self.batch_index,
                                                                                :])
                            - 0.5 * tf.square(
                                self.z_g[self.batch_index:, :] * self.z_h[self.batch_index:, :] * self.z_1_h[
                                                                                self.batch_index:, :])
                            - tf.exp(self.z_log_sigma_g_sq),
                            1,
                        )
                    )

                    self.loss = self.beta * (self.kl_loss - self.mutual_information) + self.loss

                # Adam optimizer with gradient clipping
                optimizer = tf.compat.v1.train.AdamOptimizer(self.lr)
                gradients, variables = zip(*optimizer.compute_gradients(self.loss))
                gradients, _ = tf.compat.v1.clip_by_global_norm(gradients, 1.0)
                self.step = optimizer.apply_gradients(zip(gradients, variables))
                #self.step = tf.compat.v1.train.AdamOptimizer(learning_rate=self.lr).minimize(self.loss)

                if self.use_vairl:
                    self.make_beta_update()

        self.vs = vs
        self.saver = tf.compat.v1.train.Saver(max_to_keep=None)

    ## Layers
    def linear(self, inp, inner_size, name='linear', bias=True, activation = None, init = None):
        with tf.compat.v1.variable_scope(name):

            lin = tf.compat.v1.layers.dense(inp, inner_size, name=name, activation=activation, use_bias=bias,
                                            kernel_initializer=init)
            return lin

    def conv_layer_2d(self, input, filters, kernel_size, strides=(1, 1), padding="SAME", name='conv', activation=None,
                      bias = True):
        with tf.compat.v1.variable_scope(name):

            conv = tf.compat.v1.layers.conv2d(input, filters, kernel_size, strides, padding=padding, name=name,
                                              activation=activation, use_bias=bias)
            return conv

    def embedding(self, input, indices, size, name='embs'):
        with tf.compat.v1.variable_scope(name):
            shape = (indices, size)
            stddev = min(0.1, sqrt(2.0 / (product(xs=shape[:-1]) + shape[-1])))
            initializer = tf.random.normal(shape=shape, stddev=stddev, dtype=tf.float32)
            W = tf.Variable(
                initial_value=initializer, trainable=True, validate_shape=True, name='W',
                dtype=tf.float32, shape=shape
            )
            return tf.nn.tanh(tf.compat.v1.nn.embedding_lookup(params=W, ids=input, max_norm=None))

    # Netowrk specification
    def conv_net(self, global_state, local_state, local_two_state, agent_stats, target_stats, z_sigma=None,
                 use_noise=None, with_action=False):

        conv_10 = self.conv_layer_2d(global_state, 32, [1, 1], name='conv_10', activation=tf.nn.tanh)
        conv_11 = self.conv_layer_2d(conv_10, 32, [3, 3], name='conv_11', activation=tf.nn.leaky_relu)
        conv_12 = self.conv_layer_2d(conv_11, 32, [3, 3], name='conv_12', activation=tf.nn.leaky_relu)
        fc11 = tf.reshape(conv_12, [-1,10*10*32])

        embs_41 = tf.nn.tanh(self.embedding(agent_stats, 129, 32, name='embs_41'))
        embs_41 = tf.reshape(embs_41, [-1, 16 * 32])
        fc_41 = self.linear(embs_41, 100, name = 'fc_41', activation=tf.nn.leaky_relu)

        embs_51 = self.embedding(target_stats, 125, 32, name='embs_51')
        embs_51 = tf.reshape(embs_51, [-1, 15 * 32])
        fc_51 = self.linear(embs_51, 100, name = 'fc_51', activation = tf.nn.leaky_relu)

        all_flat = tf.concat([fc11, fc_41, fc_51], axis=1)

        all_flat = self.linear(all_flat, 32, name='fc1', activation=tf.nn.leaky_relu)

        if with_action:
            hot_acts = tf.one_hot(self.acts, self.actions_size)
            hot_acts = tf.reshape(hot_acts, [-1, self.actions_size])
            all_flat = tf.concat([all_flat, hot_acts], axis=1)

        z_mean = None
        fc2 = self.linear(all_flat, 32, name='fc2', activation = tf.nn.leaky_relu)

        # In case we want to use V-AIRL
        if self.use_vairl:

            z_mean = self.linear(fc2, 32, name='z_mean', init=tf.compat.v1.initializers.variance_scaling(0.01))
            noise = tf.compat.v1.random_normal(tf.compat.v1.shape(z_mean), dtype=tf.float32)

            z = z_mean + z_sigma * noise * use_noise
            fc2 = z

            return self.linear(fc2, 1, name='out'), z_mean
        else:
            return self.linear(fc2, 1, name='out'), None

    # Train method of the discriminator
    def train(self):

        losses = []

        # Update discriminator
        for it in range(self.num_itr):

            expert_batch_idxs = random.sample(range(len(self.expert_traj['obs'])), self.batch_size)
            policy_batch_idxs = random.sample(range(len(self.buffer['obs'])), self.batch_size)

            #expert_batch_idxs = np.random.randint(0, len(expert_traj['obs']), batch_size)
            #policy_batch_idxs = np.random.randint(0, len(policy_traj['obs']), batch_size)

            expert_obs = [self.expert_traj['obs'][id] for id in expert_batch_idxs]
            policy_obs = [self.buffer['obs'][id] for id in policy_batch_idxs]

            expert_obs_n = [self.expert_traj['obs_n'][id] for id in expert_batch_idxs]
            policy_obs_n = [self.buffer['obs_n'][id] for id in policy_batch_idxs]

            expert_acts = [self.expert_traj['acts'][id] for id in expert_batch_idxs]
            policy_acts = [self.buffer['acts'][id] for id in policy_batch_idxs]

            expert_probs = []
            for (index, state) in enumerate(expert_obs):
                _, probs = self.select_action(state)
                expert_probs.append(probs[expert_acts[index]])

            policy_probs = []
            for (index, state) in enumerate(policy_obs):
                _, probs = self.select_action(state)
                policy_probs.append(probs[policy_acts[index]])

            expert_probs = np.asarray(expert_probs)
            policy_probs = np.asarray(policy_probs)

            labels = np.ones((self.batch_size, 1))
            labels = np.concatenate([labels, np.zeros((self.batch_size, 1))])

            e_states = self.obs_to_state(expert_obs)
            p_states = self.obs_to_state(policy_obs)

            all_global = np.concatenate([e_states[0], p_states[0]], axis=0)
            all_local = np.concatenate([e_states[1], p_states[1]], axis=0)
            all_local_two = np.concatenate([e_states[2], p_states[2]], axis=0)
            all_agent_stats = np.concatenate([e_states[3], p_states[3]], axis=0)
            all_target_stats = np.concatenate([e_states[4], p_states[4]], axis=0)

            e_states_n = self.obs_to_state(expert_obs_n)
            p_states_n = self.obs_to_state(policy_obs_n)

            all_global_n = np.concatenate([e_states_n[0], p_states_n[0]], axis=0)
            all_local_n = np.concatenate([e_states_n[1], p_states_n[1]], axis=0)
            all_local_two_n = np.concatenate([e_states_n[2], p_states_n[2]], axis=0)
            all_agent_stats_n = np.concatenate([e_states_n[3], p_states_n[3]], axis=0)
            all_target_stats_n = np.concatenate([e_states_n[4], p_states_n[4]], axis=0)

            all_probs = np.concatenate([expert_probs, policy_probs], axis=0)
            all_probs = np.expand_dims(all_probs, axis=1)

            feed_dict = {

                self.local_state: all_local,
                self.local_two_state: all_local_two,
                self.agent_stats: all_agent_stats,
                self.target_stats: all_target_stats,

                self.local_state_n: all_local_n,
                self.local_two_state_n: all_local_two_n,
                self.agent_stats_n: all_agent_stats_n,
                self.target_stats_n: all_target_stats_n,

                self.probs: all_probs,
                self.labels: labels,

                self.use_noise: [1]
            }

            if self.with_action:
                all_acts = np.concatenate([expert_acts, policy_acts], axis=0)
                all_acts = np.expand_dims(all_acts, axis=1)

                feed_dict[self.acts] = all_acts

            feed_dict[self.global_state] = all_global
            feed_dict[self.global_state_n] = all_global_n

            if self.use_vairl:
                loss, f, _, _, kl_loss = self.sess.run([self.loss, self.f, self.step, self.update_beta, self.kl_loss],
                                                       feed_dict=feed_dict)
            else:
                loss, f, disc, _ = self.sess.run([self.loss, self.f, self.discriminator, self.step],
                                                       feed_dict=feed_dict)

            losses.append(loss)

        # Update nomralization parameters
        for it in range(self.num_itr):

            expert_batch_idxs = random.sample(range(len(self.expert_traj['obs'])), self.batch_size)
            policy_batch_idxs = random.sample(range(len(self.buffer['obs'])), self.batch_size)

            expert_obs = [self.expert_traj['obs'][id] for id in expert_batch_idxs]
            policy_obs = [self.buffer['obs'][id] for id in policy_batch_idxs]

            expert_obs_n = [self.expert_traj['obs_n'][id] for id in expert_batch_idxs]
            policy_obs_n = [self.buffer['obs_n'][id] for id in policy_batch_idxs]

            expert_acts = [self.expert_traj['acts'][id] for id in expert_batch_idxs]
            policy_acts = [self.buffer['acts'][id] for id in policy_batch_idxs]

            e_states = self.obs_to_state(expert_obs)
            p_states = self.obs_to_state(policy_obs)

            all_global = np.concatenate([e_states[0], p_states[0]], axis=0)
            all_local = np.concatenate([e_states[1], p_states[1]], axis=0)
            all_local_two = np.concatenate([e_states[2], p_states[2]], axis=0)
            all_agent_stats = np.concatenate([e_states[3], p_states[3]], axis=0)
            all_target_stats = np.concatenate([e_states[4], p_states[4]], axis=0)

            e_states_n = self.obs_to_state(expert_obs_n)
            p_states_n = self.obs_to_state(policy_obs_n)

            all_global_n = np.concatenate([e_states_n[0], p_states_n[0]], axis=0)
            all_local_n = np.concatenate([e_states_n[1], p_states_n[1]], axis=0)
            all_local_two_n = np.concatenate([e_states_n[2], p_states_n[2]], axis=0)
            all_agent_stats_n = np.concatenate([e_states_n[3], p_states_n[3]], axis=0)
            all_target_stats_n = np.concatenate([e_states_n[4], p_states_n[4]], axis=0)

            expert_probs = []
            for (index, state) in enumerate(expert_obs):
                _, probs = self.select_action(state)
                expert_probs.append(probs[expert_acts[index]])

            policy_probs = []
            for (index, state) in enumerate(policy_obs):
                _, probs = self.select_action(state)
                policy_probs.append(probs[policy_acts[index]])

            expert_probs = np.asarray(expert_probs)
            policy_probs = np.asarray(policy_probs)

            probs = np.concatenate([expert_probs, policy_probs], axis=0)
            probs = np.expand_dims(probs, axis=1)

            feed_dict = {

                self.global_state: all_global,
                self.local_state: all_local,
                self.local_two_state: all_local_two,
                self.agent_stats: all_agent_stats,
                self.target_stats: all_target_stats,

                self.global_state_n: all_global_n,
                self.local_state_n: all_local_n,
                self.local_two_state_n: all_local_two_n,
                self.agent_stats_n: all_agent_stats_n,
                self.target_stats_n: all_target_stats_n,
            }

            if self.use_vairl:
                feed_dict[self.use_noise] = [0]

            if self.with_action:
                all_acts = np.concatenate([expert_acts, policy_acts], axis=0)
                all_acts = np.expand_dims(all_acts, axis=1)

                feed_dict[self.acts] = all_acts

            feed_dict[self.global_state] = all_global
            feed_dict[self.global_state_n] = all_global_n

            f = self.sess.run([self.f], feed_dict=feed_dict)
            f -= self.entropy_weight*np.log(probs)
            f = np.squeeze(f)
            self.push_reward(f)

        # Update Dynamic Running Stat
        if isinstance(self.r_norm, DynamicRunningStat):
            self.r_norm.reset()

        return np.mean(losses), 0

    # Eval without discriminator - only reward function
    def eval(self, obs, obs_n, acts=None, probs=None):

        states = self.obs_to_state(obs)

        feed_dict = {
            self.global_state: states[0],
            self.local_state: states[1],
            self.local_two_state: states[2],
            self.agent_stats: states[3],
            self.target_stats: states[4],

            self.use_noise: [0]
        }

        if self.with_action and acts is not None:
            acts = np.expand_dims(acts, axis=1)
            feed_dict[self.acts] = acts

        reward = self.sess.run([self.reward], feed_dict=feed_dict)
        if probs != None:
            reward -= self.entropy_weight*np.log(probs)

        # Normalize the reward
        #self.r_norm.push(reward[0][0])
        #reward = [[self.normalize_rewards(reward[0][0])]]
        #if self.r_norm.n == 0:
        #    reward = [[0]]

        return reward[0][0]

    # Eval with the discriminator - it returns an entropy regularized objective
    def eval_discriminator(self, obs, obs_n, probs, acts=None):
        states = self.obs_to_state(obs)
        states_n = self.obs_to_state(obs_n)

        probs = np.expand_dims(probs, axis=1)

        feed_dict = {
            self.global_state: states[0],
            self.local_state: states[1],
            self.local_two_state: states[2],
            self.agent_stats: states[3],
            self.target_stats: states[4],

            self.global_state_n: states[0],
            self.local_state_n: states_n[1],
            self.local_two_state_n: states_n[2],
            self.agent_stats_n: states_n[3],
            self.target_stats_n: states_n[4],

            self.use_noise: [0]
        }

        if self.with_action and acts is not None:
            acts = np.expand_dims(acts, axis=1)
            feed_dict[self.acts] = acts

        f = self.sess.run([self.f], feed_dict=feed_dict)
        f -= self.entropy_weight*np.log(probs)
        f = self.normalize_rewards(f)
        return f

    # Transform a DeepCrawl obs to state
    def obs_to_state(self, obs):

        global_batch = np.stack([np.asarray(state['global_in']) for state in obs])
        local_batch = np.stack([np.asarray(state['local_in']) for state in obs])
        local_two_batch = np.stack([np.asarray(state['local_in_two']) for state in obs])
        agent_stats_batch = np.stack([np.asarray(state['agent_stats']) for state in obs])
        target_stats_batch = np.stack([np.asarray(state['target_stats']) for state in obs])

        return global_batch, local_batch, local_two_batch, agent_stats_batch, target_stats_batch

    # For V-AIRL
    def make_beta_update(self):

        new_beta = tf.maximum(
            self.beta + self.alpha * (self.kl_loss - self.mutual_information), eps
        )
        with tf.control_dependencies([self.step]):
            self.update_beta = tf.compat.v1.assign(self.beta, new_beta)

    # Normalize the reward for each frame of the sequence
    def push_reward(self, rewards):
        for r in rewards:
            self.r_norm.push(r)

    def normalize_rewards(self, rewards):
        rewards -= self.r_norm.mean
        rewards /= (self.r_norm.std + 1e-12)
        rewards *= 0.05

        return rewards

    # Select action from the policy and fetch the probability distribution over the action space
    def select_action(self, state):

        act, _, probs = self.policy.eval([state])

        return (act, probs[0])

    # Update demonstrations
    def set_demonstrations(self, demonstrations, validations):
        self.expert_traj = demonstrations

        if validations is not None:
            self.validation_traj = validations

    # Create the replay buffer with which train the discriminator
    def create_buffer(self):
        self.buffer['obs'] = []
        self.buffer['obs_n'] = []
        self.buffer['acts'] = []

    # Add a transition to the buffer
    def add_to_buffer(self, obs, obs_n, acts, buffer_length=100000):

        if len(self.buffer['obs']) >= buffer_length:
            random_index = np.random.randint(0, len(self.buffer['obs']))
            del self.buffer['obs'][random_index]
            del self.buffer['obs_n'][random_index]
            del self.buffer['acts'][random_index]

        self.buffer['obs'].append(obs)
        self.buffer['obs_n'].append(obs_n)
        self.buffer['acts'].append(acts)

    # Create and return some demonstrations [(states, actions, frames)]. The order of the demonstrations must be from
    # best to worst. The number of demonstrations is given by the user
    def create_demonstrations(self, env, save_demonstrations=True, inference=False, verbose=False,
                              with_policy=False, num_episode=31, max_timestep=20):
        end = False

        # Initialize trajectories buffer
        expert_traj = {
            'obs': [],
            'obs_n': [],
            'acts': [],
        }

        val_traj = {
            'obs': [],
            'obs_n': [],
            'acts': []
        }

        if with_policy is None:
            num_episode = None

        episode = 1

        while not end:
            # Make another demonstration
            print('Demonstration n° ' + str(episode))
            # Reset the environment
            state = env.reset()
            states = [state]
            actions = []
            done = False
            step = 0
            cum_reward = 0
            # New sequence of states and actions
            while not done:
                try:
                    # Input the action and save the new state and action
                    step += 1
                    print("Timestep: " + str(step))
                    if verbose:
                        env.print_observation(state)
                    if not with_policy:
                        action = input('action: ')
                        if action == "f":
                            done = True
                            continue
                        while env.command_to_action(action) >= 99:
                            action = input('action: ')
                    else:
                        action, probs = self.select_action(state)
                    print(action)
                    state_n, done, reward = env.execute(action)

                    cum_reward += reward
                    if not with_policy:
                        action = env.command_to_action(action)
                    # If inference is true, print the reward
                    if inference:
                        _, probs = self.select_action(state)
                        reward = self.eval([state], [state_n], [action])
                        # print('Discriminator probability: ' + str(disc))
                        print('Unnormalize reward: ' + str(reward))
                        reward = self.normalize_rewards(reward)
                        print('Normalize reward: ' + str(reward))
                        print('Probability of state space: ')
                        print(probs)
                    state = state_n
                    states.append(state)
                    actions.append(action)
                    if step >= max_timestep:
                        done = True
                except Exception as e:
                    print(e)
                    continue

            if not inference:
                y = None

                print('Demonstration number: ' + str(episode))
                if with_policy:
                    if episode < num_episode:
                        print(state_n['target_stats'][0])
                        if True:
                            y = 'y'
                        else:
                            y = 'n'
                while y != 'y' and y != 'n':
                    y = input('Do you want to save this demonstration? [y/n] ')

                if y == 'y':
                    # Update expert trajectories
                    expert_traj['obs'].extend(np.array(states[:-1]))
                    expert_traj['obs_n'].extend(np.array(states[1:]))
                    expert_traj['acts'].extend(np.array(actions))
                    episode += 1
                else:

                    if with_policy:
                        if episode > num_episode - 1:
                            y = input('Do you want to save this demonstration as validation? [y/n] ')
                        else:
                            y = 'n'
                    else:
                        y = input('Do you want to save this demonstration as validation? [y/n] ')

                    if y == 'y':
                        val_traj['obs'].extend(np.array(states[:-1]))
                        val_traj['obs_n'].extend(np.array(states[1:]))
                        val_traj['acts'].extend(np.array(actions))
                        episode += 1

            y = None
            if num_episode is None:
                while y != 'y' and y != 'n':

                    if not inference:
                        y = input('Do you want to create another demonstration? [y/n] ')
                    else:
                        y = input('Do you want to try another episode? [y/n] ')

                    if y == 'n':
                        end = True
            else:
                if episode >= num_episode + 1:
                    end = True

        if len(val_traj['obs']) <= 0:
            val_traj = None

        # Save demonstrations to file
        if save_demonstrations and not inference:
            print('Saving the demonstrations...')
            self.save_demonstrations(expert_traj, val_traj)
            print('Demonstrations saved!')

        if not inference:
            self.set_demonstrations(expert_traj, val_traj)

        return expert_traj, val_traj

    # Save demonstrations dict to file
    @staticmethod
    def save_demonstrations(demonstrations, validations=None, name='dems_potions.pkl'):
        with open('reward_model/dems/' + name, 'wb') as f:
            pickle.dump(demonstrations, f, pickle.HIGHEST_PROTOCOL)
        if validations is not None:
            with open('reward_model/dems/vals.pkl', 'wb') as f:
                pickle.dump(validations, f, pickle.HIGHEST_PROTOCOL)

    # Load demonstrations from file
    def load_demonstrations(self, name='dems.pkl'):
        with open('reward_model/dems/' + name, 'rb') as f:
            expert_traj = pickle.load(f)

        with open('reward_model/dems/vals.pkl', 'rb') as f:
            val_traj = pickle.load(f)

        self.set_demonstrations(expert_traj, val_traj)

        return expert_traj, val_traj

    # Save the entire model
    def save_model(self, name=None):
        self.saver.save(self.sess, 'reward_model/models/{}'.format(name))
        return

    # Load entire model
    def load_model(self, name=None):
        self.saver = tf.compat.v1.train.import_meta_graph('reward_model/models/' + name + '.meta')
        self.saver.restore(self.sess, 'reward_model/models/' + name)
        return