import corner
import jax
from jax._src import config

config.update("jax_platforms", "cpu")

import jax.numpy as jnp
import jax.random as random
import numpy as np
from packaging import version

if version.parse(jax.__version__) >= version.parse("0.7.0"):
    # https://github.com/pyro-ppl/numpyro/issues/2051
    import jax.experimental.pjit
    from jax.extend.core.primitives import jit_p

    jax.experimental.pjit.pjit_p = jit_p

import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS
from prettyPlot.plotting import *

from bird import logger


def c_model(
    theta: list, t: np.ndarray | jnp.ndarray, t0: float, c0: float
) -> np.ndarray | jnp.ndarray:
    """
    Model of concentration vs time function involving kLa
    """
    cstar, kla = theta
    return (cstar - c0) * (1.0 - jnp.exp(-kla * (t - t0))) + c0


def bayes_step(
    time_obs: np.ndarray | jnp.ndarray,
    concentration_obs: np.ndarray | jnp.ndarray,
):
    """
    Prior and likelihood definition for MCMC sampling
    """
    # define parameters (incl. prior ranges)
    cstar = numpyro.sample("cstar", dist.Uniform(1e-8, 1000))
    kla = numpyro.sample("kla", dist.Uniform(1e-8, 1000))
    sigma = numpyro.sample("sigma", dist.Uniform(1e-8, 1))

    # initial values
    t0 = time_obs[0]
    c0 = concentration_obs[0]

    # concentration predictions
    c_pred = c_model([cstar, kla], t=time_obs, t0=t0, c0=c0)

    # likelihood uncertainty (assumed homoskedastic)
    std_c = jnp.ones(c_pred.shape[0]) * sigma

    # MCMC sampling with multivariate normal
    numpyro.sample("obs", dist.Normal(c_pred, std_c), obs=concentration_obs)


def hmc_samples_2_np_samples(hmc_samples: dict) -> np.ndarray:
    """
    Convert hmc samples generated by numpyro to numpy
    """
    labels = list(hmc_samples.keys())
    nsamples = len(hmc_samples[labels[0]])
    nparams = len(labels)
    np_hmc_samples = np.zeros((nsamples, nparams))
    labels_np = ["cstar", "kla", "sigma"]

    for ilabel, label in enumerate(labels):
        if label == "cstar":
            nplabel = labels_np.index("cstar")
        if label == "kla":
            nplabel = labels_np.index("kla")
        if label == "sigma":
            nplabel = labels_np.index("sigma")
        np_hmc_samples[:, nplabel] = np.array(hmc_samples[label])

    return np_hmc_samples


def np_samples_2_pred(
    np_hmc_samples: np.ndarray,
    time_obs: np.ndarray,
    concentration_obs: np.ndarray,
) -> dict:
    """
    Go from parameter samples to predictions.
    Useful for accuracy calculation and plotting.
    """
    # Uncertainty propagation
    nsamples = np_hmc_samples.shape[0]
    realization = []
    for i in range(nsamples):
        c_pred = c_model(
            np_hmc_samples[i, :-1], time_obs, time_obs[0], concentration_obs[0]
        )
        realization.append(c_pred)
    realization = np.array(realization)

    mean_real = np.mean(realization, axis=0)
    p90_real = np.percentile(realization, 90, axis=0)
    p10_real = np.percentile(realization, 10, axis=0)

    return {"mean_pred": mean_real, "p10_pred": p10_real, "p90_pred": p90_real}


def read_kla_cstar(hmc_samples: dict) -> dict:
    """
    go from hmc samples generated by numpyro to relevant arrays
    """
    np_hmc_samples = hmc_samples_2_np_samples(hmc_samples)
    return {
        "kla_samples": np_hmc_samples[:, 1],
        "cstar_samples": np_hmc_samples[:, 0],
    }


def compute_accuracy(
    hmc_samples: dict, time_obs: np.ndarray, concentration_obs: np.ndarray
) -> float:
    """
    Compute how well the sampled parameter explain the observed concentration data
    """
    np_hmc_samples = hmc_samples_2_np_samples(hmc_samples)
    pred_dict = np_samples_2_pred(np_hmc_samples, time_obs, concentration_obs)
    mean_real = pred_dict["mean_pred"]

    return np.mean(
        abs(mean_real - concentration_obs)
        / np.clip(concentration_obs, a_min=1e-12, a_max=None)
    )


