### Imports

import numpy as np
import sympy as sp
import scipy
import scipy.special
from scipy import linalg
import matplotlib
import matplotlib.pyplot as plt
import random
from numpy.random import seed
from sympy.matrices import Matrix
from sympy import *


### Linear transformation process
print("This results correspond to the linear transformation process of the first figure.")

k = 1

plt.rcParams["font.family"] = 'Dejavu Serif'
plt.rcParams['mathtext.fontset'] = 'dejavuserif'
plt.rcParams.update({'font.size': 16})
plt.rcParams["figure.figsize"] = [8, 4.8]

grid = np.linspace(0,3,1000)
P2 = (1+np.exp(-4*grid)+2*np.exp(-2*grid))/4
P1 = (1-np.exp(-4*grid))/2
P0 = (1+np.exp(-4*grid)-2*np.exp(-2*grid))/4

p1 = (1 + np.exp(-2*grid))/2 
p0 = (1 - np.exp(-2*grid))/2 

plt.figure()
plt.xlabel(r"time $t$")
plt.plot(grid,P2,linewidth = 3.0,c = [0,0,1],label = r"$P(X=2,t)$")
plt.plot(grid,P1,linewidth = 3.0,c = [1,0,0],label = r"$P(X=1,t)$")
plt.plot(grid,P0,linewidth = 3.0,c = [1.0,0.5,0.1],label = r"$P(X=0,t)$")
plt.legend()
plt.show()

plt.figure()
plt.xlabel(r"time $t$")
plt.plot(grid,p1*p1,linewidth = 3.0,c = [0,0.8,0.8],label = r"$ \ \ \ P(Y^1=1,t) P(Y^2=1,t)$")
plt.plot(grid,2*p0*p1,linewidth = 3.0,c = [0.4,0.4,0.8],label = r"$2 \ P(Y^1=1,t) P(Y^2=0,t)$")
plt.plot(grid,p0*p0,linewidth = 3.0,c = [0.8,0.0,0.8],label = r"$ \ \ \ P(Y^1=0,t)P(Y^2=0,t)$")
plt.legend()
plt.show()

### epigenetic regulation via DNA methylation

plt.rcParams["font.family"] = 'Dejavu Serif'
plt.rcParams['mathtext.fontset'] = 'dejavuserif'
plt.rcParams.update({'font.size': 16})
plt.rcParams["figure.figsize"] = [8, 4.8]


t=Symbol('t')
A=Matrix([[-1,1,0,0],
          [1,-2,0,0],
          [0,0,-1,1],
          [0,1,1,-1]])

B=Matrix([[-1,1],
          [1,-1]])

C=Matrix([[-2,1,0,0,0,0],
          [2,-3,2,0,0,0],
          [0,1,-4,0,0,0],
          [0,0,0,-2,1,0],
          [0,1,0,2,-2,2],
          [0,0,2,0,1,-2]])


eA = exp(A*t)
eB = exp(B*t)
eC = exp(C*t)
print(r"P(X,t)=")
display(eC[:,0])

print(r"exp(tS)=")
display(eA)

comb = Matrix([[eA[0,0]*eA[0,0]],[2*eA[0,0]*eA[1,0]],[eA[1,0]*eA[1,0]],[eA[2,0]*(eA[0,0] + eB[0,0])],[eA[2,0]*(eB[1,0] +eA[1,0]) + eA[3,0]*(eB[0,0]+ eA[0,0])],[eA[3,0]*(eA[1,0] + eB[1,0])]])


numTimeSteps = 100
tEnd = 10.0

single = np.zeros((6,numTimeSteps))
double = np.zeros((6,numTimeSteps))

i = 0
for T in np.linspace(0,tEnd,numTimeSteps):
    single[0,i] = (eA[0,0]*eA[0,0]).evalf(subs={t: T})
    double[0,i] = eC[0,0].evalf(subs={t: T})
    
    single[1,i] = (2*eA[0,0]*eA[1,0]).evalf(subs={t: T})
    double[1,i] = eC[1,0].evalf(subs={t: T})
    
    single[2,i] = (eA[1,0]*eA[1,0]).evalf(subs={t: T})
    double[2,i] = eC[2,0].evalf(subs={t: T})
    
    single[3,i] = (eA[2,0]*(eA[0,0] + eB[0,0])).evalf(subs={t: T})
    double[3,i] = eC[3,0].evalf(subs={t: T})

    single[4,i] = (eA[2,0]*(eB[1,0] +eA[1,0]) + eA[3,0]*(eB[0,0]+ eA[0,0])).evalf(subs={t: T})
    double[4,i] = eC[4,0].evalf(subs={t: T})
    
    single[5,i] = (eA[3,0]*(eA[1,0] + eB[1,0])).evalf(subs={t: T})
    double[5,i] = eC[5,0].evalf(subs={t: T})

    i += 1
    
print("This results correspond to the model for epigenetic regulation via DNA methylation in the second figure")

