from __future__ import division
import scipy.sparse as sp 
import scipy.io
import numpy as np
import copy
from scipy import pi
import sys
import pandas as pd
from os.path import join
import gzip
from scipy.io import mmwrite,mmread
from sklearn.decomposition import PCA,TruncatedSVD
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.metrics import pairwise_distances
from sklearn.cluster import KMeans
from sklearn.metrics.cluster import normalized_mutual_info_score, adjusted_rand_score
from sklearn.metrics.cluster import homogeneity_score, adjusted_mutual_info_score
from sklearn.preprocessing import MinMaxScaler,MaxAbsScaler
import metric
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_style("whitegrid", {'axes.grid' : False})
matplotlib.rc('xtick', labelsize=20) 
matplotlib.rc('ytick', labelsize=20) 
matplotlib.rcParams.update({'font.size': 22})

def compute_inertia(a, X, norm=True):
    if norm:
        W = [np.sum(pairwise_distances(X[a == c, :]))/(2.*sum(a == c)) for c in np.unique(a)]
        return np.sum(W)
    else:
        W = [np.mean(pairwise_distances(X[a == c, :])) for c in np.unique(a)]
        return np.mean(W)

#gap statistic 
def compute_gap(clustering, data, k_max=10, n_references=100):
    if len(data.shape) == 1:
        data = data.reshape(-1, 1)
    reference_inertia = []
    for k in range(2, k_max+1):
        reference = np.random.rand(*data.shape)
        mins = np.min(data,axis=0)
        maxs = np.max(data,axis=0)
        reference = reference*(maxs-mins)+mins
        local_inertia = []
        for _ in range(n_references):
            clustering.n_clusters = k
            assignments = clustering.fit_predict(reference)
            local_inertia.append(compute_inertia(assignments, reference))
        reference_inertia.append(np.mean(local_inertia))
    
    ondata_inertia = []
    for k in range(2, k_max+1):
        clustering.n_clusters = k
        assignments = clustering.fit_predict(data)
        ondata_inertia.append(compute_inertia(assignments, data))
        
    gap = np.log(reference_inertia)-np.log(ondata_inertia)
    return gap, np.log(reference_inertia), np.log(ondata_inertia)



#scATAC data
class scATAC_Sampler(object):
    def __init__(self,name,dim=20,low=0.03,has_label=True):
        self.name = name
        self.dim = dim
        self.has_label = has_label
        X = pd.read_csv('datasets/%s/sc_mat.txt'%name,sep='\t',header=0,index_col=[0]).values
        if has_label:
            labels = [item.strip() for item in open('datasets/%s/label.txt'%name).readlines()]
            uniq_labels = list(np.unique(labels))
            Y = np.array([uniq_labels.index(item) for item in labels])
            #X,Y = self.filter_cells(X,Y,min_peaks=10)
            self.Y = Y
        X = self.filter_peaks(X,low)
        #TF-IDF transformation
        nfreqs = 1.0 * X / np.tile(np.sum(X,axis=0), (X.shape[0],1))
        X  = nfreqs * np.tile(np.log(1 + 1.0 * X.shape[1] / np.sum(X,axis=1)).reshape(-1,1), (1,X.shape[1]))
        X = X.T #(cells, peaks)
        X = MinMaxScaler().fit_transform(X)
        #PCA transformation
        pca = PCA(n_components=dim, random_state=3456).fit(X)
        X = pca.transform(X)
        self.X = X
        self.total_size = self.X.shape[0]


    def filter_peaks(self,X,ratio):
        ind = np.sum(X>0,axis=1) > X.shape[1]*ratio
        return X[ind,:]

    def filter_cells(self,X,Y,min_peaks):
        ind = np.sum(X>0,axis=0) > min_peaks
        return X[:,ind], Y[ind]

    def correlation(self,X,Y,heatmap=False):
        nb_classes = len(set(Y))
        print(nb_classes)
        km = KMeans(n_clusters=nb_classes,random_state=0).fit(X)
        label_kmeans = km.labels_
        purity = metric.compute_purity(label_kmeans, Y)
        nmi = normalized_mutual_info_score(Y, label_kmeans)
        ari = adjusted_rand_score(Y, label_kmeans)
        homogeneity = homogeneity_score(Y, label_kmeans)
        ami = adjusted_mutual_info_score(Y, label_kmeans)
        print('NMI = {}, ARI = {}, Purity = {},AMI = {}, Homogeneity = {}'.format(nmi,ari,purity,ami,homogeneity))
 
    def train(self, batch_size):
        indx = np.random.randint(low = 0, high = self.total_size, size = batch_size)

        if self.has_label:
            return self.X[indx, :], self.Y[indx]
        else:
            return self.X[indx, :]

    def load_all(self):
        if self.has_label:
            return self.X, self.Y
        else:
            return self.X 


