import numpy as np
import math
import random
import os
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from tree import Tree

__outputDir = None
__verbose = None


def init(outputDir='./RNAscrutiny-output',
         seed=np.random.randint(1e8),
         verbose=False):

    global __outputDir
    global __verbose

    # set output directory
    __outputDir = outputDir
    try:
        os.makedirs(outputDir)
    except:
        raise Exception('Unable to create output directory.')

    # seed random number generators
    np.random.seed(seed)
    random.seed(seed)

    # set verbosity
    __verbose = verbose


def generateGeneCorrelationMatrix(n,
                                  networkStructure='SimulatedPowerLaw',
                                  maxAlphaBeta=10,
                                  normalizeWRows=False,
                                  powerLawExponent=2.05,
                                  networkSparsity=0.999):
    """Generate correlation matrix of gene regulatory network.

    Parameters
    ----------
    n : int
        number of genes
    maxAlphaBeta : float
        Maximum value of the beta distribution alpha and beta parameters.
    normalizeWRows : logical
        Choose whether to divide each row of :math:`W` by the number of connections.
    networkStructure : string
        Type of network structure, must be one of 'SimulatedPowerLaw', 'TRRUST', 'Random'.
    powerLawExponent : float
        Exponent for simulated power law network, higher values imply more connections, recommended value between 2 and 3.
    networkSparsity : float
        Fraction of genes unrelated in network.

    Returns
    -------
    gc : ndarray(n,n)
        Gene correlation matrix, entry i,j is the correlation between genes i and j.
    tf : ndarray(*,)
        List of transcription factors. A transcription factor is defined as a gene with a nonzero outdegree in the gene network.
    """

    # adjacency matrix of gene network
    networkMatrix = np.zeros(shape=(n, n))

    # Generate the network structure
    if networkStructure == "SimulatedPowerLaw":

        #Generate number of in and out connections per gene
        inConnections = np.empty(n, dtype=int)
        outConnections = np.empty(n, dtype=int)
        for i in range(n):
            inConnections[i] = int(0.5 * (1 - np.random.uniform())**
                                   (-1 / (powerLawExponent - 1)) + 0.5)
            outConnections[i] = int(0.5 * (1 - np.random.uniform())**
                                    (-1 / (powerLawExponent - 1)) + 0.5)

        # Make sure number of in and out connections is not greater than the number of genes
        inConnections[inConnections > n] = n
        outConnections[outConnections > n] = n

        # Set up a multiset of the number of outconnections each gene has
        # then, for each gene, sample the inconnections from
        # the multiset of outconnections and remove duplicates
        outNodesSamplePool = []
        for col in range(n):
            outNodesSamplePool += [col] * outConnections[col]

        for row in range(n):
            connectionIndexes = random.sample(outNodesSamplePool,
                                              inConnections[row])
            connectionIndexes = list(set(connectionIndexes))
            networkMatrix[row, connectionIndexes] = 1

    elif networkStructure == "TRRUST":
        # This can be generated by using AnalyzeTRRUST.py and the original file
        networkMatrix = np.load("TRRUST_Network.npy")
        n = networkMatrix.shape[0]

    elif networkStructure == "Random":
        #Randomly generate network using sparsity and binomial distribution for the number of in connections per gene
        for row in range(n):
            numInConnections = np.random.binomial(n, (1 - networkSparsity))
            connectionIndexes = random.sample(range(n), numInConnections)
            networkMatrix[row, connectionIndexes] = 1

    else:
        raise Exception(
            'Invalid network structure. Parameter "networkStructure" must be one of "SimulatedPowerLaw", "TRRUST", or "Random".'
        )

    valueMatrix = np.zeros(shape=(n, n))

    # Generate values for matrix from a Beta distribution using randomly generated
    # alpha and beta in the range [0, 1, 2, ... , maxAlphaBeta],
    # then transform the actual values from the (0, 1) range to the (-1, 1) range
    for col in range(n):
        alpha = np.random.randint(1, maxAlphaBeta + 1)
        beta = np.random.randint(1, maxAlphaBeta + 1)
        valueMatrix[col, :] = np.random.choice([-1.0, 1.0]) * np.random.beta(
            a=alpha, b=beta, size=n)

    # Element-wise multiply network structure matrix and value matrix
    gc = networkMatrix * valueMatrix

    # If set to true, will divide each value of each row (in connections) by the number of connections
    if normalizeWRows:
        rowCount = np.count_nonzero(gc, axis=1)
        rowCount[rowCount == 0] = 1
        gc /= rowCount[:, np.newaxis]

    # Transcription Factors are defined as genes with a nonzero outdegree
    outDegrees = np.sum(networkMatrix, axis=0)
    tf = np.nonzero(outDegrees)[0]

    if __outputDir is not None:

        inDegrees = np.sum(networkMatrix, axis=1)
        outDegrees = np.sum(networkMatrix, axis=0)

        #inDegreesXVals, inDegreesFreq = np.unique(inDegrees, return_counts=True)
        #outDegreesXVals, outDegreesFreq = np.unique(outDegrees, return_counts=True)

        fig = plt.figure()
        ax = fig.gca()
        ax.hist(inDegrees, bins=100, range=[0, 100])
        ax.set_xlabel("Indegree")
        ax.set_ylabel("Frequency")
        plt.savefig(
            os.path.join(__outputDir, networkStructure +
                         "_Network_Indegree.tiff"))
        plt.close()

        fig = plt.figure()
        ax = fig.gca()
        ax.hist(outDegrees, bins=100, range=[0, 100])
        ax.set_xlabel("Outdegree")
        ax.set_ylabel("Frequency")
        plt.savefig(
            os.path.join(__outputDir, networkStructure +
                         "_Network_Outdegree.tiff"))
        plt.close()

    if __verbose:
        print("Number of Connections: ", np.count_nonzero(networkMatrix))
        sparsityProportion = float(np.count_nonzero(networkMatrix)) / n**2
        print("Network Density: ", sparsityProportion)

    return gc, tf