plt.figure()
for k in range(6):
    plt.xlabel(r"time $t$")
    plt.ylabel(r"$P(\bf{{{Y}}}$ $,t)$")
    plt.plot(np.linspace(0,tEnd,numTimeSteps),single[k,:],color = [1.0-0.1*k,0.0,0.2+0.1*k],linewidth = 3.0,label = r"$\bf{{{Y}}}$ $=("+str(np.mod(k,3))+r","+str(np.floor_divide(k,3))+r")$")
    plt.legend()
    
plt.figure()
for l in range(6):
    plt.xlabel(r"time $t$")
    plt.ylabel(r"$P(\bf{{{X}}}$ $,t)$")
    plt.plot(np.linspace(0,tEnd,numTimeSteps),double[l,:],color = [1.0-0.1*l,0.2+0.1*l,0.0],linewidth = 3.0,label = r"$\bf{{{X}}}$ $=("+str(np.mod(l,3))+r","+str(np.floor_divide(l,3))+r")$")
    plt.legend()

### Implementation of Classic SSA and Single Catalyst SSA

def gillespie(Xinit, tEnd, nP, nS, nD, k):
    X = Xinit
    t = 0
    
    v = np.zeros((3*nD,2*nD))
    v[:nD,:nD] = np.eye(nD)
    v[nD:2*nD,:nD] = -np.eye(nD)
    v[2*nD:3*nD,nD:2*nD] = np.eye(nD)
    
    while (t < tEnd):
        A = np.zeros(3*nD)
        
        for r in range(nD):
            A[r] = k[0]*(nP-np.sum(X[0,:nD]))
            A[r+nD] = k[1]*X[0,r]
            A[r+2*nD] = k[2]*X[0,r]*(1-X[0,r+nD]/nS)
            
 
        possibleReactions = np.linspace(0,3*nD-1,3*nD,dtype='int')
        probabilities = A/np.sum(A)
        i = np.random.choice(possibleReactions, 1, p=probabilities)
        
        t = t + np.log(1/np.random.rand(1))/np.sum(A)
        
        if (t < tEnd):
            X = X + v[i,:]
            
    return X

def singleMolGillespie(Xinit, tEnd, nP, nS, nD, k):
    Ym = Xinit[0,nD:]
    Yp = np.zeros((nP,nD))
    
    for p in range(nP):
        
        Y = np.array([np.concatenate((Xinit[0,:nD],Ym))])
        Y = gillespie(Y, tEnd, 1, nS, nD, k)
        
        Yp[p,:] = Y[0,:nD]
        Ym = Y[0,nD:]
        
    return np.array([np.concatenate((np.sum(Yp,0),Ym))])

def comparison(Xinit, tEnd, nP, nS, nD, k, numSamples):
    
    plt.rcParams["font.family"] = 'Dejavu Serif'
    plt.rcParams['mathtext.fontset'] = 'dejavuserif'
    plt.rcParams.update({'font.size': 16})
    plt.rcParams["figure.figsize"] = [8, 4.8]
    
    Xres = np.zeros((numSamples,2*nD))
    Yres = np.zeros((numSamples,2*nD))
    
    for s in range(numSamples):
        Xres[s,:] = gillespie(Xinit, tEnd, nP, nS, nD, k)
        Yres[s,:] = singleMolGillespie(Xinit, tEnd, nP, nS, nD, k)
        
    error = 0
    for i in range(nD):
        
        plt.figure()
        plt.title(r"$P(\bf{{{X}}}$$,t=3.0)$ from Classic SSA")
        plt.xlabel(r"# protein molecules bound to DNA " + str(i+1))
        plt.ylabel(r"# methylations at DNA " + str(i+1))
        Xhist,_,_,_ = plt.hist2d(Xres[:,i], Xres[:,i+nD], cmap='plasma', density = True)
        plt.colorbar()
        plt.show()

        plt.figure()
        plt.title(r"$P(\bf{{{Y}}}$$,t=3.0)$ from Single Catalyst SSA")
        plt.xlabel(r"# protein molecules bound to DNA " + str(i+1))
        plt.ylabel(r"# methylations at DNA " + str(i+1))
        Yhist,_,_,_ = plt.hist2d(Yres[:,i], Yres[:,i+nD], cmap='plasma', density = True)
        plt.colorbar()
        plt.show()

        error += np.sum(np.sum(np.power(Xhist-Yhist,2)))
    return error

print("This results correspond to the gillespie simulations in Figure 3")
np.random.seed(1)

tEnd = 3.0
nP = 10
nS = 10
nD = 2
k = np.array([2.0,1.0,0.5])
Xinit = np.zeros((1,2*nD))

numSamples = 10**5
comparison(Xinit, tEnd, nP, nS, nD, k, numSamples)

np.random.seed(1)

tEnd = 3.0
nP = 100
nS = 100
nD = 5
k = np.array([2.0,1.0,0.5])
Xinit = np.zeros((1,2*nD))

numSamples = 10**5
comparison(Xinit, tEnd, nP, nS, nD, k, numSamples)