"""Methods to communicate with PyMC."""

import warnings
from copy import copy
from sys import modules

import numpy as np

try:
    from pymc.pytensorf import compile, join_nonshared_inputs
    from pymc.util import get_untransformed_name, is_transformed_name
    from pytensor import function
    from pytensor.graph.traversal import ancestors
    from pytensor.tensor import TensorConstant, matrix
except ImportError:
    warnings.warn("PyMC not installed. PyMC related functions will not work.")

from preliz.internal.distribution_helper import get_distributions


def back_fitting_pymc(prior, preliz_model, var_info, new_families=None):
    """
    Fit the samples from prior into user provided model's prior.

    From the perspective of ppe "prior" is actually an approximated posterior
    but from the users perspective is its prior.
    We need to "backfit" because we can not use arbitrary samples as priors.
    We need probability distributions.
    """
    new_priors = {}
    for rv_name, (_, size, *_) in var_info.items():
        if size > 1:
            params = []
            for i in range(size):
                opt_values = prior[rv_name][:, i]
                # Not sure how to fit alternative families.
                dist = preliz_model[rv_name]
                dist._fit_mle(opt_values)
                params.append(dist.params)
            dist._parametrization(*[np.array(x) for x in zip(*params)])
        else:
            opt_values = prior[rv_name]
            dists = set_families(preliz_model[rv_name], rv_name, new_families)
            mle = getattr(modules["preliz.unidimensional"], "mle")
            idx, _ = mle(dists, opt_values, plot=False)
            dist = dists[idx[0]]

        new_priors[rv_name] = dist

    return new_priors


def set_families(dist, var, new_families):
    dists = [dist]
    if new_families is not None:
        if new_families == "auto":
            alt = [
                getattr(modules["preliz.distributions"], d)
                for d in ["Normal", "HalfNormal", "Gamma"]
            ]
            dists += [a for a in alt if dist.__class__.__name__ != a.__class__.__name__]
        elif isinstance(new_families, list):
            dists += new_families
        elif isinstance(new_families, dict):
            dists += new_families.get(var, [])

    return dists


def compile_mllk(model):
    """
    Compile the log-likelihood for a pymc model.

    The compiled function allow us to condition on both data and parameters.
    """
    obs_rvs = model.observed_RVs[0]
    old_y_value = model.rvs_to_values[obs_rvs]
    new_y_value = obs_rvs.type()
    model.rvs_to_values[obs_rvs] = new_y_value

    vars_ = model.value_vars
    initial_point = model.initial_point()

    [logp], raveled_inp = join_nonshared_inputs(
        point=initial_point, outputs=[model.datalogp], inputs=vars_
    )

    rv_logp_fn = compile([raveled_inp, new_y_value], logp)
    rv_logp_fn.trust_input = True

    def fmodel(params, obs):
        return -rv_logp_fn(params, obs).sum()

    return fmodel, old_y_value, obs_rvs


def get_initial_guess(model):
    """Get initial guess for optimization routine."""
    return np.concatenate([np.ravel(value) for value in model.initial_point().values()])


def extract_preliz_distributions(model):
    """
    Extract the corresponding PreliZ distributions from a PyMC model.

    Parameters
    ----------
    model : a PyMC model

    Returns
    -------
    preliz_model : a dictionary of RVs names as keys and PreliZ distributions as values
    num_draws : the sample size of the observed
    """
    all_distributions = [
        dist
        for dist in modules["preliz.distributions"].__all__
        if dist not in ["Truncated", "Censored", "Hurdle", "Mixture"]
    ]
    pymc_to_preliz = dict(
        zip([dist.lower() for dist in all_distributions], get_distributions(all_distributions)),
    )

    preliz_model = {}
    for r_v in model.free_RVs:
        dist_name = (
            r_v.owner.op.name if r_v.owner.op.name else str(r_v.owner.op).split("RV", 1)[0].lower()
        )
        dist = copy(pymc_to_preliz[dist_name])
        preliz_model[r_v.name] = dist

    return preliz_model


def retrieve_variable_info(model):
    """Get shape, size, transformation and parents of each free RV in a PyMC model."""
    var_info = {}
    initial_point = model.initial_point()
    for v_var in model.value_vars:
        name = v_var.name
        rvs = model.values_to_rvs[v_var]
        nc_parents = non_constant_parents(rvs, model)
        idx_parents = []
        if nc_parents:
            idx_parents = [model.free_RVs.index(var_) for var_ in nc_parents]

        if is_transformed_name(name):
            name = get_untransformed_name(name)
            x_var = matrix(f"{name}_transformed")
            z_var = model.rvs_to_transforms[rvs].backward(x_var)
            transformation = function(inputs=[x_var], outputs=z_var)
        else:
            transformation = None

        var_info[name] = (
            initial_point[v_var.name].shape,
            initial_point[v_var.name].size,
            transformation,
            idx_parents,
        )

    num_draws = model.observed_RVs[0].eval().size

    return var_info, num_draws


