from scipy.special import expit
from scipy.misc import derivative
from sklearn.manifold import TSNE
from sklearn.decomposition import PCA
from rpy2.robjects import r
from rpy2.robjects.numpy2ri import numpy2ri
from shutil import copyfile
from collections import Counter
from pylab import pcolor, show, colorbar, xticks, yticks
from numpy import arange

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import numpy as np
import random
import scipy
import time
import math
import os
import datetime
import shutil  #Doesn't work on MAC

#Global Variable#
global currentCellType
currentCellType = 0


#Generate the W matrix
def generateGeneCorrelationMatrix(n, maxAlphaBeta, normalizeWRows,
                                  networkStructure, powerLawExponent,
                                  networkSparsity, saveFolder):
    networkMatrix = np.zeros(shape=(n, n))

    #Generate the network structure
    if (networkStructure == "SimulatedPowerLaw"):
        #Generate number of in and out connections per gene
        inConnections = [
            int((1 / 2.0) * math.pow(1 - np.random.uniform(), -1 /
                                     (powerLawExponent - 1)) + 1 / (2.0))
            for i in range(n)
        ]
        outConnections = [
            int((1 / 2.0) * math.pow(1 - np.random.uniform(), -1 /
                                     (powerLawExponent - 1)) + 1 / (2.0))
            for i in range(n)
        ]

        #Make sure number of in and out connections is not greater than the number of genes
        for i in range(n):
            if (inConnections[i] > n):
                inConnections[i] = n
            if (outConnections[i] > n):
                outConnections[i] = 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):
            numOutConnection = outConnections[col]
            outNodesSamplePool += [col] * numOutConnection

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

    elif (networkStructure == "TRRUST"):
        #This file must be downloaded, or it can be generated by using the final AnalyzeTRRUST.py and the original file
        networkMatrix = np.load("TRRUST_Network.npy")

    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')

    analyzeNetworkStructure(n, networkMatrix, networkStructure, saveFolder)

    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 = int(maxAlphaBeta * np.random.uniform() + 1)
        beta = int(maxAlphaBeta * np.random.uniform() + 1)
        valueMatrix[col, :] = (
            2 * np.random.randint(0, 2) - 1) * np.random.beta(
                a=alpha, b=beta, size=n)

    #Element-wise multiply network structure matrix and value matrix
    gc = np.multiply(networkMatrix, valueMatrix)

    #If set to true, will divide each value of each row (in connections) by the number of connections
    if (normalizeWRows):
        for row in range(n):
            if (np.count_nonzero(gc[row, :]) != 0):
                gc[row, :] = gc[row, :] / np.count_nonzero(
                    gc[row, :])  #Normalize by number of connections per row

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

    return (gc, transcriptionFactors)


#Analysis of the indegree and outdegree distributions of the chosen network
def analyzeNetworkStructure(n, networkMatrix, networkStructure, saveFolder):
    os.chdir(saveFolder)

    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)

    plt.hist(inDegrees, bins=100, range=[0, 100])
    plt.xlabel("Indegree")
    plt.ylabel("Frequency")
    inDegreeString = networkStructure + "_Network_Indegree.tiff"
    plt.savefig(inDegreeString)
    plt.clf()
    plt.cla()

    plt.hist(outDegrees, bins=100, range=[0, 100])
    plt.xlabel("Outdegree")
    plt.ylabel("Frequency")
    outDegreeString = networkStructure + "_Network_Outdegree.tiff"
    plt.savefig(outDegreeString)
    plt.clf()
    plt.cla()

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


#Transform cell states from (0, 1) range to (-1, 1) range
def rangeTransformation(cellStates):
    cellStates = 2 * cellStates - 1
    return (cellStates)


#Randomly generate an initial gene expression state, and use it for all the cells
def generateInitialStates(n, cells, geneMeanLevels):
    initialStates = np.zeros(shape=(n, cells))
    baseInitialState = rangeTransformation(np.random.rand(n)) + geneMeanLevels
    for j in range(cells):
        initialStates[:, j] = baseInitialState
    return (initialStates)


#Equation to develop each cell for a time step
def developCell(n, gc, timeStep, initialState, proj, const, noiseDisp,
                geneMeanLevels):
    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


