# SBML model we want to import
sbml_file = 'model_presimulation.xml'
# Name of the model that will also be the name of the python module
model_name = 'model_presimulation'
# Directory to which the generated model code is written
model_output_dir = model_name

import libsbml
import amici
import amici.plotting
import os
import sys
import importlib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pprint import pprint

"""
Explanation: AMICI Python example "Experimental Conditions"
In this example we will explore some more options for the initialization of experimental conditions, including how to reset initial conditions based on changing values for fixedParameters as well as an additional presimulation phase on top of preequilibration. This notebook is expected to run from the python/example_presimulation directory.
End of explanation
"""


sbml_reader = libsbml.SBMLReader()
sbml_doc = sbml_reader.readSBML(sbml_file)
sbml_model = sbml_doc.getModel()

print('Species:')
pprint([(s.getId(),s.getName()) for s in sbml_model.getListOfSpecies()])

print('\nReactions:')
for reaction in sbml_model.getListOfReactions():
    reactants = ' + '.join(['%s %s'%(int(r.getStoichiometry()) if r.getStoichiometry() > 1 else '', r.getSpecies()) for r in reaction.getListOfReactants()])
    products  = ' + '.join(['%s %s'%(int(r.getStoichiometry()) if r.getStoichiometry() > 1 else '', r.getSpecies()) for r in reaction.getListOfProducts()])
    reversible = '<' if reaction.getReversible() else ''
    print('%3s: %10s %1s->%10s\t\t[%s]' % (reaction.getName(), 
                        reactants,
                        reversible,
                         products,
                        libsbml.formulaToL3String(reaction.getKineticLaw().getMath())))
    
print('Parameters:')
pprint([(p.getId(),p.getName()) for p in sbml_model.getListOfParameters()])

# Create an SbmlImporter instance for our SBML model
sbml_importer = amici.SbmlImporter(sbml_file)

"""
Explanation: Model Loading
Here we load a simple model of protein phosphorylation that can be inhibited by a drug. This model was created using PySB (see createModel.py)
End of explanation
"""


from IPython.display import IFrame
IFrame('https://amici.readthedocs.io/en/latest/glossary.html#term-fixed-parameters', width=600, height=175)

fixedParameters = ['DRUG_0','KIN_0']

"""
Explanation: For this example we want specify the initial drug and kinase concentrations as experimental conditions. Accordingly we specify them as fixedParameters. The meaning of fixedParameters is defined in the Glossary, which we display here for convenience.
End of explanation
"""


# Retrieve model output names and formulae from AssignmentRules and remove the respective rules
observables = amici.assignmentRules2observables(
        sbml_importer.sbml, # the libsbml model object
        filter_function=lambda variable: variable.getName() == 'pPROT' 
    )
print('Observables:')
pprint(observables)

"""
Explanation: The SBML model specifies a single observable named pPROT which describes the fraction of phosphorylated Protein. We load this observable using amici.assignmentRules2observables.
End of explanation
"""


sbml_importer.sbml2amici(model_name, 
                         model_output_dir, 
                         verbose=False,
                         observables=observables,
                         constant_parameters=fixedParameters)
# load the generated module
model_module = amici.import_model_module(model_name, model_output_dir)

"""
Explanation: Now the model is ready for compilation using sbml2amici. Note that we here pass fixedParameters as arguments to constant_parameters, which ensures that amici is aware that we want to have them as fixedParameters:
End of explanation
"""


# Create Model instance
model = model_module.getModel()

# Create solver instance
solver = model.getSolver()

"""
Explanation: To simulate the model we need to create an instance via the getModel() method in the generated model module.
End of explanation
"""


# Run simulation using default model parameters and solver options
model.setTimepoints(np.linspace(0, 60, 60)) 
rdata = amici.runAmiciSimulation(model, solver)
amici.plotting.plotObservableTrajectories(rdata)

"""
Explanation: The only thing we need to simulate the model is a timepoint vector, which can be specified using the setTimepoints method. If we do not specify any additional options, the default values for fixedParameters and parameters that were specified in the SBML file will be used.
End of explanation
"""


edata = amici.ExpData(rdata, 0.1, 0.0)
edata.fixedParameters = [0,2]
rdata = amici.runAmiciSimulation(model, solver, edata)
amici.plotting.plotObservableTrajectories(rdata)

"""
Explanation: Simulation options can be specified either in the Model or in an ExpData instance. The ExpData instance can also carry experimental data. To initialize an ExpData instance from simulation results, amici offers some convenient constructors. In the following we will initialize an ExpData object from simulation results, but add noise with standard deviation 0.1 and specify the standard deviation accordingly. Moreover, we will specify custom values for DRUG_0=0 and KIN_0=2. If fixedParameter is specified in an ExpData instance, runAmiciSimulation will use those parameters instead of the ones specified in the Model instance.
End of explanation
"""


edata.fixedParametersPreequilibration = [3,0]
rdata = amici.runAmiciSimulation(model, solver, edata)
amici.plotting.plotObservableTrajectories(rdata)

"""
Explanation: For many biological systems, it is reasonable to assume that they start in a
 steady state. In this example we want to specify an experiment where a pretreatment with a drug is performed before the kinase is added. We assume that the pretreatment is sufficiently long such that the system reaches steadystate before the kinase is added. To implement this in amici, we can specify fixedParametersPreequilibration in the ExpData object. This automatically adds a preequilibration phase where the model is run to steadystate, before regular simulation starts. Here we set DRUG_0=3 and KIN_0=0 for the preequilibration. This means that there is no kinase available in the preequilibration phase.
End of explanation
"""


edata.reinitializeFixedParameterInitialStates = True

"""
Explanation: The resulting trajectory is definitely not what one may expect. The problem is that the DRUG_0 and KIN_0 set initial conditions for species in the model. By default these initial conditions are only applied at the very beginning of the simulation, i.e., before the preequilibration. Accordingly, the fixedParameters that we specified do not have any effect. To fix this, we need to set the reinitializeFixedParameterInitialStates attribue to True, to spefify that AMICI reinitializes all states that have fixedParameter-dependent initial states.
End of explanation
"""


rdata = amici.runAmiciSimulation(model, solver, edata)
amici.plotting.plotObservableTrajectories(rdata)

"""
Explanation: With this option activated, the kinase concentration will be reinitialized after the preequilibration and we will see the expected change in fractional phosphorylation:
End of explanation
"""


edata.t_presim = 10
edata.fixedParametersPresimulation = [10.0, 0.0]
edata.fixedParameters = [10.0, 2.0]
print(edata.fixedParametersPreequilibration)
print(edata.fixedParametersPresimulation)
print(edata.fixedParameters)
rdata = amici.runAmiciSimulation(model, solver, edata)
amici.plotting.plotObservableTrajectories(rdata)

"""
Explanation: On top of preequilibration, we can also specify presimulation. This option can be used to specify pretreatments where the system is not assumed to reach steadystate. Presimulation can be activated by specifying t_presim and edata.fixedParametersPresimulation. If both fixedParametersPresimulation and fixedParametersPreequilibration are specified, preequilibration will be performed first, followed by presimulation, followed by regular simulation. For this example we specify DRUG_0=10 and KIN_0=0 for the presimulation and DRUG_0=10 and KIN_0=2 for the regular simulation. We do not overwrite the DRUG_0=3 and KIN_0=0 that was previously specified for preequilibration.
End of explanation
"""