def unravel_projection(prior_array, var_info, iterations):
    size = 0
    prior_dict = {}
    for key, values in var_info.items():
        shape, new_size, transformation, _ = values
        vector = prior_array[:, size : size + new_size]
        if transformation is not None:
            vector = transformation(vector)
        prior_dict[key] = vector.reshape(iterations, *shape).squeeze()
        size += new_size

    return prior_dict


def write_pymc_string(new_priors, var_info):
    """
    Return a string with the new priors for the PyMC model.

    So the user can copy and paste, ideally with none to minimal changes.
    """
    header = "with pm.Model() as model:\n"
    variables = []
    names = list(new_priors.keys())
    for key, value in new_priors.items():
        idxs = var_info[key][-1]
        if idxs:
            for i in idxs:
                nkey = names[i]
                cp_dist = copy(new_priors[nkey])
                cp_dist._fit_moments(np.mean(value.mean()), np.mean(value.std()))

                dist_name, dist_params = repr(cp_dist).split("(")
                size = var_info[nkey][1]
                if size > 1:
                    dist_params = dist_params.split(")")[0]
                    # fmt: off
                    variables[i] = f'    {nkey:} = pm.{dist_name}("{nkey}", {dist_params}, shape={size})\n' # noqa: E501
                    # fmt: on
                else:
                    variables[i] = f'    {nkey:} = pm.{dist_name}("{nkey}", {dist_params}\n'

        else:
            dist_name, dist_params = repr(value).split("(")
            size = var_info[key][1]
            if size > 1:
                dist_params = dist_params.split(")")[0]
                variables.append(
                    f'    {key:} = pm.{dist_name}("{key}", {dist_params}, shape={size})\n'
                )
            else:
                variables.append(f'    {key:} = pm.{dist_name}("{key}", {dist_params}\n')

    return "".join([header] + variables)


def non_constant_parents(rvs, model):
    """Find the parents of a variable that are not constant."""
    parents = []
    for variable in rvs.get_parents()[0].inputs[2:]:
        if not isinstance(variable, TensorConstant):
            for free_rv in model.free_RVs:
                if free_rv in list(ancestors([variable])) and free_rv not in parents:
                    parents.append(free_rv)
    return parents


def if_pymc_get_preliz(dist):
    """Check if dist is a PyMC or Prior distribution and if so convert to PreliZ."""
    if dist.__class__.__name__ == "TensorVariable":
        dist = from_pymc(dist)
    elif dist.__class__.__name__ == "Prior":
        dist = from_prior(dist)
    return dist


def from_pymc(dist):
    """Convert a PyMC distribution to a PreliZ distribution.

    Parameters
    ----------
    dist : PyMC distribution

    Returns
    -------
    PreliZ distribution
    """
    name = dist.owner.op._print_name[0]
    if name == "MultivariateNormal":
        name = "MvNormal"

    if name == "Censored":
        base_dist = dist.owner.inputs[0]
        lower = _as_scalar(dist.owner.inputs[1].eval())
        upper = _as_scalar(dist.owner.inputs[2].eval())
        if np.isnan(lower):
            lower = -np.inf
        if np.isnan(upper):
            upper = np.inf

        BaseDist = from_pymc(base_dist)
        return modules["preliz.distributions"].Censored(BaseDist, lower=lower, upper=upper)

    if "Truncated" in name and name != "TruncatedNormal":
        base_dist_name = name.split("Truncated")[1]
        base_params = [v.eval() for v in dist.owner.inputs[2:-2]]
        lower = _as_scalar(dist.owner.inputs[-2].eval())
        upper = _as_scalar(dist.owner.inputs[-1].eval())
        if np.isnan(lower):
            lower = -np.inf
        if np.isnan(upper):
            upper = np.inf

        BaseDist = getattr(modules["preliz.distributions"], base_dist_name)
        return modules["preliz.distributions"].Truncated(
            _reparametrize(BaseDist, base_dist_name, base_params), lower=lower, upper=upper
        )

    elif name == "Mixture":
        name_0 = dist.owner.inputs[2].owner.op._print_name[0]
        if name_0 == "DiracDelta":
            base_node = dist.owner.inputs[-1]
            base_name = base_node.owner.op._print_name[0]
            base_params = [v.eval() for v in base_node.owner.inputs[2:]]
            psi = _nan_to_none(dist.owner.inputs[1].eval())[1]
            ZeroInflated = getattr(modules["preliz.distributions"], f"ZeroInflated{base_name}")
            if base_name == "NegativeBinomial":
                n, p = base_params
                mu = n * (1 - p) / p
                base_params = [mu, n]

            base_params = _nan_to_none(base_params)
            return ZeroInflated(psi, *base_params)

        else:
            components = dist.owner.inputs[2:]
            weights = _nan_to_none(dist.owner.inputs[1].eval())
            PreliZ_components = [from_pymc(comp) for comp in components]
            return modules["preliz.distributions"].Mixture(PreliZ_components, weights=weights)

    elif name == "Hurdle":
        base_type_name = dist.owner.inputs[-1].owner.op._print_name[0].replace("Truncated", "")
        psi = _nan_to_none(dist.owner.inputs[1].eval())[-1]
        base_params = [v.eval() for v in dist.owner.inputs[-1].owner.inputs[2:]]
        BaseDist = getattr(modules["preliz.distributions"], base_type_name)
        return getattr(modules["preliz.distributions"], "Hurdle")(
            _reparametrize(BaseDist, base_type_name, base_params), psi
        )

    else:
        if name == "HalfNormal":
            params_inputs = [v.eval() for v in dist.owner.inputs[3:]]
        elif name == "LogitNormal":
            params_inputs = [v.eval() for v in dist.owner.inputs[1:]]
        else:
            params_inputs = [v.eval() for v in dist.owner.inputs[2:]]

        params_inputs = [
            p for p in params_inputs if isinstance(p, (int, float, np.number, np.ndarray))
        ]

        try:
            Dist = getattr(modules["preliz.distributions"], name)
        except AttributeError:
            raise NotImplementedError(f"No PreliZ distribution named {name}")

        return _reparametrize(Dist, name, params_inputs)