#Find the optimal alpha for W and the number of optimal iterations given that alpha
def findOptimalAlphaAndIterations(n, gc, timeStep, projMatrix, constMatrix,
                                  geneMeanLevels, convergenceThreshold,
                                  noiseDisp, currentStates, maxNumIterations,
                                  convDistAvgNum, numToSample, initialAlpha,
                                  alpha_step, showAlphaGeneMeans):

    #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

    global currentCellType
    proj = projMatrix[currentCellType, :]
    const = constMatrix[currentCellType, :]

    #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, avgIterations)


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

    #Initialize variables (not to be changed)
    global currentCellType
    proj = projMatrix[currentCellType, :]
    const = constMatrix[currentCellType, :]

    numKeepCellType = int(cells * keepProportion)

    cells = currentStates.shape[1]
    avgIterations = 0
    normDiffVals = [0] * maxNumIterations

    #Sample cells and save them in a new array so they don't overwrite old values
    colsToSample = random.sample(range(numKeepCellType, cells), 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)


#Develop the first "cell type" - the progenitor of all other "cell types"
def developInitialStates(n, cells, initialStates, gc, timeStep, iterations,
                         cellTypesRecord, totalCellIterations, projMatrix,
                         constMatrix, noiseDisp, geneMeanLevels, showPlots,
                         keepProportion, currentCellSplit, numCellSplits,
                         cellDevelopmentMode):
    global currentCellType
    print("Current Cell Type: ", currentCellType)
    numKeepCellType = int(
        keepProportion * cells)  #How many cells to keep at the first cell type
    numDiffCellType = cells - numKeepCellType  #How many cells to differentiate later

    cellTypesRecord[0:cells] = int(currentCellType)  #Fill in cell types
    proj1 = projMatrix[currentCellType, :]
    const1 = constMatrix[currentCellType, :]
    currentCellType = currentCellType + 1

    for i in range(0, cells):
        initialState = initialStates[:, i]
        if (not cellDevelopmentMode):  #If cells are supposed to all be fully developed
            cellIterations = np.random.poisson(
                iterations
            )  #The number of iterations should be centered around the equilibrium value
        else:  #Else if the cells are captured at different developmental time points
            if (i < numKeepCellType or currentCellSplit ==
                    numCellSplits):  #If cell's final cell type is the current cell type
                cellIterations = np.random.randint(0, iterations + 1)
            else:  #Else if the cell's cell type is not the current cell type
                cellIterations = iterations
        totalCellIterations[i] += cellIterations
        initialStates[:, i] = findCellNextState(
            gc, n, timeStep, i, proj1, const1, initialState, cellIterations,
            noiseDisp, geneMeanLevels, showPlots)

    return (initialStates)


#Generate projection and constant vectors for next cell type
def getNextProjConst(n, projInit, constInit, constProb, transcriptionFactors):
    numZeros = int(n * constProb)  #np.random.binomial(n, constProb, 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 getProjConstNextRow(n, projInit, constInit, projMatrix, constMatrix,
                        constProb, currentCellSplit, numCellSplits,
                        transcriptionFactors):
    global currentCellType

    proj1, const1 = getNextProjConst(n, projInit, constInit, constProb,
                                     transcriptionFactors)
    proj2, const2 = getNextProjConst(n, projInit, constInit, constProb,
                                     transcriptionFactors)

    projMatrix[currentCellType, :] = proj1
    constMatrix[currentCellType, :] = const1
    currentCellType = currentCellType + 1

    projMatrix[currentCellType, :] = proj2
    constMatrix[currentCellType, :] = const2
    currentCellType = currentCellType + 1

    currentCellSplit = currentCellSplit + 1
    if (currentCellSplit < numCellSplits):
        getProjConstNextRow(n, proj1, const1, projMatrix, constMatrix,
                            constProb, currentCellSplit, numCellSplits,
                            transcriptionFactors)
        getProjConstNextRow(n, proj2, const2, projMatrix, constMatrix,
                            constProb, currentCellSplit, numCellSplits,
                            transcriptionFactors)


def getProjConstMatrices(n, constProb, targetNumCellTypes, currentCellSplit,
                         numCellSplits, transcriptionFactors):
    global currentCellType
    totalNumCellTypes = 2 * targetNumCellTypes - 1
    projMatrix = np.ones(shape=(totalNumCellTypes, n))
    constMatrix = np.zeros(shape=(totalNumCellTypes, n))

    projInit, constInit = getNextProjConst(n, projMatrix[0, :],
                                           constMatrix[0, :], constProb,
                                           transcriptionFactors)

    projMatrix[currentCellType, :] = projInit
    constMatrix[currentCellType, :] = constInit
    currentCellType = currentCellType + 1

    if (currentCellSplit < numCellSplits):
        getProjConstNextRow(n, projInit, constInit, projMatrix, constMatrix,
                            constProb, currentCellSplit, numCellSplits,
                            transcriptionFactors)

    currentCellType = 0

    return projMatrix, constMatrix


