import os, tempfile
import logging
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np

from ionchannelABC import theoretical_population_size
from ionchannelABC import IonChannelDistance, EfficientMultivariateNormalTransition, IonChannelAcceptor
from ionchannelABC.experiment import setup
from ionchannelABC.visualization import plot_sim_results, plot_kde_matrix_custom
import myokit

from pyabc import Distribution, RV, History, ABCSMC
from pyabc.epsilon import MedianEpsilon
from pyabc.sampler import MulticoreEvalParallelSampler, SingleCoreSampler
from pyabc.populationstrategy import ConstantPopulationSize

"""
Explanation: ABC calibration of $I_\text{CaL}$ in Courtemanche model to original dataset.
End of explanation
"""


from experiments.ical_li import (li_act,
                                 li_inact_1000)
from experiments.ical_sun import sun_inact_kin_slow

modelfile = 'models/courtemanche_ical.mmt'

"""
Explanation: Initial set-up
Load experiments used in the publication:
 - Steady-state activation [Li1997]
 - Steady-state inactivation [Li1997]
 - Inactivation time constant [Sun1997]
End of explanation
"""


from ionchannelABC.visualization import plot_variables

sns.set_context('poster')

V = np.arange(-140, 50, 0.01)

cou_par_map = {'di': 'ical.d_inf',
            'fi': 'ical.f_inf',
            'dt': 'ical.tau_d',
            'ft': 'ical.tau_f'}

f, ax = plot_variables(V, cou_par_map, 'models/courtemanche_ical.mmt', figshape=(2,2))

"""
Explanation: Plot steady-state and time constant functions of original model
End of explanation
"""


observations, model, summary_statistics = setup(modelfile,
                                                li_act)

assert len(observations)==len(summary_statistics(model({})))

g = plot_sim_results(modelfile,
                     li_act)

"""
Explanation: Activation gate ($d$) calibration
Combine model and experiments to produce:
 - observations dataframe
 - model function to run experiments and return traces
 - summary statistics function to accept traces
End of explanation
"""


limits = {'ical.p1': (-100, 100),
          'ical.p2': (0, 50),
          'log_ical.p3': (-7, 3),
          'ical.p4': (-100, 100),
          'ical.p5': (0, 50)}
prior = Distribution(**{key: RV("uniform", a, b - a)
                        for key, (a,b) in limits.items()})

# Test this works correctly with set-up functions
assert len(observations) == len(summary_statistics(model(prior.rvs())))

"""
Explanation: Set up prior ranges for each parameter in the model.
See the modelfile for further information on specific parameters. Prepending `log_' has the effect of setting the parameter in log space.
End of explanation
"""


db_path = ("sqlite:///" + os.path.join(tempfile.gettempdir(), "courtemanche_ical_dgate_original.db"))

logging.basicConfig()
abc_logger = logging.getLogger('ABC')
abc_logger.setLevel(logging.DEBUG)
eps_logger = logging.getLogger('Epsilon')
eps_logger.setLevel(logging.DEBUG)

pop_size = theoretical_population_size(2, len(limits))
print("Theoretical minimum population size is {} particles".format(pop_size))

abc = ABCSMC(models=model,
             parameter_priors=prior,
             distance_function=IonChannelDistance(
                 exp_id=list(observations.exp_id),
                 variance=list(observations.variance),
                 delta=0.05),
             population_size=ConstantPopulationSize(2000),
             summary_statistics=summary_statistics,
             transitions=EfficientMultivariateNormalTransition(),
             eps=MedianEpsilon(initial_epsilon=100),
             sampler=MulticoreEvalParallelSampler(n_procs=16),
             acceptor=IonChannelAcceptor())

obs = observations.to_dict()['y']
obs = {str(k): v for k, v in obs.items()}

abc_id = abc.new(db_path, obs)

history = abc.run(minimum_epsilon=0., max_nr_populations=100, min_acceptance_rate=0.01)

"""
Explanation: Run ABC calibration
End of explanation
"""


df, w = history.get_distribution(m=0)

df.describe()

sns.set_context('poster')

mpl.rcParams['font.size'] = 14
mpl.rcParams['legend.fontsize'] = 14

g = plot_sim_results(modelfile,
                     li_act,
                     df=df, w=w)

plt.tight_layout()

import pandas as pd
N = 100
cou_par_samples = df.sample(n=N, weights=w, replace=True)
cou_par_samples = cou_par_samples.set_index([pd.Index(range(N))])
cou_par_samples = cou_par_samples.to_dict(orient='records')

sns.set_context('talk')
mpl.rcParams['font.size'] = 14
mpl.rcParams['legend.fontsize'] = 14

