import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from scipy.integrate import odeint
import pprint
import textwrap

def mmreaction(initial, t, kSin, kON, kOFF, kcat, kEout, kSout, kCout, b):
    """Return derivatives at t = 0 for each molecule.

    This is just a helper function for the scipy.integrate.odeint solver
    See the documentation for scipy.integrate.odeint for more explanation.
    """
    Eo, So, Co = initial
    dedt = - (kON * Eo * So) + (kOFF * Co) + (2 * kcat * Co) - (kEout * Eo)
    dsdt = - (kON * Eo * So) + (kOFF * Co) + (kSin * b) - (kSout * So)
    dcdt = (kON * Eo * So) - (kOFF * Co) - (kcat * Co) - (kCout * Co)
    return [dedt, dsdt, dcdt]

pops = {'E': 20, 'S': 10, 'C': 0}
ratesC = {'kSin': 0.25, 'kON': 0.19, 'kOFF': 0.45,
          'kcat': 0.010, 'kEout': 0.004, 'kSout': 0.004,
          'kCout': 0.002, 'b': 1}
# Used to adjust kSin for each burst size later
start_kSin = ratesC['kSin']
ratesC['Km'] = (ratesC['kOFF'] + ratesC['kcat']) / ratesC['kON']
endtime = (10 * np.log(2)
        / min([ratesC['kEout'], ratesC['kSout'], ratesC['kCout']]))

init = [pops['E'], pops['S'], pops['C']]
t = np.linspace(0, endtime, 10000)

endstates = []
pp = pprint.PrettyPrinter(indent = 4, width = 76)

for i in range(50):
    ratesC['b'] = i + 1
    ratesC['kSin'] = start_kSin / ratesC['b']

    constants = (ratesC['kSin'], ratesC['kON'], ratesC['kOFF'], ratesC['kcat'],
                 ratesC['kEout'], ratesC['kSout'], ratesC['kCout'],
                 ratesC['b'])
    endstates.append(odeint(mmreaction, init, t, args = constants)[-1])

endstates = pd.DataFrame(endstates, index=range(1,51),
                         columns=['Mean Enzyme', 'Mean Substrate',
                                  'Mean Complex'])
endstates.index.name = 'Burst'
wrapper = textwrap.TextWrapper()
with open('BURST_StaticMeanControl.txt', 'w') as f:
    pp = pprint.PrettyPrinter(width=80, stream=f)
    f.write('Static Mean Control')
    f.write('\n' + '=' * 80 + '\n\n')
    intro = ('In theory, as the burst parameter increases, it would also '
         + 'change the dS/dt because the rate of S formation events, '
         + 'kSin is not changing, but the burst size of S formation '
         + 'is increasing. To avoid this and ensure that the only cause '
         + 'of a varying mean population of S and consequently E and C '
         + 'is the burst size, we updated kSin as follows:\n\n    kSin = '
         + 'Original_kSin / Burst Size\n\nThus, the mean populations for '
         + 'all three species should remain constant in a deterministic '
         + 'model, despite the varying burst size. The results of the '
         + 'control are below.')
    f.write(wrapper.fill(intro))
    f.write('\n\n' + '=' * 80 + '\n\n')
    pp.pprint(endstates)