def formCellTypeTree(n, cells, cellTypeParents, cellTypeNumCells,
                     cellTypeConstProps, transcriptionFactors, cellTypeMeans):
    """
    Parameters
    ----------
    n : int
        Number of genes.
    cells : int
        Number of cells.
    cellTypeParents :
    cellTypeNumCells :
    cellTypeConstProps :
    transcriptionFactors :
    cellTypeMeans :

    Returns
    -------
	cellTypeTree :
    projMatrix :
    constMatrix :
    """
    cellTypeTree = Tree()

    parentCellType = -1
    projInit = np.ones(shape=n)
    constInit = np.zeros(shape=n)
    cellTypeMember = []
    cellTypeChildMember = list(range(cells))
    cellTypeMean = 0
    numCellTypes = len(cellTypeParents)

    projMatrix = np.zeros(shape=(numCellTypes, n))
    constMatrix = np.zeros(shape=(numCellTypes, n))

    headNode = cellTypeTree.add_head(parentCellType, cellTypeMember, constInit,
                                     projInit, cellTypeMean)
    headNode.setChildCellTypeMembers(cellTypeChildMember)

    numCellTypes = len(cellTypeParents)
    cellTypeMembers = [[] for i in range(numCellTypes)]

    cellNumPool = range(cells)
    for cellType in range(numCellTypes):
        numCells = cellTypeNumCells[cellType]
        cellTypeIndices = random.sample(cellNumPool, numCells)
        cellTypeMembers[cellType] = cellTypeIndices
        cellNumPool = list(set(cellNumPool) - set(cellTypeIndices))

    addTreeChildren(n, cellTypeTree, cellTypeParents, cellTypeMembers,
                    cellTypeConstProps, projInit, constInit, projMatrix,
                    constMatrix, cellTypeMeans, parentCellType,
                    transcriptionFactors)

    #maxCellType = int(max(cellTypeParents)) #Highest cell type number corresponds to the highest cell type
    # for node in cellTypeTree.traverse(-1):
    #     for field in node:
    #         print(field)

    return cellTypeTree, projMatrix, constMatrix