def generateGeneMeanLevels(n, expScale):
    geneMeanLevels = np.zeros(shape=(n))
    return (geneMeanLevels)


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.plot(xVals, normDiffVals)
        plt.show()

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

        plt.hist(yVals, bins=int(n / 25))
        plt.show()

    return (currentState)


def findAllNextStates(gc, n, timeStep, numCellTypes, numCellSplits,
                      currentCellSplit, projMatrix, constMatrix, currentStates,
                      cellTypesRecord, totalCellIterations, cellNumStart,
                      cellNumEnd, iterations, noiseDisp, geneMeanLevels,
                      showPlots, keepProportion, cellDevelopmentMode):
    localNumCells = cellNumEnd - cellNumStart + 1  #CellNumEnd is inclusive
    numKeepCellType = int(keepProportion * localNumCells)
    numOfCellType1 = int(
        (1 / 2) * (localNumCells - numKeepCellType
                   ))  #int(np.random.uniform() * localNumCells)
    numOfCellType2 = localNumCells - numKeepCellType - numOfCellType1
    cellNumStart = cellNumStart + numKeepCellType

    global currentCellType
    print("Current Cell Type: ", currentCellType)
    currentCellSplit = currentCellSplit + 1
    #Fill in cell types
    cellTypesRecord[cellNumStart:(cellNumStart + numOfCellType1)] = int(
        currentCellType)
    proj1 = projMatrix[currentCellType, :]
    const1 = constMatrix[currentCellType, :]
    currentCellType = currentCellType + 1

    for i in range(cellNumStart, cellNumStart + numOfCellType1):
        initialState = currentStates[:, i]
        if (not cellDevelopmentMode):
            cellIterations = np.random.poisson(iterations)
        elif ((i - cellNumStart) < int(keepProportion * numOfCellType1)
              or currentCellSplit == numCellSplits):
            cellIterations = np.random.randint(0, iterations + 1)
        else:
            cellIterations = iterations
        totalCellIterations[i] = totalCellIterations[i] + cellIterations
        nextCellState = findCellNextState(gc, n, timeStep, i, proj1, const1,
                                          initialState, cellIterations,
                                          noiseDisp, geneMeanLevels, showPlots)
        currentStates[:, i] = nextCellState

    print(currentCellType)
    cellTypesRecord[(cellNumStart + numOfCellType1):cellNumEnd + 1] = int(
        currentCellType)
    proj2 = projMatrix[currentCellType, :]
    const2 = constMatrix[currentCellType, :]
    currentCellType = currentCellType + 1

    for i in range(cellNumStart + numOfCellType1, cellNumEnd + 1):
        initialState = currentStates[:, i]
        if (not cellDevelopmentMode):
            cellIterations = np.random.poisson(iterations)
        elif ((i - (cellNumStart + numOfCellType1)) < int(
                keepProportion * numOfCellType2)
              or currentCellSplit == numCellSplits):
            cellIterations = np.random.randint(0, iterations + 1)
        else:
            cellIterations = iterations
        totalCellIterations[i] = totalCellIterations[i] + cellIterations
        nextCellState = findCellNextState(gc, n, timeStep, i, proj2, const2,
                                          initialState, cellIterations,
                                          noiseDisp, geneMeanLevels, showPlots)
        currentStates[:, i] = nextCellState

    if (currentCellSplit < numCellSplits):
        findAllNextStates(gc, n, timeStep, numCellTypes, numCellSplits,
                          currentCellSplit, projMatrix, constMatrix,
                          currentStates, cellTypesRecord, totalCellIterations,
                          cellNumStart, cellNumStart + numOfCellType1 - 1,
                          iterations, noiseDisp, geneMeanLevels, showPlots,
                          keepProportion, cellDevelopmentMode)
        findAllNextStates(gc, n, timeStep, numCellTypes, numCellSplits,
                          currentCellSplit, projMatrix, constMatrix,
                          currentStates, cellTypesRecord, totalCellIterations,
                          cellNumStart + numOfCellType1, cellNumEnd,
                          iterations, noiseDisp, geneMeanLevels, showPlots,
                          keepProportion, cellDevelopmentMode)