def from_prior(prior):
    """Convert a Prior distribution (from pymc-extras).

    Parameters
    ----------
    dist : Prior distribution

    Returns
    -------
    PreliZ distribution
    """
    dist_name = prior.distribution
    kwargs = prior.to_dict().get("kwargs", {})

    if dist_name in ["Truncated", "Censored"]:
        dist_arg = kwargs["dist"]
        base = getattr(modules["preliz.distributions"], dist_arg["dist"])(
            **dist_arg.get("kwargs", {})
        )
        return getattr(modules["preliz.distributions"], dist_name)(
            base, lower=kwargs.get("lower"), upper=kwargs.get("upper")
        )

    if dist_name.startswith("Hurdle"):
        base_name = dist_name.replace("Hurdle", "")
        base_kwargs = {k: v for k, v in kwargs.items() if k != "psi"}
        base = getattr(modules["preliz.distributions"], base_name)(**base_kwargs)
        return getattr(modules["preliz.distributions"], "Hurdle")(base, psi=kwargs["psi"])

    if dist_name == "Mixture":
        comp_dists = [from_prior(d) for d in kwargs["comp_dists"]]
        w = kwargs["w"]
        return getattr(modules["preliz.distributions"], "Mixture")(comp_dists, weights=w)

    if hasattr(modules["preliz.distributions"], dist_name):
        return getattr(modules["preliz.distributions"], dist_name)(**kwargs)

    raise ValueError(f"Unknown PreliZ distribution: {dist_name}")


def _as_scalar(x):
    x = np.asarray(x)
    return x.item() if x.shape == () or x.size == 1 else x


def _reparametrize(Dist, name, params_inputs):
    if name == "AsymmetricLaplace":
        b, kappa, mu = _nan_to_none(params_inputs)
        return Dist(mu=mu, b=b, kappa=kappa)
    if name == "Exponential":
        scale = _nan_to_none(params_inputs)[0]
        if scale is not None:
            lam_ = 1 / scale
        else:
            lam_ = None
        return Dist(lam_)
    if name == "Gamma":
        alpha, inv_beta = _nan_to_none(params_inputs)
        if inv_beta is not None:
            beta = 1 / inv_beta
        else:
            beta = None
        return Dist(alpha=alpha, beta=beta)
    if name == "Rice":
        b, sigma = params_inputs
        nu, sigma = _nan_to_none((b * sigma, sigma))
        return Dist(nu=nu, sigma=sigma)
    if name == "SkewStudentT":
        a, b, mu, sigma = _nan_to_none(params_inputs)
        return Dist(mu=mu, sigma=sigma, a=a, b=b)
    if name == "Wald":
        mu, lam, _ = _nan_to_none(params_inputs)
        return Dist(mu=mu, lam=lam)
    if name == "BetaBinomial":
        n, alpha, beta = _nan_to_none(params_inputs)
        return Dist(alpha=alpha, beta=beta, n=n)
    if name == "NegativeBinomial":
        n, p = _nan_to_none(params_inputs)
        if p is not None and n is not None:
            mu = n * (1 - p) / p
        else:
            mu = None
        return Dist(mu=mu, alpha=n)
    if name == "HyperGeometric":
        good, bad, n = _nan_to_none(params_inputs)
        if good is not None and bad is not None:
            N = good + bad
        else:
            N = None
        return Dist(N=N, k=good, n=n)

    return Dist(*_nan_to_none(params_inputs))


def _nan_to_none(params):
    if np.isscalar(params):
        return None if np.isnan(params) else params

    result = []
    for p in params:
        arr = np.asarray(p)
        if arr.size == 1:
            val = arr.item()
            result.append(None if np.isnan(val) else val)
        else:
            mask = np.isnan(arr)
            if np.any(mask):
                out = arr.astype(object)
                out[mask] = None
                result.append(out)
            else:
                result.append(p)

    return result