def addTreeChildren(n, cellTypeTree, cellTypeParents, cellTypeMembers,
                    cellTypeConstProps, projInit, constInit, projMatrix,
                    constMatrix, cellTypeMeans, parentCellType,
                    transcriptionFactors):
    """
    Add children to tree

    Parameters
    ----------
    n : int
        Number of genes
    cellTypeTree : tree
        A tree representation of the cell types
    cellTypeParents :
    cellTypeMembers :
    cellTypeConstProps :
    projInit :
    constInit :
    projMatrix :
    constMatrix :
    cellTypeMeans :
    parentCellType :
    transcriptionFactors :

    Returns
    -------
    list
    """
    childCellTypes = np.where(np.array(cellTypeParents) == parentCellType)[0]
    parentCellTypeMember = []
    for i in range(len(childCellTypes)):
        childCellType = childCellTypes[i]

        constProp = cellTypeConstProps[childCellType]
        proj, const = getNextProjConst(n, projInit, constInit, constProp,
                                       transcriptionFactors)
        cellTypeMean = cellTypeMeans[childCellType]
        cellTypeMember = cellTypeMembers[childCellType]

        projMatrix[childCellType, :] = proj
        constMatrix[childCellType, :] = const

        cellTypeNode = cellTypeTree.add_node(childCellType, cellTypeMember,
                                             const, proj, cellTypeMean,
                                             parentCellType)

        childCellTypeMember = addTreeChildren(
            n, cellTypeTree, cellTypeParents, cellTypeMembers,
            cellTypeConstProps, proj, const, projMatrix, constMatrix,
            cellTypeMeans, childCellType, transcriptionFactors)
        parentCellTypeMember += childCellTypeMember
    cellTypeTree[parentCellType].setChildCellTypeMembers(parentCellTypeMember)

    return parentCellTypeMember + cellTypeMembers[parentCellType]


def getNextProjConst(n, projInit, constInit, constProp, transcriptionFactors):
    """
    Generate projection and constant vectors for next cell type


    Parameters
    ----------
    n : int
        number of genes
    projInit : ndarray
    constInit : ndarray
    constProp : float
    transcriptionFactors : ndarray

    Returns
    -------
    proj : ndarray
    const : ndarray

    """
    numZeros = int(len(transcriptionFactors) *
                   constProp)  #np.random.binomial(n, constProp, 1)[0]
    zeroGenes = random.sample(
        np.ndarray.tolist(transcriptionFactors), numZeros)

    proj = np.copy(projInit)
    proj[zeroGenes] = 0

    const = np.copy(constInit)
    for i in range(0, n):
        if (proj[i] == 0
                and projInit[i] != 0):  #Only if cell is zero in only new proj
            const[i] = rangeTransformation(random.randint(0, 1))
    return proj, const


def generateInitialStates(n, cells, sameInitialCell=True, geneMeanLevels=0):
    """
    Randomly generate an initial gene expression state, and use it for all the cells.

    Parameters
    ----------
    n : int
        Number of genes
    cells : int
        Number of cells
    sameInitialCell : logical
        If true, use same initial state for each cell, otherwise generate different initial state for each cell.

    Returns
    -------
    initialStates : ndarray()
        initial states of cells expression levels
    """
    initialStates = np.zeros(shape=(n, cells))
    baseInitialState = (2.0 * np.random.rand(n) - 1.0) + geneMeanLevels
    for j in range(cells):
        if sameInitialCell:
            initialStates[:, j] = baseInitialState
        else:
            initialStates[:,
                          j] = (2.0 * np.random.rand(n) - 1.0) + geneMeanLevels
    return initialStates