def analyzeDataBeforeTransformation(n, cells, finalCellStates, saveFolder,
                                    alpha, randomGenes, pseudotimes,
                                    networkStructure):
    os.chdir(saveFolder)

    #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 = ""

    fig = plt.figure()
    SGeneMeans = np.mean(finalCellStates, axis=1)
    geneMeansString = "S_Gene_Means" + testingString + ".tiff"
    plt.hist(SGeneMeans, bins=int(n / 50))
    ax = fig.add_subplot(111)
    ax.set_xlabel("S Gene Means")
    ax.set_ylabel("Frequency")
    plt.savefig(geneMeansString)
    plt.clf()
    plt.cla()

    fig = plt.figure()
    SGeneVars = np.var(finalCellStates, axis=1)
    geneVarsString = "S_Gene_Vars" + testingString + ".tiff"
    plt.hist(SGeneVars, bins=int(n / 50))
    ax = fig.add_subplot(111)
    ax.set_xlabel("S Gene Variances")
    ax.set_ylabel("Frequency")
    plt.savefig(geneVarsString)
    plt.clf()
    plt.cla()

    fig = plt.figure()
    SCellMeans = np.mean(finalCellStates, axis=0)
    cellMeansString = "S_Cell_Means" + testingString + ".tiff"
    plt.hist(SCellMeans, bins=int(cells / 10))
    ax = fig.add_subplot(111)
    ax.set_xlabel("S Cell Means")
    ax.set_ylabel("Frequency")
    plt.savefig(cellMeansString)
    plt.clf()
    plt.cla()

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

    for randomGene in randomGenes:
        geneExpr = finalCellStates[randomGene, :]
        fig = plt.figure()
        genePseudotimeString = "Gene_" + str(
            randomGene) + "_Pseudotime" + testingString + ".tiff"
        plt.scatter(pseudotimes, geneExpr)
        ax = fig.add_subplot(111)
        ax.set_xlabel("Pseudotime")
        ax.set_ylabel("S Gene Expression")
        plt.savefig(genePseudotimeString)
        plt.clf()
        plt.cla()


def transformData(n, cells, finalCellStates, dropoutProportion, cellRadiusDisp,
                  geneScale, expScale, expLambda, saveFolder, shape, rate,
                  gammaDistMean):

    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 plotSCRNAseqData(n, cells, cellStates, cellTypeString, saveFolder):
    numBinsScalar = 1
    newpath = saveFolder + "Cell_Type_" + cellTypeString
    if not os.path.exists(newpath):
        os.makedirs(newpath)

    os.chdir(newpath)

    geneExpr = cellStates[int(np.random.uniform() * n), :]
    plt.hist(geneExpr, bins=int(cells / numBinsScalar))
    plt.title("Expression of Random Gene")
    plt.savefig("Expression_Of_Gene")
    plt.clf()
    plt.cla()

    cellExpr = cellStates[:, int(np.random.uniform() * cells)]
    plt.hist(cellExpr, bins=int(n / numBinsScalar))
    plt.title("Expression of Random Cell")
    plt.savefig("Expression_Of_Cell")
    plt.clf()
    plt.cla()

    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)
    plt.hist(geneMeans, bins=int(n / numBinsScalar))
    plt.title("Gene Means")  #Change back to Log[10]
    plt.savefig(geneMeansString)
    plt.clf()
    plt.cla()

    librarySizes = np.sum(cellStates, axis=0)
    plt.hist(librarySizes, bins=int(n / numBinsScalar))
    plt.title("Library Sizes")
    plt.savefig(librarySizesString)
    plt.clf()
    plt.cla()

    numExpressedGenes = np.zeros(shape=int(cells))
    for cellNum in range(cells):
        numExpressedGenes[cellNum] = np.count_nonzero(cellStates[:, cellNum])
    plt.hist(numExpressedGenes, bins=int(n / numBinsScalar))
    plt.title("Num Expressed Genes")
    plt.savefig(numExpressedGenesString)
    plt.clf()
    plt.cla()