def check_data_size(time_obs: np.ndarray, concentration_obs: np.ndarray):
    """
    Sanity checks
    """
    assert time_obs.shape == concentration_obs.shape
    assert len(time_obs) > 5
    assert len(time_obs.shape) == 1


def compute_kla(
    data_t: np.ndarray,
    data_c: np.ndarray,
    num_warmup: int = 4000,
    num_samples: int = 1000,
    max_chop: int | None = None,
    bootstrap: bool = True,
) -> dict:
    """
    Main loop to compute kla
    """
    rng_key = random.PRNGKey(0)
    rng_key, rng_key_ = random.split(rng_key)

    acc = []
    kla = []
    cstar = []
    kla_err = []
    cstar_err = []
    ind = []
    check_data_size(data_t, data_c)

    data_t_tmp = data_t.copy()
    data_c_tmp = data_c.copy()

    # Find where to start in the timeseries
    for ind_start in range(len(data_t_tmp) - 5):
        if max_chop is not None:
            max_chop -= 1
        logger.debug(f"Chopping index = {ind_start}")
        data_t = data_t_tmp[ind_start:]
        data_c = data_c_tmp[ind_start:]

        # Hamiltonian Monte Carlo (HMC) with no u turn sampling (NUTS)
        kernel = NUTS(bayes_step, target_accept_prob=0.9)
        mcmc = MCMC(
            kernel,
            num_warmup=num_warmup,
            num_samples=num_samples,
            progress_bar=False,
        )
        mcmc.run(
            rng_key_,
            time_obs=jnp.array(data_t),
            concentration_obs=jnp.array(data_c),
        )
        # mcmc.print_summary()

        # Read samples
        hmc_samples = mcmc.get_samples()
        samp_dict = read_kla_cstar(hmc_samples)
        kla_samples = samp_dict["kla_samples"]
        cstar_samples = samp_dict["cstar_samples"]

        accuracy = compute_accuracy(hmc_samples, data_t, data_c)

        ind.append(ind_start)
        acc.append(accuracy)
        kla.append(np.mean(kla_samples))
        cstar.append(np.mean(cstar_samples))
        kla_err.append(np.std(kla_samples))
        cstar_err.append(np.std(cstar_samples))

        # post_cal(hmc_samples, data_t, data_c)

        # If accuracy has significantly improved, we have chopped enough
        if len(acc) >= 3 and abs(acc[-1] - acc[-2]) < 0.1 * acc[0]:
            break
        # If we exceed the maximum number of timesteps to ignore, break
        if max_chop == 0:
            break

    bootstrapped = False
    # Data bootstrapping
    ## Include the uncertainty due to the number of data point available
    ## Reduce the number of data points following 4 arbitrary scenarios
    ## scenario 1: remove the first point
    ## scenario 2: remove the first 2 points
    ## scenario 3: remove the last point
    ## scenario 4: remove the last 2 points

    # Don't bootstrap if there aren't enough data points
    if bootstrap and len(data_t_tmp[ind[-1] :]) > 10:
        kla_boot = []
        kla_err_boot = []
        cstar_boot = []
        cstar_err_boot = []
        bootstrapped = True
        logger.info("Doing data bootstrapping")

        for i in range(4):
            logger.info(f"\t scenario {i}")
            if i == 0:
                data_t = data_t_tmp[ind_start + 1 :]
                data_c = data_c_tmp[ind_start + 1 :]
            elif i == 1:
                data_t = data_t_tmp[ind_start + 2 :]
                data_c = data_c_tmp[ind_start + 2 :]
            elif i == 2:
                data_t = data_t_tmp[ind_start:-1]
                data_c = data_c_tmp[ind_start:-1]
            elif i == 3:
                data_t = data_t_tmp[ind_start:-2]
                data_c = data_c_tmp[ind_start:-2]

            # Hamiltonian Monte Carlo (HMC) with no u turn sampling (NUTS)
            kernel = NUTS(bayes_step, target_accept_prob=0.9)
            mcmc = MCMC(
                kernel,
                num_warmup=num_warmup,
                num_samples=num_samples,
                progress_bar=False,
            )
            mcmc.run(
                rng_key_,
                time_obs=jnp.array(data_t),
                concentration_obs=jnp.array(data_c),
            )
            # mcmc.print_summary()

            # Read samples
            hmc_samples = mcmc.get_samples()
            samp_dict = read_kla_cstar(hmc_samples)
            kla_samples = samp_dict["kla_samples"]
            cstar_samples = samp_dict["cstar_samples"]

        kla_boot.append(np.mean(kla_samples))
        kla_err_boot.append(np.std(kla_samples))
        cstar_boot.append(np.mean(cstar_samples))
        cstar_err_boot.append(np.std(cstar_samples))

    else:
        kla_boot = [kla[-1]]
        kla_err_boot = [kla_err[-1]]
        cstar_boot = [cstar[-1]]
        cstar_err_boot = [cstar_err[-1]]

    return {
        "kla": np.mean(np.array(kla_boot)),
        "kla_err": np.mean(np.array(kla_err_boot)),
        "cstar": np.mean(np.array(cstar_boot)),
        "cstar_err": np.mean(np.array(cstar_err_boot)),
        "kla_noboot": np.mean(np.array([kla[-1]])),
        "kla_err_noboot": np.mean(np.array([kla_err[-1]])),
        "cstar_noboot": np.mean(np.array([cstar[-1]])),
        "cstar_err_noboot": np.mean(np.array([cstar_err[-1]])),
        "bootstrapped": bootstrapped,
    }