def findOptimalAlpha(n,
                     gc,
                     cellTypeTree,
                     currentStates,
                     numToSample,
                     timeStep=0.01,
                     convergenceThreshold=0.1,
                     noiseDisp=0.05,
                     maxNumIterations=500,
                     convDistAvgNum=20,
                     initialAlpha=0.01,
                     alpha_step=0.02,
                     showAlphaGeneMeans=False,
                     geneMeanLevels=0):
    """
    Find alpha based on optimal development for progenitor cell type (parent = -1)

    Paramters
    ---------
    n : int
        Number of genes
    gc : ndarray(n,n)
        Gene expression correlation matrix
    cellTypeTree : tree
        cell types represented as tree
    currentStates : ndarray()
        current states of cells' expression levels
    numToSample : int
        number of cells to sample, must be less than or equal to number of cells
    timeStep : float
    convergenceThreshold : float
    noiseDisp : float
    maxNumIterations : int
        maximum numbe of iterations
    convDistAvgNum : float
    initialAlpha : float
    alpha_step : float
    showAlphaGeneMeans : logical

    Returns
    -------
    alpha : float
    """

    #Initialize variables (not to be changed)
    cells = currentStates.shape[1]
    successRateNeeded = 1
    numSuccessesNeeded = int(numToSample * successRateNeeded)
    successRate = 0
    if (initialAlpha >= alpha_step):
        alpha = initialAlpha - alpha_step
    else:
        alpha = 0
    avgIterations = 0
    normDiffVals = [0] * maxNumIterations

    treeHead = cellTypeTree.head()
    if (len(treeHead.children) == 1):  #If only one cell lineage, use head
        origCellType = cellTypeTree[(treeHead.children)[0]]
        proj = origCellType.proj
        const = origCellType.const
    else:  #Multiple cell lineages
        proj = np.zeros(shape=n)
        const = np.zeros(shape=n)

    #While proportion of cells that have converged is less than the required proportion
    while (successRate < successRateNeeded):
        geneMeans = np.zeros(shape=n)

        #Sample cells and save them in a new array so they don't overwrite old values
        cellsToSample = random.sample(range(cells), numToSample)
        simulatedStates = np.zeros(shape=(n, numToSample))
        for i in range(numToSample):
            simulatedStates[:, i] = currentStates[:, cellsToSample[i]]

        #Initialize these parameters at the start of every run of each alpha
        numSuccesses = 0
        if (successRate >
                0):  #Slow down alpha_step once at least one cell converges
            alpha_step = 0.01
        alpha = alpha + alpha_step
        avgIterations = 0

        #For each cell being sampled
        for i in range(numToSample):
            initialState = simulatedStates[:, i]
            cellIteration = 0
            normDiffVals = [0] * maxNumIterations
            convDist = float("inf")
            gc_alpha = gc / alpha

            #Develop each cell until the maximum number of iterations is reached or until it converges
            while (cellIteration < maxNumIterations
                   and convDist > convergenceThreshold):
                currentState = developCell(n, gc_alpha, timeStep, initialState,
                                           proj, const, noiseDisp,
                                           geneMeanLevels)
                normDiffVals[cellIteration] = np.mean(
                    np.absolute(currentState - initialState) / timeStep
                )  #np.linalg.norm((currentState - initialState)) / n #vector norm

                #Average average value of derivative over previous "convDistAvgNum" iterations
                if (cellIteration >= (convDistAvgNum - 1)):
                    convDist = np.mean(
                        normDiffVals[(cellIteration - convDistAvgNum + 1):(
                            cellIteration + 1)])
                initialState = currentState
                cellIteration += 1
            if (cellIteration != maxNumIterations):  #If the cell has converged
                numSuccesses += 1
            avgIterations += cellIteration
            geneMeans += currentState
        geneMeans /= numToSample
        if (showAlphaGeneMeans):
            plt.hist(geneMeans)
            plt.title("Gene Means for alpha: " + str(alpha))
            plt.show()
        successRate = numSuccesses / numToSample
        avgIterations = int(avgIterations / numToSample)
        avgIterations = avgIterations - convDistAvgNum  #Subtract iterations where it is below convergence distance
        # print("Alpha: ", alpha)
        # print("Success Rate: ", successRate)
        # print("Average Iterations: ", avgIterations)
    #xVals = [x for x in range(0, cellIteration)]
    #plt.scatter(xVals, normDiffVals[0:cellIteration])
    #plt.plot(xVals, normDiffVals[0:cellIteration])
    #plt.show()
    return alpha


