import jax
import jax.numpy as jnp
import numpy as np
from numpyro import deterministic, sample
from numpyro.distributions import Gamma, Normal, Unit

__all__ = [
    "gwpopulation_likelihood_model",
    "posterior_predictive_resample",
]


def gwpopulation_likelihood_model(
    likelihood,
    hyper_params,
    varmax=np.inf,
    apply_selection=True,
    predictive_resample=True,
    fit_keys=None,
):
    """
    A :code:`numpyro` implementation of :func:`gwpopulation.hyperpe.HyperparameterLikelihood`.

    Parameters
    ==========
    likelihood: gwpopulation.hyperpe.HyperparameterLikelihood
        The likelihood to evaluate.
    hyper_params: dict
        The population parameters to use.
    varmax: float, optional
        The maximum allowed variance in the likelihood.
    apply_selection: bool, optional
        Whether to apply selection effects. Default is :code:`True`.
    predictive_resample: bool, optional
        Whether to draw samples from the predicitve distribution.
        Default is :code:`True`.
    fit_keys: list, optional
        The keys to resample. Default is :code:`None`, which fits all keys in :code:data.
    """
    parameters, added_keys = likelihood.conversion_function(hyper_params.copy())

    weights = (
        likelihood.hyper_prior.prob(likelihood.data, **parameters)
        / likelihood.sampling_prior
    )
    expectations = jnp.mean(weights, axis=-1)
    square_expectations = jnp.mean(weights**2, axis=-1)
    variances = deterministic(
        "variances",
        (square_expectations - expectations**2)
        / (likelihood.samples_per_posterior * expectations**2),
    )
    ln_bayes_factors = deterministic("ln_bayes_factors", jnp.log(expectations))

    ln_l_events = jnp.sum(ln_bayes_factors)
    variance_events = jnp.sum(variances)

    if apply_selection:
        selector = likelihood.selection_function
        injection_weights = (
            selector.model.prob(selector.data, **parameters) / selector.data["prior"]
        )
        efficiency = jnp.sum(injection_weights) / selector.total_injections
        selection_uncertainty = (
            jnp.sum(injection_weights**2) / selector.total_injections**2
            - efficiency**2 / selector.total_injections
        )
        selection = deterministic(
            "selection", -jnp.log(efficiency) * likelihood.n_posteriors
        )
        selection_variance = deterministic(
            "selection_variance",
            selection_uncertainty / efficiency**2 * likelihood.n_posteriors**2,
        )

        total_hypervolume = selector.surveyed_hypervolume(parameters)
        vt = total_hypervolume * efficiency
        _ = sample("rate", Gamma(likelihood.n_posteriors, 1 / vt))
        if predictive_resample:
            posterior_predictive_resample(
                selector.data,
                injection_weights,
                "injections",
                shape=(likelihood.n_posteriors,),
                fit_keys=fit_keys,
            )
    else:
        selection = 0
        selection_variance = 0

    ln_l = deterministic("ln_l", ln_l_events + selection)
    variance = deterministic("variance", variance_events + selection_variance)

    if predictive_resample:
        for event in range(likelihood.n_posteriors):
            posterior_predictive_resample(
                likelihood.data, weights, event, event=event, fit_keys=fit_keys
            )
    if varmax < np.inf:
        ln_l = jnp.where(
            jnp.isnan(ln_l) | (variance > varmax),
            jnp.nan_to_num(-jnp.inf),
            ln_l,
        )
    return ln_l


def posterior_predictive_resample(
    data, weights, label, event=None, shape=(), fit_keys=None
):
    """
    Draw a sample from the predictive distribution given the weights for each sample.

    Parameters
    ----------
    data: dict
        The data dictionary.
    weights: array-like
        The weights for each sample.
    label: str
        The label to use for the new sample.
    event: int, optional
        The event number to use. Default is :code:`None`, assumes the input is
        one dimensional.
    shape: tuple, optional
        The shape of the output sample. Default is :code:`()`.
    fit_keys: list, optional
        The keys to fit. Default is :code:`None`, which fits all keys in :code:data.
    """
    if event is not None:
        weights = weights[:, event]
    weights /= jnp.sum(weights)
    n_points = int(weights.shape[0])
    idx = jax.random.choice(
        jax.random.PRNGKey(0),
        jnp.arange(n_points),
        p=weights,
        shape=shape,
        replace=False,
    )
    if fit_keys is None:
        fit_keys = data.keys()
    for key in fit_keys:
        if event is None:
            deterministic(f"{key}_{label}", data[key][idx])
        else:
            deterministic(f"{key}_{label}", data[key][event, idx])