#load data from 10x Genomic paired ARC technology
class ARC_Sampler(object):
    def __init__(self,name='PBMC10k',n_components=50,scale=10000,filter_feat=True,filter_cell=False,random_seed=1234,mode=1, \
        min_rna_c=0,max_rna_c=None,min_atac_c=0,max_atac_c=None):
        #c:cell, g:gene, l:locus
        self.name = name
        self.mode = mode
        self.min_rna_c = min_rna_c
        self.max_rna_c = max_rna_c
        self.min_atac_c = min_atac_c
        self.max_atac_c = max_atac_c

        self.rna_mat, self.atac_mat, self.genes, self.peaks = self.load_data(filter_feat,filter_cell)

        self.rna_mat = self.rna_mat*scale/self.rna_mat.sum(1)
        self.rna_mat = np.log10(self.rna_mat+1)
        self.atac_mat = self.atac_mat*scale/self.atac_mat.sum(1)
        self.atac_mat = np.log10(self.atac_mat+1)

        self.rna_reducer = PCA(n_components=n_components, random_state=random_seed)
        self.rna_reducer.fit(self.rna_mat)
        self.pca_rna_mat = self.rna_reducer.transform(self.rna_mat)

        self.atac_reducer = PCA(n_components=n_components, random_state=random_seed)
        self.atac_reducer.fit(self.atac_mat)
        self.pca_atac_mat = self.atac_reducer.transform(self.atac_mat)

    def load_data(self,filter_feat,filter_cell):
        mtx_file = 'datasets/%s/matrix.mtx'%self.name
        feat_file = 'datasets/%s/features.tsv'%self.name
        barcode_file = 'datasets/%s/barcodes.tsv'%self.name
        combined_mat = mmread(mtx_file).T.tocsr() #(cells, genes+peaks)
        cells = [item.strip() for item in open(barcode_file).readlines()] 
        genes = [item.split('\t')[1] for item in open(feat_file).readlines() if item.split('\t')[2]=="Gene Expression"]
        peaks = [item.split('\t')[1] for item in open(feat_file).readlines() if item.split('\t')[2]=="Peaks"]
        assert len(genes)+len(peaks) == combined_mat.shape[1]
        rna_mat = combined_mat[:,:len(genes)]
        atac_mat = combined_mat[:,len(genes):]
        print('scRNA-seq: ', rna_mat.shape, 'scATAC-seq: ', atac_mat.shape)

        if filter_feat:
            rna_mat, atac_mat, genes, peaks = self.filter_feats(rna_mat, atac_mat, genes, peaks)
            print('scRNA-seq filtered: ', rna_mat.shape, 'scATAC-seq filtered: ', atac_mat.shape)
        return rna_mat, atac_mat, genes, peaks

    def filter_feats(self, rna_mat_sp, atac_mat_sp, genes, peaks):
        #filter genes
        gene_select = np.array((rna_mat_sp>0).sum(axis=0)).squeeze() > self.min_rna_c
        if self.max_rna_c is not None:
            gene_select *= np.array((rna_mat_sp>0).sum(axis=0)).squeeze() < self.max_rna_c
        rna_mat_sp = rna_mat_sp[:,gene_select]
        genes = np.array(genes)[gene_select]
        #filter peaks
        locus_select = np.array((atac_mat_sp>0).sum(axis=0)).squeeze() > self.min_atac_c
        if self.max_atac_c is not None:
            locus_select *= np.array((atac_mat_sp>0).sum(axis=0)).squeeze() < self.max_atac_c
        atac_mat_sp = atac_mat_sp[:,locus_select]
        peaks = np.array(peaks)[locus_select]
        return rna_mat_sp, atac_mat_sp, genes, peaks

    def get_batch(self,batch_size):
        idx = np.random.randint(low = 0, high = self.atac_mat.shape[0], size = batch_size)
        if self.mode == 1:
            return self.pca_rna_mat[idx,:]
        elif self.mode == 2:
            return self.pca_atac_mat[idx,:]
        elif self.mode == 3:
            return np.hstack((self.pca_rna_mat, self.pca_atac_mat))[idx,:]
        else:
            print('Wrong mode!')
            sys.exit()
    
    def load_all(self):
        #mode: 1 only scRNA-seq, 2 only scATAC-seq, 3 both
        if self.mode == 1:
            return self.pca_rna_mat
        elif self.mode == 2:
            return self.pca_atac_mat
        elif self.mode == 3:
            return np.hstack((self.pca_rna_mat, self.pca_atac_mat))
        else:
            print('Wrong mode!')
            sys.exit()