def developCell(n, gc, timeStep, initialState, proj, const, noiseDisp,
                geneMeanLevels):
    """Equation to develop each cell for a time step

    Parameters
    ----------
    n : int
        Number of genes
    gc : ndarray(n,n)
        gene correlation matrix
    timeStep : float
        change in time for consecutive iterations
    initialState : ndarray()
        initial expression levels of the cells
    proj :
    const :
    noiseDisp :

    Returns
    -------
    ndarray()

    """
    tau = np.ones(shape=n)
    nextState = np.tanh(
        gc.dot(initialState - geneMeanLevels) + np.random.normal(
            loc=0, scale=noiseDisp, size=(n)))
    derivative = (nextState - initialState) / tau
    currentState = initialState + timeStep * derivative
    currentState = np.multiply(proj, currentState) + const + geneMeanLevels

    return currentState


def findAllNextStates(gc,
                      cellTypeNode,
                      cellTypeTree,
                      currentStates,
                      cellTypesRecord,
                      totalCellIterations,
                      timeStep=0.01,
                      noiseDisp=0.05,
                      convergenceThreshold=0.1,
                      maxNumIterations=500,
                      convDistAvgNum=20,
                      numToSample=50,
                      showPlots=False,
                      cellDevelopmentMode=True,
                      geneMeanLevels=0):
    n = gc.shape[0]
    currentCellType = cellTypeNode.identifier
    cellTypeMembers = cellTypeNode.cellTypeMembers
    proj = cellTypeNode.proj
    const = cellTypeNode.const
    childCellTypeMembers = cellTypeNode.childCellTypeMembers

    cellTypesRecord[cellTypeMembers] = int(currentCellType)
    # print("Current Cell Type: ", currentCellType)

    if len(cellTypeMembers) == 0:
        cellTypeIterations = 0
    else:
        #print("Constant Gene Proportion: ", 1 - float(np.count_nonzero(proj))/np.size(proj))
        cellTypeIterations = findOptimalIterations(
            n, gc, timeStep, cellTypeMembers, proj, const, geneMeanLevels,
            convergenceThreshold, noiseDisp, currentStates, maxNumIterations,
            convDistAvgNum, numToSample)
        for i in range(len(cellTypeMembers)):
            if cellDevelopmentMode:
                cellIterations = np.random.randint(0, cellTypeIterations + 1)
            else:
                cellIterations = np.random.poisson(cellTypeIterations)
            cellTypeMember = cellTypeMembers[i]
            initialState = currentStates[:, cellTypeMember]
            currentStates[:, cellTypeMember] = findCellNextState(
                gc, n, timeStep, i, proj, const, initialState, cellIterations,
                noiseDisp, geneMeanLevels, showPlots)
            totalCellIterations[cellTypeMember] += cellIterations
        for i in range(len(childCellTypeMembers)):
            cellIterations = cellTypeIterations
            childCellTypeMember = childCellTypeMembers[i]
            initialState = currentStates[:, childCellTypeMember]
            currentStates[:, childCellTypeMember] = findCellNextState(
                gc, n, timeStep, i, proj, const, initialState, cellIterations,
                noiseDisp, geneMeanLevels, showPlots)
            totalCellIterations[childCellTypeMember] += cellIterations

    cellTypeNodeChildrenIDs = cellTypeNode.children
    for cellTypeNodeChildID in cellTypeNodeChildrenIDs:
        cellTypeNodeChild = cellTypeTree[cellTypeNodeChildID]

        findAllNextStates(gc, cellTypeNodeChild, cellTypeTree, currentStates,
                          cellTypesRecord, totalCellIterations)


