from jax import jit
import jax.numpy as jnp
from jax.random import PRNGKey as Key
import numpy as np
import numpyro
from numpyro import distributions as dist
from numpyro.infer import log_likelihood, MCMC, NUTS, Predictive
import pandas as pd
from scipy.special import logsumexp
from metaDMG.fit import fit_utils


#%%

numpyro.enable_x64()

priors = fit_utils.get_priors()
A_prior = priors["A"]  # mean = 0.2, concentration = 5
q_prior = priors["q"]  # mean = 0.2, concentration = 5
c_prior = priors["c"]  # mean = 0.1, concentration = 10
phi_prior = priors["phi"]

#%%


def model_PMD(x, N, k=None):
    x_abs = jnp.abs(x)

    A = numpyro.sample("A", dist.Beta(A_prior[0], A_prior[1]))
    q = numpyro.sample("q", dist.Beta(q_prior[0], q_prior[1]))
    c = numpyro.sample("c", dist.Beta(c_prior[0], c_prior[1]))
    Dx = jnp.clip(numpyro.deterministic("Dx", A * (1 - q) ** (x_abs - 1) + c), 0, 1)

    delta = numpyro.sample("delta", dist.Exponential(1 / phi_prior[1]))
    phi = numpyro.deterministic("phi", delta + phi_prior[0])

    alpha = numpyro.deterministic("alpha", Dx * phi)
    beta = numpyro.deterministic("beta", (1 - Dx) * phi)

    numpyro.sample("obs", dist.BetaBinomial(alpha, beta, N), obs=k)


def model_null(x, N, k=None):
    c = numpyro.sample("c", dist.Beta(c_prior[0], c_prior[1]))
    Dx = numpyro.deterministic("Dx", c)
    delta = numpyro.sample("delta", dist.Exponential(1 / phi_prior[1]))
    phi = numpyro.deterministic("phi", delta + phi_prior[0])

    alpha = numpyro.deterministic("alpha", Dx * phi)
    beta = numpyro.deterministic("beta", (1 - Dx) * phi)

    numpyro.sample("obs", dist.BetaBinomial(alpha, beta, N), obs=k)


#%%


def filter_out_k(data):
    return {key: value for key, value in data.items() if key != "k"}


def is_model_PMD(model):
    name = model.__name__.lower()
    if "pmd" in name:
        return True
    elif "null" in name:
        return False
    raise AssertionError(f"Model should be PMD or null, got {model}")


#%%


@jit
def _get_posterior_PMD(rng_key, samples, *args, **kwargs):
    return Predictive(model_PMD, samples)(rng_key, *args, **kwargs)


@jit
def _get_posterior_null(rng_key, samples, *args, **kwargs):
    return Predictive(model_null, samples)(rng_key, *args, **kwargs)


def get_posterior_predictive(mcmc, data):
    posterior_samples = mcmc.get_samples()
    rng_key = Key(0)
    data_no_k = filter_out_k(data)
    if is_model_PMD(mcmc.sampler.model):
        return _get_posterior_PMD(rng_key, posterior_samples, **data_no_k)
    else:
        return _get_posterior_null(rng_key, posterior_samples, **data_no_k)


def get_posterior_predictive_obs(mcmc, data):
    return get_posterior_predictive(mcmc, data)["obs"]


#%%


def compute_posterior(
    mcmc, data, func_avg=np.mean, func_dispersion=lambda x: np.std(x, axis=0)
):
    """func = central tendency function, e.g. np.mean or np.median"""
    posterior_predictive = get_posterior_predictive(mcmc, data)
    predictions_fraction = posterior_predictive["obs"] / data["N"]
    y_average = func_avg(predictions_fraction, axis=0)
    # y_dispersion = numpyro.diagnostics.hpdi(predictions_fraction, prob=0.68)
    y_dispersion = func_dispersion(predictions_fraction)
    return y_average, y_dispersion


def compute_D_max(mcmc, data):
    posterior = get_posterior_predictive(mcmc, data)
    c = mcmc.get_samples()["c"]
    f = posterior["obs"] / data["N"]
    f = f[:, 0]
    D_max_samples = f - c
    D_max_mu = np.mean(D_max_samples).item()
    D_max_std = np.std(D_max_samples).item()

    # A = np.mean(mcmc.get_samples()["A"]).item()
    # phi = np.mean(mcmc.get_samples()["phi"]).item()
    # N = data["N"][0]
    # np.sqrt(A*(1-A)*(phi+N)/((phi+1)*N))

    return {"mu": D_max_mu, "std": D_max_std}


#%%


@jit
def _compute_log_likelihood_PMD(posterior_samples, data):
    return log_likelihood(model_PMD, posterior_samples, **data)["obs"]


@jit
def _compute_log_likelihood_null(posterior_samples, data):
    return log_likelihood(model_null, posterior_samples, **data)["obs"]


def compute_log_likelihood(mcmc, data):
    posterior_samples = mcmc.get_samples()
    if is_model_PMD(mcmc.sampler.model):
        return _compute_log_likelihood_PMD(posterior_samples, data)
    else:
        return _compute_log_likelihood_null(posterior_samples, data)


#%%