def post_cal(
    hmc_samples: dict, time_obs: np.ndarray, concentration_obs: np.ndarray
) -> None:
    """
    Plot MCMC results
    """
    np_hmc_samples = hmc_samples_2_np_samples(hmc_samples)
    labels_np = ["cstar", "kla", "sigma"]
    nsamples = np_hmc_samples.shape[0]
    labels = list(hmc_samples.keys())
    nparams = len(labels)

    # Distribution of the parameter samples
    fig = corner.corner(np_hmc_samples, labels=labels_np)

    # Convergence
    fig, axes = plt.subplots(nparams, sharex=True)
    for i in range(nparams):
        ax = axes[i]
        ax.plot(np_hmc_samples[:, i], "k", alpha=0.3, rasterized=True)
        ax.set_ylabel(labels[i])

    # Uncertainty propagation
    pred_dict = np_samples_2_pred(np_hmc_samples, time_obs, concentration_obs)
    mean_real = pred_dict["mean_pred"]
    p10_real = pred_dict["p10_pred"]
    p90_real = pred_dict["p90_pred"]

    fig = plt.figure()
    plt.plot(
        time_obs, mean_real, color="k", linewidth=3, label="mean predictions"
    )
    plt.plot(
        time_obs,
        p90_real,
        "--",
        color="k",
        linewidth=3,
        label="10th and 90th percentile",
    )
    plt.plot(time_obs, p10_real, "--", color="k", linewidth=3)
    plt.plot(
        time_obs,
        concentration_obs,
        "o",
        color="k",
        markersize=7,
        label="Data observed",
    )
    pretty_labels("t", "c[mol/m3]", 14)
    pretty_legend()
    plt.show()


def print_res_dict(res_dict: dict) -> None:
    """
    Log kla and cstar to screen
    """
    kla = res_dict["kla"]
    kla_err = res_dict["kla_err"]
    cs = res_dict["cstar"]
    cs_err = res_dict["cstar_err"]

    kla_nb = res_dict["kla_noboot"]
    kla_err_nb = res_dict["kla_err_noboot"]
    cs_nb = res_dict["cstar_noboot"]
    cs_err_nb = res_dict["cstar_err_noboot"]

    bs = res_dict["bootstrapped"]

    if bs:
        logger.info(f"\tkla = {kla*3600:.4g} +/- {kla_err*3600:.4g} [h-1]")
        logger.info(f"\tcstar = {cs:.4g} +/- {cs_err:.4g} [mol/m3]")
        logger.info(f"Without data bootstrap")
    logger.info(f"\tkla = {kla_nb*3600:.4g} +/- {kla_err_nb*3600:.4g} [h-1]")
    logger.info(f"\tcstar = {cs_nb:.4g} +/- {cs_err_nb:.4g} [mol/m3]")