#Given a W matrix and the projection/constant vectors, find the minimum number of iterations until convergence
def findOptimalIterations(n, gc, timeStep, cellTypeMembers, proj, const,
                          geneMeanLevels, convergenceThreshold, noiseDisp,
                          currentStates, maxNumIterations, convDistAvgNum,
                          numToSample):

    #Initialize variables (not to be changed)
    avgIterations = 0
    normDiffVals = [0] * maxNumIterations

    #Sample cells and save them in a new array so they don't overwrite old values
    colsToSample = random.sample(cellTypeMembers, numToSample)
    simulatedStates = np.zeros(shape=(n, numToSample))
    for i in range(numToSample):
        simulatedStates[:, i] = currentStates[:, colsToSample[i]]

    #For each cell being sampled
    for i in range(numToSample):
        initialState = simulatedStates[:, i]
        cellIteration = 0
        normDiffVals = [0] * maxNumIterations
        convDist = float("inf")

        #Develop each cell until the maximum number of iterations is reached or until it converges
        while (cellIteration < maxNumIterations
               and convDist > convergenceThreshold):
            currentState = developCell(n, gc, timeStep, initialState, proj,
                                       const, noiseDisp, geneMeanLevels)
            normDiffVals[cellIteration] = np.mean(
                np.absolute(currentState - initialState) / timeStep
            )  #np.linalg.norm((currentState - initialState)) / n #vector norm

            #Average average value of derivative over previous "convDistAvgNum" iterations
            if (cellIteration >= (convDistAvgNum - 1)):
                convDist = np.mean(normDiffVals[(
                    cellIteration - convDistAvgNum + 1):(cellIteration + 1)])
            initialState = currentState
            cellIteration += 1
        avgIterations += cellIteration
    avgIterations = int(avgIterations / numToSample)
    avgIterations = avgIterations - convDistAvgNum  #Subtract iterations where it is below convergence distance
    # print("Average Iterations: ", avgIterations)
    #xVals = [x for x in range(0, cellIteration)]
    #plt.scatter(xVals, normDiffVals[0:cellIteration])
    #plt.plot(xVals, normDiffVals[0:cellIteration])
    #plt.show()
    return (avgIterations)


def findCellNextState(gc, n, timeStep, i, proj, const, initialState,
                      cellIterations, noiseDisp, geneMeanLevels, showPlots):
    normDiffVals = [0] * cellIterations
    currentState = initialState

    gene = int(np.random.uniform() * n)
    geneExpr = np.zeros(cellIterations)
    for x in range(0, cellIterations):
        yVals = initialState
        geneExpr[x] = initialState[gene]
        currentState = developCell(n, gc, timeStep, initialState, proj, const,
                                   noiseDisp, geneMeanLevels)
        normDiffVals[x] = np.mean(
            np.absolute(currentState - initialState
                        ))  #np.linalg.norm((currentState - initialState))
        initialState = currentState

    if (showPlots):
        xVals = [x for x in range(0, cellIterations)]

        plt.title("Convergence Graph")
        plt.plot(xVals, normDiffVals)
        plt.show()

        title = "Gene " + str(gene) + " Over Time"
        plt.scatter(xVals, geneExpr)
        plt.title(title)
        plt.plot(xVals, geneExpr)
        plt.show()

        plt.title("Current Distribution of Gene Expression")
        plt.hist(yVals, bins=int(n / 25))
        plt.show()

    return (currentState)


def analyzeDataBeforeTransformation(n,
                                    cells,
                                    finalCellStates,
                                    alpha,
                                    pseudotimes,
                                    networkStructure,
                                    saveFolder=None):

    randomGenes = np.random.randint(0, n, size=5)

    #What do you want to name the graphs?
    testingAlpha = False
    testingNetworkStructure = True

    if (testingAlpha):
        testingString = "_Alpha_" + str(alpha)  #What is being tested
    elif (testingNetworkStructure):
        testingString = "_" + networkStructure + "_Network"  #What is being tested
    else:
        testingString = ""

    if saveFolder is not None:
        fig = plt.figure()
        ax = plt.gca()
        SGeneMeans = np.mean(finalCellStates, axis=1)
        geneMeansString = os.path.join(
            saveFolder, "S_Gene_Means" + testingString + ".tiff")
        ax.hist(SGeneMeans, bins=int(n / 50))
        ax.set_xlabel("S Gene Means")
        ax.set_ylabel("Frequency")
        plt.savefig(geneMeansString)
        plt.close()

        fig = plt.figure()
        ax = plt.gca()
        SGeneVars = np.var(finalCellStates, axis=1)
        geneVarsString = os.path.join(saveFolder,
                                      "S_Gene_Vars" + testingString + ".tiff")
        ax.hist(SGeneVars, bins=int(n / 50))
        ax.set_xlabel("S Gene Variances")
        ax.set_ylabel("Frequency")
        plt.savefig(geneVarsString)
        plt.close()

        fig = plt.figure()
        ax = plt.gca()
        SCellMeans = np.mean(finalCellStates, axis=0)
        cellMeansString = os.path.join(
            saveFolder, "S_Cell_Means" + testingString + ".tiff")
        ax.hist(SCellMeans, bins=int(cells / 10))
        ax.set_xlabel("S Cell Means")
        ax.set_ylabel("Frequency")
        plt.savefig(cellMeansString)
        plt.close()

        # print("Random Genes Being Graphed: ", randomGenes)

        for randomGene in randomGenes:
            geneExpr = finalCellStates[randomGene, :]
            fig = plt.figure()
            ax = plt.gca()
            genePseudotimeString = os.path.join(
                saveFolder, "Gene_" + str(randomGene) + "_Pseudotime" +
                testingString + ".tiff")
            ax.scatter(pseudotimes, geneExpr)
            ax.set_xlabel("Pseudotime")
            ax.set_ylabel("S Gene Expression")
            plt.savefig(genePseudotimeString)
            plt.close()