f, ax = plot_variables(V, cou_par_map, 
                       'models/courtemanche_ical.mmt', 
                       [cou_par_samples],
                       figshape=(2,2))
plt.tight_layout()

m,_,_ = myokit.load(modelfile)

originals = {}
for name in limits.keys():
    if name.startswith("log"):
        name_ = name[4:]
    else:
        name_ = name
    val = m.value(name_)
    if name.startswith("log"):
        val_ = np.log10(val)
    else:
        val_ = val
    originals[name] = val_

sns.set_context('paper')
g = plot_kde_matrix_custom(df, w, limits=limits, refval=originals)
plt.tight_layout()

"""
Explanation: Analysis of results
End of explanation
"""


observations, model, summary_statistics = setup(modelfile,
                                                li_inact_1000,
                                                sun_inact_kin_slow)

assert len(observations)==len(summary_statistics(model({})))

g = plot_sim_results(modelfile,
                     li_inact_1000,
                     sun_inact_kin_slow)

"""
Explanation: Voltage-dependent inactivation gate ($f$) calibration
End of explanation
"""


limits = {'log_ical.q1': (0, 3),
          'log_ical.q2': (-2, 3),
          'log_ical.q3': (-4, 0),
          'ical.q4': (-100, 100),
          'log_ical.q5': (-4, 0),
          'ical.q6': (-100, 100),
          'ical.q7': (0, 50)}
prior = Distribution(**{key: RV("uniform", a, b - a)
                        for key, (a,b) in limits.items()})

# Test this works correctly with set-up functions
assert len(observations) == len(summary_statistics(model(prior.rvs())))

"""
Explanation: Set up prior ranges for each parameter in the model.
See the modelfile for further information on specific parameters. Prepending `log_' has the effect of setting the parameter in log space.
End of explanation
"""


db_path = ("sqlite:///" + os.path.join(tempfile.gettempdir(), "courtemanche_ical_fgate_original.db"))

logging.basicConfig()
abc_logger = logging.getLogger('ABC')
abc_logger.setLevel(logging.DEBUG)
eps_logger = logging.getLogger('Epsilon')
eps_logger.setLevel(logging.DEBUG)

pop_size = theoretical_population_size(2, len(limits))
print("Theoretical minimum population size is {} particles".format(pop_size))

abc = ABCSMC(models=model,
             parameter_priors=prior,
             distance_function=IonChannelDistance(
                 exp_id=list(observations.exp_id),
                 variance=list(observations.variance),
                 delta=0.05),
             population_size=ConstantPopulationSize(2000),
             summary_statistics=summary_statistics,
             transitions=EfficientMultivariateNormalTransition(),
             eps=MedianEpsilon(initial_epsilon=100),
             sampler=MulticoreEvalParallelSampler(n_procs=8),
             acceptor=IonChannelAcceptor())

obs = observations.to_dict()['y']
obs = {str(k): v for k, v in obs.items()}

abc_id = abc.new(db_path, obs)

history = abc.run(minimum_epsilon=0., max_nr_populations=100, min_acceptance_rate=0.01)

"""
Explanation: Run ABC calibration
End of explanation
"""


history = History(db_path)

history.all_runs()

df, w = history.get_distribution(m=0)

df.describe()

sns.set_context('poster')

mpl.rcParams['font.size'] = 14
mpl.rcParams['legend.fontsize'] = 14

g = plot_sim_results(modelfile,
                     li_inact_1000,
                     sun_inact_kin_slow,
                     df=df, w=w)

plt.tight_layout()

m,_,_ = myokit.load(modelfile)

originals = {}
for name in limits.keys():
    if name.startswith("log"):
        name_ = name[4:]
    else:
        name_ = name
    val = m.value(name_)
    if name.startswith("log"):
        val_ = np.log10(val)
    else:
        val_ = val
    originals[name] = val_

sns.set_context('paper')
g = plot_kde_matrix_custom(df, w, limits=limits, refval=originals)
plt.tight_layout()

import pandas as pd
N = 100
cou_par_samples = df.sample(n=N, weights=w, replace=True)
cou_par_samples = cou_par_samples.set_index([pd.Index(range(N))])
cou_par_samples = cou_par_samples.to_dict(orient='records')

sns.set_context('poster')
mpl.rcParams['font.size'] = 14
mpl.rcParams['legend.fontsize'] = 14

f, ax = plot_variables(V, cou_par_map, 
                       'models/courtemanche_ical.mmt', 
                       [cou_par_samples],
                       figshape=(2,2))
plt.tight_layout()

"""
Explanation: Analysis of results
End of explanation
"""