def analyzeSCRNAseqData(n, cells, perplexityParam, cellTypesRecord,
                        finalCellStates, saveFolder):
    cellTypeString = "All"
    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)

    os.chdir(saveFolder)
    pca = PCA(n_components=15)
    pcaFit = pca.fit(finalCellStates)
    pcaResult = pca.components_
    print(pca.explained_variance_ratio_)
    np.save("Variance_Explained", np.asarray(pca.explained_variance_ratio_))
    xVals = pcaResult[0, :]
    yVals = pcaResult[1, :]
    plt.plot(xVals, yVals, 'bo')
    plt.savefig("PCA")

    # 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 makeFolder(n, cells):
    cwd = os.getcwd()
    dataPath = cwd + "/Saved_Runs"
    experimentName = datetime.datetime.now().strftime(
        "%m-%d-%Y+%H:%M:%S") + "_N:" + str(n) + "_C:" + str(cells)
    copyFolder = dataPath + "/" + experimentName + "/"
    if not os.path.exists(copyFolder):
        os.makedirs(copyFolder)

    # shutil.copy(__file__, copyFolder + "MatrixGenerator.py")

    return (copyFolder)


def loadSavedData():
    gc = np.load("gc.npy")
    projMatrix = np.load("projMatrix.npy")
    constMatrix = np.load("constMatrix.npy")
    #currentStates = np.load("currentStates.npy")

    return gc, projMatrix, constMatrix


def saveData(synthscRNAseq, cellTypesRecord, gc, projMatrix, constMatrix,
             cellSizeFactors, synthPseudotimes, copyFolder, targetDirectory,
             R_Directory):

    os.chdir(copyFolder)

    itemsToSave = [
        synthscRNAseq, cellTypesRecord, gc, projMatrix, constMatrix,
        cellSizeFactors, synthPseudotimes
    ]
    itemNames = [
        "synthscRNAseq", "cellTypesRecord", "gc", "projMatrix", "constMatrix",
        "cellSizeFactors", "synthPseudotimes"
    ]

    for i in range(len(itemsToSave)):
        np.save(itemNames[i], itemsToSave[i])
        np.save(targetDirectory + itemNames[i], itemsToSave[i])
        ro = numpy2ri(itemsToSave[i])
        r.assign(itemNames[i], ro)
        rSaveString = "save(" + itemNames[i] + ", file='" + itemNames[i] + ".gzip', compress=TRUE)"
        r(rSaveString)
        # rSaveString = "save(" + itemNames[i] + ", file='" + R_Directory + ".gzip', compress=TRUE)"
        # r(rSaveString)