def transformData(n,
                  cells,
                  finalCellStates,
                  dropoutProportion=0.0,
                  cellRadiusDisp=0.14,
                  geneScale=0.0,
                  expScale=1.0,
                  expLambda=0.5,
                  shape=0.35,
                  rate=0.19,
                  gammaDistMean=True):

    originalGeneMeans = np.mean(finalCellStates, axis=1)
    geneMeanRankings = np.argsort(originalGeneMeans)

    if (gammaDistMean):
        newGeneMeans = np.log(
            np.random.gamma(shape, scale=1 / rate, size=n)) / expScale
        newGeneMeansSorted = np.sort(newGeneMeans)

        for i in range(len(geneMeanRankings)):
            geneIndex = geneMeanRankings[i]
            finalCellStates[geneIndex, :] += newGeneMeansSorted[i] - np.mean(
                finalCellStates[geneIndex, :])
    else:
        newGeneMeanOffsets = (
            geneScale + expLambda) / expScale + -1 * np.random.exponential(
                expLambda, size=n)
        newGeneMeansSorted = np.sort(newGeneMeanOffsets)

        for i in range(len(geneMeanRankings)):
            geneIndex = geneMeanRankings[i]
            finalCellStates[geneIndex, :] += newGeneMeansSorted[i]

    finalCellStates = np.exp(expScale * finalCellStates)

    cellSizeFactors = (np.random.normal(
        loc=1, scale=cellRadiusDisp, size=cells)**3)
    for j in range(cells):
        finalCellStates[:, j] *= cellSizeFactors[j]

    for i in range(n):
        for j in range(cells):
            if finalCellStates[i, j] < 0:
                finalCellStates[i, j] = 0
            finalCellStates[i, j] = np.random.poisson(
                lam=finalCellStates[i, j])

        #numCellKeep = numCellsKeep[i]
        #zeroCells = random.sample(range(cells), cells - numCellKeep)
        #finalCellStates[i, zeroCells] = 0

    return finalCellStates, cellSizeFactors