def get_lppd_and_waic(mcmc, data):
    d_results = {}
    # get the log likehood for each (point, num_samples)
    logprob = np.asarray(compute_log_likelihood(mcmc, data))
    # lppd for each observation
    lppd_i = logsumexp(logprob, 0) - np.log(logprob.shape[0])
    d_results["lppd_i"] = lppd_i
    # lppd
    lppd = lppd_i.sum()
    d_results["lppd"] = lppd
    # waic penalty for each observation
    pWAIC_i = np.var(logprob, 0)
    d_results["pWAIC_i"] = pWAIC_i
    # waic penalty # the effective number of parameters penalty
    pWAIC = pWAIC_i.sum()
    d_results["pWAIC"] = pWAIC
    # waic  for each observation
    waic_i = -2 * (lppd_i - pWAIC_i)
    d_results["waic_i"] = waic_i
    # waic # prediction of  out-of-sample deviance
    waic = waic_i.sum()
    d_results["waic"] = waic
    # standard error of waic
    # waic_vec = -2 * (lppd_i - pWAIC_i)
    # waic_uncertainty = jnp.sqrt(logprob.shape[1] * jnp.var(waic_vec))
    return d_results


#%%


def get_mean_of_variable(mcmc, variable, axis=0):
    values = mcmc.get_samples()[variable]
    return np.mean(values, axis=axis).item()


def get_std_of_variable(mcmc, variable, axis=0):
    values = mcmc.get_samples()[variable]
    return np.std(values, axis=axis).item()


def compute_rho_Ac(mcmc):
    A = mcmc.get_samples()["A"]
    c = mcmc.get_samples()["c"]
    return np.corrcoef(A, c)[0, 1]


def compute_z(d_results_PMD, d_results_null):
    n = len(d_results_PMD["waic_i"])
    waic_i_PMD = d_results_PMD["waic_i"]
    waic_i_null = d_results_null["waic_i"]
    dse = np.sqrt(n * np.var(waic_i_PMD - waic_i_null))
    d_waic = d_results_null["waic"] - d_results_PMD["waic"]
    z = d_waic / dse
    return z.item()


#%%


def init_mcmc(model, **kwargs):

    mcmc_kwargs = dict(
        progress_bar=False,
        num_warmup=500,
        num_samples=1000,
        num_chains=1,  # problem when setting to 2
        chain_method="sequential",
        # http://num.pyro.ai/en/stable/_modules/numpyro/infer/mcmc.html#MCMC
    )

    return MCMC(NUTS(model), jit_model_args=True, **mcmc_kwargs, **kwargs)


def init_mcmc_PMD(**kwargs):
    mcmc_PMD = init_mcmc(model_PMD, **kwargs)
    return mcmc_PMD


def init_mcmc_null(**kwargs):
    mcmc_null = init_mcmc(model_null, **kwargs)
    return mcmc_null


def init_mcmcs(config, **kwargs):
    if config["bayesian"]:
        mcmc_PMD = init_mcmc_PMD(**kwargs)
        mcmc_null = init_mcmc_null(**kwargs)
    else:
        mcmc_PMD = None
        mcmc_null = None
    return mcmc_PMD, mcmc_null


def fit_mcmc(mcmc, data, seed=0):
    mcmc.run(Key(seed), **data)


def use_last_state_as_warmup_state(mcmc):
    # https://github.com/pyro-ppl/numpyro/issues/539
    mcmc._warmup_state = mcmc._last_state


def add_Bayesian_fit_result(
    fit_result,
    data,
    mcmc_PMD,
    mcmc_null,
):

    d_results_PMD = get_lppd_and_waic(mcmc_PMD, data)
    d_results_null = get_lppd_and_waic(mcmc_null, data)

    z = compute_z(d_results_PMD, d_results_null)
    fit_result["Bayesian_z"] = z

    D_max = compute_D_max(mcmc_PMD, data)
    fit_result["Bayesian_D_max"] = D_max["mu"]
    fit_result["Bayesian_D_max_std"] = D_max["std"]

    fit_result["Bayesian_A"] = get_mean_of_variable(mcmc_PMD, "A")
    fit_result["Bayesian_A_std"] = get_std_of_variable(mcmc_PMD, "A")

    fit_result["Bayesian_q"] = get_mean_of_variable(mcmc_PMD, "q")
    fit_result["Bayesian_q_std"] = get_std_of_variable(mcmc_PMD, "q")

    fit_result["Bayesian_c"] = get_mean_of_variable(mcmc_PMD, "c")
    fit_result["Bayesian_c_std"] = get_std_of_variable(mcmc_PMD, "c")

    fit_result["Bayesian_phi"] = get_mean_of_variable(mcmc_PMD, "phi")
    fit_result["Bayesian_phi_std"] = get_std_of_variable(mcmc_PMD, "phi")

    fit_result["Bayesian_rho_Ac"] = compute_rho_Ac(mcmc_PMD)


def make_fits(fit_result, data, mcmc_PMD, mcmc_null):
    fit_mcmc(mcmc_PMD, data)
    fit_mcmc(mcmc_null, data)
    add_Bayesian_fit_result(fit_result, data, mcmc_PMD, mcmc_null)
    # mcmc_PMD.print_summary(prob=0.68)

    # if False:
    #     use_last_state_as_warmup_state(mcmc_PMD)
    #     use_last_state_as_warmup_state(mcmc_null)
