import optparse
import yaml
import time
import pickle as pkl
import pandas as pd
import tensorflow as tf
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd
from tensorflow_probability import bijectors as tfb
import matplotlib.pyplot as plt

from covid.rdata import load_population, load_age_mixing, load_mobility_matrix
from covid.model import CovidUKODE, covid19uk_logp
from covid.util import *


def plotting(dates, sim):
    print("Initial R0:", simulator.eval_R0(param))
    print("Doubling time:", doubling_time(dates, sim.numpy(), '2020-02-27','2020-03-13'))

    fig = plt.figure()
    removals = tf.reduce_sum(sim[:, 3, :], axis=1)
    infected = tf.reduce_sum(sim[:, 1:3, :], axis=[1,2])
    exposed = tf.reduce_sum(sim[:, 1, :], axis=1)
    date = np.squeeze(np.where(dates == np.datetime64('2020-03-13'))[0])
    print("Daily incidence 2020-03-13:", exposed[date]-exposed[date-1])

    plt.plot(dates, removals*0.10, '-', label='10% reporting')
    plt.plot(dates, infected, '-', color='red', label='Total infected')
    plt.plot(dates, removals, '-', color='gray', label='Total recovered/detected/died')

    plt.scatter(np.datetime64('2020-03-13'), 600, label='gov.uk cases 13th March 2020')
    plt.legend()
    plt.grid(True)
    fig.autofmt_xdate()
    plt.show()


if __name__ == '__main__':

    parser = optparse.OptionParser()
    parser.add_option("--config", "-c", dest="config",
                      help="configuration file")
    options, args = parser.parse_args()
    with open(options.config, 'r') as ymlfile:
        config = yaml.load(ymlfile)

    K_tt, age_groups = load_age_mixing(config['data']['age_mixing_matrix_term'])
    K_hh, _ = load_age_mixing(config['data']['age_mixing_matrix_hol'])

    T, la_names = load_mobility_matrix(config['data']['mobility_matrix'])
    np.fill_diagonal(T, 0.)

    N, n_names = load_population(config['data']['population_size'])

    param = sanitise_parameter(config['parameter'])
    settings = sanitise_settings(config['settings'])

    case_reports = pd.read_csv(config['data']['reported_cases'])
    case_reports['DateVal'] = pd.to_datetime(case_reports['DateVal'])
    date_range = [case_reports['DateVal'].min(), case_reports['DateVal'].max()]
    y = case_reports['CumCases'].to_numpy()
    y_incr = np.round((y[1:] - y[:-1]) * 0.8)

    simulator = CovidUKODE(K_tt, K_hh, T, N, date_range[0]-np.timedelta64(1,'D'),
                           date_range[1], settings['holiday'], settings['bg_max_time'], int(settings['time_step']))

    seeding = seed_areas(N, n_names)  # Seed 40-44 age group, 30 seeds by popn size
    state_init = simulator.create_initial_state(init_matrix=seeding)

    #@tf.function
    def logp(epsilon, beta):
        p = param
        p['epsilon'] = epsilon
        p['beta1'] = beta
        epsilon_logp = tfd.Gamma(concentration=1., rate=1.).log_prob(p['epsilon'])
        beta_logp = tfd.Gamma(concentration=1., rate=1.).log_prob(p['beta1'])
        t, sim, solve = simulator.simulate(p, state_init)
        y_logp = covid19uk_logp(y_incr, sim, 0.1)
        logp = epsilon_logp + beta_logp + tf.reduce_sum(y_logp)
        return logp

    def trace_fn(_, pkr):
      return (
          pkr.inner_results.log_accept_ratio,
          pkr.inner_results.accepted_results.target_log_prob,
          pkr.inner_results.accepted_results.step_size)


    unconstraining_bijector = [tfb.Exp(), tfb.Exp()]
    initial_mcmc_state = [0.00001, 0.03]
    print("Initial log likelihood:", logp(*initial_mcmc_state))

    @tf.function
    def sample(n_samples, step_size=[0.01, 0.1]):
        return tfp.mcmc.sample_chain(
            num_results=n_samples,
            num_burnin_steps=0,
            current_state=initial_mcmc_state,
            kernel=tfp.mcmc.SimpleStepSizeAdaptation(
                tfp.mcmc.TransformedTransitionKernel(
                    inner_kernel=tfp.mcmc.HamiltonianMonteCarlo(
                        target_log_prob_fn=logp,
                        step_size=step_size,
                        num_leapfrog_steps=5),
                    bijector=unconstraining_bijector),
                num_adaptation_steps=500),
            trace_fn=lambda _, pkr: pkr.inner_results.inner_results.accepted_results.step_size)

    with tf.device("/CPU:0"):
        start = time.perf_counter()
        pi_beta, results = sample(1000, [0.1, 0.1])
        # sd = tfp.stats.stddev(pi_beta, sample_axis=1)
        # sd = sd / tf.reduce_sum(sd)
        # new_step_size = [results[-1] * sd[0], results[-1] * sd[1]]
        # print("New step size:",new_step_size)
        # pi_beta, results = sample(1000, step_size=new_step_size)
        end = time.perf_counter()
        print(f"Simulation complete in {end-start} seconds")
        print(results)



    fig, ax = plt.subplots(1, 2)
    ax[0].plot(pi_beta[0])
    ax[1].plot(pi_beta[1])
    plt.show()
    print(f"Posterior mean: {np.mean(pi_beta[0])}")

    with open('pi_beta_2020-03-15.pkl', 'wb') as f:
        pkl.dump(pi_beta, f)

    #dates = settings['start'] + t.numpy().astype(np.timedelta64)
    #plotting(dates, sim)