def main():
    ###Step 0 - choose a directory to work in###
    targetDirectory = "."
    R_Directory = "./R_files/"
    os.chdir(targetDirectory)

    ###Step 1 - Choose the number of genes (n) and cells###

    n = 500  #2718 #19027
    cells = 100  #864

    ###Step 2 - Choose the network structure you would like:###

    #SimulatedPowerLaw - Simulated power law network with powerLawExponent as parameter
    #TRRUST - Real network from TRRUST database of transcription factor interactions
    #Random - Random network with networkSparsity as a parameter

    networkStructure = "SimulatedPowerLaw"

    powerLawExponent = 2.05  #Increase for more connections, but must stay between 2-3
    networkSparsity = 0.999

    ###Step 3 - Choose the value of alpha (alpha must be greater than or equal to 0)#
    #and the maximum value of the Beta distribution alpha and beta parameters###

    initialAlpha = 1  #Set to 0 otherwise
    maxAlphaBeta = 10

    ###Choose whether to divide each row of W by the number of connections:###

    normalizeWRows = False

    ###Step 4 - Choose the cell development structure:###

    cellDevelopmentMode = True
    targetNumCellTypes = 2  #4#8
    constProb = 0.05
    keepProportion = 1.0 / 4

    ###Step 5 - Choose the mean and cell size mapping parameters:###

    gammaDistMean = True

    #Gamma mean parameters
    shape = 0.35  #0.429743565557383
    rate = 0.19

    #Exp modified mean parameters
    expLambda = 0  #0.50
    geneScale = 0  #0.47

    #Expontial base and cell size parameters
    expScale = 1  #4#6.81
    cellRadiusDisp = 0.14

    ###Step 6 - Choose the timeStep, noise, convergence, and iteration sampling parameters:###

    timeStep = 0.01  #0.01

    noiseDisp = 5e-2
    convergenceThreshold = 1e-1

    alpha_step = 0.02  #How much to increase alpha by in each step
    maxNumIterations = 500
    numToSample = 50
    convDistAvgNum = 20

    ###Display parameters###

    loadData = False
    showPlots = False
    showAlphaGeneMeans = False
    saveFolder = None  #Set to None if you want to make it dynamic
    randomGenes = np.random.randint(0, n, size=5)

    ###Unused parameters###
    dropoutProportion = 0.0

    ###Parameters that are usually not changed
    perplexityParam = 30
    randomSeed = 393  #123
    numCellTypes = 2  #Doesn't work with other values right now - num cell types to make at each split
    numCellSplits = (int)(math.log(targetNumCellTypes, numCellTypes))

    np.random.seed(randomSeed)
    random.seed(randomSeed)
    currentCellSplit = 0
    cellTypesRecord = np.zeros(shape=cells, dtype="int64")
    np.set_printoptions(threshold=20)
    totalCellIterations = np.zeros(shape=cells, dtype="int64")
    maxNumIterations = maxNumIterations + convDistAvgNum

    if saveFolder == None:
        saveFolder = makeFolder(n, cells)

    gc, transcriptionFactors = generateGeneCorrelationMatrix(
        n, maxAlphaBeta, normalizeWRows, networkStructure, powerLawExponent,
        networkSparsity, saveFolder)
    projMatrix, constMatrix = getProjConstMatrices(
        n, constProb, targetNumCellTypes, currentCellSplit, numCellSplits,
        transcriptionFactors)
    geneMeanLevels = generateGeneMeanLevels(n, expScale)

    initialStates = generateInitialStates(n, cells, geneMeanLevels)
    print("Initial States: ", initialStates)
    alpha, iterations = findOptimalAlphaAndIterations(
        n, gc, timeStep, projMatrix, constMatrix, geneMeanLevels,
        convergenceThreshold, noiseDisp, initialStates, maxNumIterations,
        convDistAvgNum, numToSample, initialAlpha, alpha_step,
        showAlphaGeneMeans)
    gc = gc / alpha  #Update gc with new alpha

    if (loadData):
        gc, projMatrix, constMatrix = loadSavedData()
    print("Loaded Data: ", loadData)
    print("GC :", gc)

    currentStates = developInitialStates(
        n, cells, initialStates, gc, timeStep, iterations, cellTypesRecord,
        totalCellIterations, projMatrix, constMatrix, noiseDisp,
        geneMeanLevels, showPlots, keepProportion, currentCellSplit,
        numCellSplits, cellDevelopmentMode)
    print("Alpha: ", alpha, "\n Branch 0 Iterations: ", iterations)
    print("Current State at Branch 0: \n", currentStates)
    np.save("initialDevelopedStates_saved", currentStates)

    #The alpha and iterations will only work for the first branch, as it runs on ALL cells (not just subsection)!
    if (numCellSplits != 0):  #Make sure that cell should even split
        iterations = findOptimalIterations(
            n, cells, gc, timeStep, projMatrix, constMatrix, geneMeanLevels,
            convergenceThreshold, noiseDisp, currentStates, maxNumIterations,
            convDistAvgNum, numToSample, keepProportion)
        findAllNextStates(gc, n, timeStep, numCellTypes, numCellSplits,
                          currentCellSplit, projMatrix, constMatrix,
                          currentStates, cellTypesRecord, totalCellIterations,
                          0, cells - 1, iterations, noiseDisp, geneMeanLevels,
                          showPlots, keepProportion, cellDevelopmentMode)

    finalCellStates = currentStates
    print("Final Cell States Before Transformations: \n", finalCellStates)

    cellSizeFactors = np.zeros(shape=cells)
    pseudotimes = timeStep * totalCellIterations
    analyzeDataBeforeTransformation(n, cells, finalCellStates, saveFolder,
                                    alpha, randomGenes, pseudotimes,
                                    networkStructure)
    finalCellStates, cellSizeFactors = transformData(
        n, cells, finalCellStates, dropoutProportion, cellRadiusDisp,
        geneScale, expScale, expLambda, saveFolder, shape, rate, gammaDistMean)
    print("Final States After Transformations: \n", finalCellStates)

    saveData(finalCellStates, cellTypesRecord, gc, projMatrix, constMatrix,
             cellSizeFactors, pseudotimes, saveFolder, targetDirectory,
             R_Directory)
    analyzeSCRNAseqData(n, cells, perplexityParam, cellTypesRecord,
                        finalCellStates, saveFolder)


main()