def analyzeSCRNAseqData(n,
                        cells,
                        cellTypesRecord,
                        finalCellStates,
                        perplexityParam=30,
                        saveFolder=None):
    cellTypeString = "All"
    if saveFolder is not None:
        plotSCRNAseqData(n, cells, finalCellStates, cellTypeString, saveFolder)

    cellTypesSet = list(set(cellTypesRecord))
    for cellTypeNum in cellTypesSet:
        cellTypeString = str(cellTypeNum)
        cellTypeIndices = [
            i for i in range(cells) if cellTypesRecord[i] == cellTypeNum
        ]
        finalCellStatesOfType = finalCellStates[:, cellTypeIndices]
        plotSCRNAseqData(n,
                         len(cellTypeIndices), finalCellStatesOfType,
                         cellTypeString, saveFolder)

    pca = PCA(n_components=15)
    pcaFit = pca.fit(finalCellStates)
    pcaResult = pca.components_
    # print("PCA Explained Var: ", pca.explained_variance_ratio_)
    if saveFolder is not None:
        np.save(
            os.path.join(saveFolder, "Variance_Explained"),
            np.asarray(pca.explained_variance_ratio_))
    xVals = pcaResult[0, :]
    yVals = pcaResult[1, :]

    if saveFolder is not None:
        fig = plt.figure()
        ax = plt.gca()
        ax.plot(xVals, yVals, 'bo')
        plt.savefig(os.path.join(saveFolder, "PCA"))
        plt.close()

    # model = TSNE(n_components=2, perplexity=perplexityParam)
    # np.set_printoptions(suppress=True)
    # tsneResult = model.fit_transform(np.transpose(finalCellStates)) #Transpose because cells are data we are clustering

    # palette = np.array(sns.color_palette("hls", max(cellTypesRecord)+1))

    # # We create a scatter plot.
    # f = plt.figure(figsize=(8, 8))
    # ax = plt.subplot(aspect='equal')

    # colorHandles = []
    # for cellTypeNum in cellTypesSet:
    # 	colorHandles.append(mpatches.Patch(color=palette[cellTypeNum], label=cellTypeNum))

    # plt.legend(handles=colorHandles, borderaxespad=0)

    # sc = ax.scatter(tsneResult[:,0], tsneResult[:,1], lw=0, s=40, c=palette[cellTypesRecord.astype(np.int)])
    # ax.axis('off')
    # ax.axis('tight')
    # #plt.show()
    # plt.savefig("t-SNE")


def plotSCRNAseqData(n, cells, cellStates, cellTypeString, saveFolder):
    """Plot scRNA-seq data

    n : int
        Number of genes
    cells : int
        Number of cells
    """
    numBinsScalar = 1
    newpath = os.path.join(saveFolder, "Cell_Type_" + cellTypeString)
    try:
        os.makedirs(newpath)
    except:
        pass

    geneExpr = cellStates[int(np.random.uniform() * n), :]
    fig = plt.figure()
    ax = plt.gca()
    ax.hist(geneExpr, bins=int(cells / numBinsScalar))
    plt.title("Expression of Random Gene")
    plt.savefig(os.path.join(saveFolder, "Expression_Of_Gene"))
    plt.close()

    cellExpr = cellStates[:, int(np.random.uniform() * cells)]
    fig = plt.figure()
    ax = plt.gca()
    ax.hist(cellExpr, bins=int(n / numBinsScalar))
    plt.title("Expression of Random Cell")
    plt.savefig(os.path.join(saveFolder, "Expression_Of_Cell"))
    plt.close()

    geneMeansString = "Gene_Means_" + "Cell_Type_" + cellTypeString
    librarySizesString = "Library_Sizes_" + "Cell_Type_" + cellTypeString
    numExpressedGenesString = "Num_Expressed_Genes_" + "Cell_Type_" + cellTypeString
    geneMeans = np.mean(cellStates, axis=1)
    fig = plt.figure()
    ax = plt.gca()
    ax.hist(geneMeans, bins=int(n / numBinsScalar))
    plt.title("Gene Means")  #Change back to Log[10]
    plt.savefig(os.path.join(saveFolder, geneMeansString))
    plt.close()

    librarySizes = np.sum(cellStates, axis=0)
    fig = plt.figure()
    ax = plt.gca()
    ax.hist(librarySizes, bins=int(n / numBinsScalar))
    plt.title(os.path.join(saveFolder, "Library Sizes"))
    plt.savefig(librarySizesString)
    plt.close()

    numExpressedGenes = np.zeros(shape=int(cells))
    for cellNum in range(cells):
        numExpressedGenes[cellNum] = np.count_nonzero(cellStates[:, cellNum])
    fig = plt.figure()
    ax = plt.gca()
    ax.hist(numExpressedGenes, bins=int(n / numBinsScalar))
    plt.title(os.path.join(saveFolder, "Num Expressed Genes"))
    plt.savefig(numExpressedGenesString)
    plt.close()