#sample continuous (Gaussian) and discrete (Catagory) latent variables together
class Mixture_sampler(object):
    def __init__(self, nb_classes, N, dim, sd, scale=1):
        self.nb_classes = nb_classes
        self.total_size = N
        self.dim = dim
        self.sd = sd 
        self.scale = scale
        np.random.seed(1024)
        self.X_c = self.scale*np.random.normal(0, self.sd**2, (self.total_size,self.dim))
        #self.X_c = self.scale*np.random.uniform(-1, 1, (self.total_size,self.dim))
        self.label_idx = np.random.randint(low = 0 , high = self.nb_classes, size = self.total_size)
        self.X_d = np.eye(self.nb_classes)[self.label_idx]
        self.X = np.hstack((self.X_c,self.X_d))
    
    def train(self,batch_size,weights=None):
        X_batch_c = self.scale*np.random.normal(0, 1, (batch_size,self.dim))
        #X_batch_c = self.scale*np.random.uniform(-1, 1, (batch_size,self.dim))
        if weights is None:
            weights = np.ones(self.nb_classes, dtype=np.float64) / float(self.nb_classes)
        label_batch_idx =  np.random.choice(self.nb_classes, size=batch_size, replace=True, p=weights)
        X_batch_d = np.eye(self.nb_classes)[label_batch_idx]
        return X_batch_c, X_batch_d

    def load_all(self):
        return self.X_c, self.X_d

#sample continuous (Gaussian Mixture) and discrete (Catagory) latent variables together
class Mixture_sampler_v2(object):
    def __init__(self, nb_classes, N, dim, weights=None,sd=0.5):
        self.nb_classes = nb_classes
        self.total_size = N
        self.dim = dim
        np.random.seed(1024)
        if nb_classes<=dim:
            self.mean = np.random.uniform(-5,5,size =(nb_classes, dim))
            #self.mean = np.zeros((nb_classes,dim))
            #self.mean[:,:nb_classes] = np.eye(nb_classes)
        else:
            if dim==2:
                self.mean = np.array([(np.cos(2*np.pi*idx/float(self.nb_classes)),np.sin(2*np.pi*idx/float(self.nb_classes))) for idx in range(self.nb_classes)])
            else:
                self.mean = np.zeros((nb_classes,dim))
                self.mean[:,:2] = np.array([(np.cos(2*np.pi*idx/float(self.nb_classes)),np.sin(2*np.pi*idx/float(self.nb_classes))) for idx in range(self.nb_classes)])
        self.cov = [sd**2*np.eye(dim) for item in range(nb_classes)]
        if weights is None:
            weights = np.ones(self.nb_classes, dtype=np.float64) / float(self.nb_classes)
        self.Y = np.random.choice(self.nb_classes, size=N, replace=True, p=weights)
        self.X_c = np.array([np.random.multivariate_normal(mean=self.mean[i],cov=self.cov[i]) for i in self.Y],dtype='float64')
        self.X_d = np.eye(self.nb_classes)[self.Y]
        self.X = np.hstack((self.X_c,self.X_d))

    def train(self, batch_size, label = False):
        indx = np.random.randint(low = 0, high = self.total_size, size = batch_size)
        if label:
            return self.X_c[indx, :], self.X_d[indx, :], self.Y[indx, :]
        else:
            return self.X_c[indx, :], self.X_d[indx, :]

    def get_batch(self,batch_size,weights=None):
        if weights is None:
            weights = np.ones(self.nb_classes, dtype=np.float64) / float(self.nb_classes)
        label_batch_idx =  np.random.choice(self.nb_classes, size=batch_size, replace=True, p=weights)
        return self.X_c[label_batch_idx, :], self.X_d[label_batch_idx, :]
    def predict_onepoint(self,array):#return component index with max likelyhood
        from scipy.stats import multivariate_normal
        assert len(array) == self.dim
        return np.argmax([multivariate_normal.pdf(array,self.mean[idx],self.cov[idx]) for idx in range(self.nb_classes)])

    def predict_multipoints(self,arrays):
        assert arrays.shape[-1] == self.dim
        return map(self.predict_onepoint,arrays)
    def load_all(self):
        return self.X_c, self.X_d, self.label_idx

#get a batch of data from previous 50 batches, add stochastic
class DataPool(object):
    def __init__(self, maxsize=50):
        self.maxsize = maxsize
        self.nb_batch = 0
        self.pool = []

    def __call__(self, data):
        if self.nb_batch < self.maxsize:
            self.pool.append(data)
            self.nb_batch += 1
            return data
        if np.random.rand() > 0.5:
            results=[]
            for i in range(len(data)):
                idx = int(np.random.rand()*self.maxsize)
                results.append(copy.copy(self.pool[idx])[i])
                self.pool[idx][i] = data[i]
            return results
        else:
            return data