class Likelihood:
    """Representation of a Likelihood function for a Bambi model.

    Notes:
    * ``parent`` must not be in ``kwargs``.
    * ``parent`` is inferred from the ``name`` if it is a known name

    Parameters
    ----------
    name: str
        Name of the likelihood function. Must be a valid PyMC3 distribution name.
    parent: str
        Optional specification of the name of the mean parameter in the likelihood.
        This is the parameter whose transformation is modeled by the linear predictor.
    kwargs:
        Keyword arguments that indicate prior distributions for auxiliary parameters in the
        likelihood.
    """

    DISTRIBUTIONS = DISTRIBUTIONS

    def __init__(self, name, parent=None, **kwargs):
        if name in self.DISTRIBUTIONS:
            self.name = name
            self.parent = parent
            self.priors = self._check_priors(kwargs)
        else:
            # On your own risk
            self.name = name
            # Check priors passed are in fact of class Prior
            check_all_are_priors(kwargs)
            self.priors = kwargs
            self.parent = parent

    @property
    def parent(self):
        return self._parent

    @parent.setter
    def parent(self, x):
        # Checks are made when using a known distribution
        if self.name in self.DISTRIBUTIONS:
            if x is None:
                x = self.DISTRIBUTIONS[self.name]["parent"]
            elif x not in self.DISTRIBUTIONS[self.name]["params"]:
                raise ValueError(f"'{x}' is not a valid parameter for the likelihood '{self.name}'")
        # Otherwise, no check is done. At your own risk!
        self._parent = x

    def _check_priors(self, priors):
        args = self.DISTRIBUTIONS[self.name]["args"]

        # The function requires priors but none were passed
        if priors == {} and args is not None:
            raise ValueError(f"'{self.name}' requires priors for the parameters {args}.")

        # The function does not require priors, but at least one was passed
        if priors != {} and args is None:
            raise ValueError(f"'{self.name}' does not require any additional prior.")

        # The function requires priors, priors were passed, but they differ from the required
        if priors and args:
            difference = set(args) - set(priors)
            if len(difference) > 0:
                raise ValueError(f"'{self.name}' misses priors for the parameters {difference}")

            # And check priors passed are in fact of class Prior
            check_all_are_priors(priors)

        return priors

    def __str__(self):
        args = [f"name: {self.name}", f"parent: {self.parent}", f"priors: {self.priors}"]
        return f"{self.__class__.__name__}({spacify(multilinify(args))}\n)"

    def __repr__(self):
        return self.__str__() class Prior:
    """Abstract specification of a term prior.

    Parameters
    ----------
    name : str
        Name of prior distribution. Must be the name of a PyMC3 distribution
        (e.g., ``"Normal"``, ``"Bernoulli"``, etc.)
    auto_scale: bool
        Whether to adjust the parameters of the prior or use them as passed. Default to ``True``.
    scale: num or str
        How to scale the prior. If a number, must be within 0 and 1. If a string, must be one of
        ``"narrow"``, ``"medium"``, ``"wide"``, or ``"superwide"``. Only used when the model was
        created with ``automatic_priors="mle"``.
    kwargs : dict
        Optional keywords specifying the parameters of the named distribution.
    """

    def __init__(self, name, auto_scale=True, scale=None, **kwargs):
        self.name = name
        self.auto_scale = auto_scale
        self.scale = scale
        self.args = {}
        self.update(**kwargs)

    def update(self, **kwargs):
        """Update the arguments of the prior with additional arguments.

        Parameters
        ----------
        kwargs : dict
            Optional keyword arguments to add to prior args.
        """
        # The backend expect numpy arrays, so make sure all numeric values are represented as such.
        kwargs_ = {}
        for key, val in kwargs.items():
            if isinstance(val, (int, float)):
                val = np.array(val)
            elif isinstance(val, np.ndarray):
                val = val.squeeze()
            kwargs_[key] = val
        self.args.update(kwargs_)

    def __eq__(self, other):
        if not isinstance(other, type(self)):
            return False
        else:
            return self.__dict__ == other.__dict__

    def __str__(self):
        args = ", ".join(
            [
                f"{k}: {np.round_(v, 4)}" if not isinstance(v, type(self)) else f"{k}: {v}"
                for k, v in self.args.items()
            ]
        )
        return f"{self.name}({args})"

    def __repr__(self):
        return self.__str__() class Bernoulli(UnivariateFamily):
    SUPPORTED_LINKS = ["identity", "logit", "probit", "cloglog"]

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
        "Sample from posterior predictive distribution"
        mean = self.link.linkinv(linear_predictor)
        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        mean = mean[:, idxs, :]
        return np.random.binomial(1, mean) class Beta(UnivariateFamily):
    SUPPORTED_LINKS = ["identity", "logit", "probit", "cloglog"]

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
        mean = self.link.linkinv(linear_predictor)
        kappa = posterior[model.response.name + "_kappa"].values

        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        mean = mean[:, idxs, :]
        kappa = kappa[:, idxs, np.newaxis]

        alpha = mean * kappa
        beta = (1 - mean) * kappa
        return np.random.beta(alpha, beta) class Binomial(UnivariateFamily):
    SUPPORTED_LINKS = ["identity", "logit", "probit", "cloglog"]

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n, trials=None):
        if trials is None:
            trials = model.response.data[:, 1]

        mean = self.link.linkinv(linear_predictor)
        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        mean = mean[:, idxs, :]
        np.random.binomial(trials.squeeze(), mean) class Gamma(UnivariateFamily):
    SUPPORTED_LINKS = ["identity", "log", "inverse"]

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
        mean = self.link.linkinv(linear_predictor)
        alpha = posterior[model.response.name + "_alpha"].values

        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        mean = mean[:, idxs, :]
        alpha = alpha[:, idxs, np.newaxis]
        beta = alpha / mean
        return np.random.gamma(alpha, 1 / beta) class Gaussian(UnivariateFamily):
    SUPPORTED_LINKS = ["identity", "log", "inverse"]

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
        "Sample from posterior predictive distribution"
        mean = self.link.linkinv(linear_predictor)
        sigma = posterior[model.response.name + "_sigma"].values

        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        mean = mean[:, idxs, :]
        sigma = sigma[:, idxs, np.newaxis]

        return np.random.normal(mean, sigma) class NegativeBinomial(UnivariateFamily):
    SUPPORTED_LINKS = ["identity", "log", "cloglog"]

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
        mean = self.link.linkinv(linear_predictor)
        n = posterior[model.response.name + "_alpha"].values

        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        mean = mean[:, idxs, :]
        n = n[:, idxs, np.newaxis]
        p = n / (mean + n)

        return np.random.negative_binomial(n, p) class Poisson(UnivariateFamily):
    SUPPORTED_LINKS = ["identity", "log"]

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
        mean = self.link.linkinv(linear_predictor)
        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        mean = mean[:, idxs, :]
        return np.random.poisson(mean) class StudentT(UnivariateFamily):
    SUPPORTED_LINKS = ["identity", "log", "inverse"]

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
        mean = self.link.linkinv(linear_predictor)
        sigma = posterior[model.response.name + "_sigma"].values

        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        if isinstance(self.likelihood.priors["nu"], (int, float)):
            nu = self.likelihood.priors["nu"]
        else:
            nu = posterior[model.response.name + "_nu"].values[:, idxs, np.newaxis]

        mean = mean[:, idxs, :]
        sigma = sigma[:, idxs, np.newaxis]

        return stats.t.rvs(nu, mean, sigma) class VonMises(UnivariateFamily):
    SUPPORTED_LINKS = ["identity", "tan_2"]

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
        mean = self.link.linkinv(linear_predictor)
        kappa = posterior[model.response.name + "_kappa"].values

        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        mean = mean[:, idxs, :]
        kappa = kappa[:, idxs, np.newaxis]

        return np.random.vonmises(mean, kappa) class Wald(UnivariateFamily):
    SUPPORTED_LINKS = ["inverse", "inverse_squared", "identity", "log"]

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
        mean = self.link.linkinv(linear_predictor)
        lam = posterior[model.response.name + "_lam"].values

        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        mean = mean[:, idxs, :]
        lam = lam[:, idxs, np.newaxis]

        return np.random.wald(mean, lam) class Categorical(MultivariateFamily):
    SUPPORTED_LINKS = ["softmax"]

    def predict(self, model, posterior, linear_predictor):
        # This is only 'softmax' for now.
        # Second axis is the one where the response coord is inserted
        # We need to append zeros for the reference category
        shape = linear_predictor.shape
        linear_predictor = np.dstack(
            (np.zeros(shape=(shape[0], shape[1], 1, shape[3])), linear_predictor)
        )

        mean = self.link.linkinv(linear_predictor, axis=2)

        obs_n = mean.shape[-1]
        name = model.response.name + "_mean"
        coord_name = model.response.name + "_obs"

        # Drop var/dim if already present
        if name in posterior.data_vars:
            posterior = posterior.drop_vars(name).drop_dims(coord_name)

        response_coord_name = model.response.name + "_mean_coord"
        coords = ("chain", "draw", response_coord_name, coord_name)
        posterior[name] = (coords, mean)
        posterior = posterior.assign_coords({response_coord_name: model.response.levels})
        posterior = posterior.assign_coords({coord_name: list(range(obs_n))})
        return posterior

    def posterior_predictive(self, model, posterior, linear_predictor, draws, draw_n):
        # https://stackoverflow.com/questions/34187130
        def draw_categorical_samples(probability_matrix, items):
            cumsum = probability_matrix.cumsum(axis=0)
            idx = np.random.rand(probability_matrix.shape[1])
            idx = (cumsum < idx).sum(axis=0)
            return items[idx]

        shape = linear_predictor.shape
        linear_predictor = np.dstack(
            (np.zeros(shape=(shape[0], shape[1], 1, shape[3])), linear_predictor)
        )

        mean = self.link.linkinv(linear_predictor, axis=2)
        idxs = np.random.randint(low=0, high=draw_n, size=draws)
        mean = mean[:, idxs, :, :]
        shape = mean.shape

        mean = mean.reshape((np.prod(mean.shape[:2]), mean.shape[2], mean.shape[3]))
        draws_n = mean.shape[0]
        obs_n = mean.shape[-1]

        pps = np.empty((draws_n, obs_n), dtype=int)
        response_levels = np.arange(len(model.response.levels))
        for idx in range(obs_n):
            pps[:, idx] = draw_categorical_samples(mean[..., idx].T, response_levels)

        return pps.reshape((shape[0], shape[1], obs_n))
from bambi.families import Likelihood
from bambi.priors import Prior
from bambi.families.univariate import (
    Bernoulli,
    Beta,
    Binomial,
    Gamma,
    Gaussian,
    NegativeBinomial,
    Poisson,
    StudentT,
    VonMises,
    Wald,
)
from bambi.families.multivariate import Categorical



## NOTE: Check docs/api_reference.rst links the right lines from this document
# Default parameters for PyMC3 distributions
SETTINGS_DISTRIBUTIONS = {
    "Bernoulli": {"p": 0.5},
    "Beta": {"alpha": 1, "beta": 1},
    "Binomial": {"n": 1, "p": 0.5},
    "Cauchy": {"alpha": 0, "beta": 1},
    "Flat": {},
    "Gamma": {"alpha": 2, "beta": 0.1},
    "HalfCauchy": {"beta": 1},
    "HalfFlat": {},
    "HalfNormal": {"sigma": 1},
    "Normal": {"mu": 0, "sigma": 1},
    "NegativeBinomial": {"alpha": 1, "mu": 1},
    "Poisson": {"mu": 1},
    "StudentT": {"lam": 1, "nu": 1},
    "VonMises": {"mu": 0, "kappa": 1},
    "Wald": {"mu": 1, "lam": 1},
}

# fmt: off
SETTINGS_FAMILIES = {
    "bernoulli": {
        "likelihood": {
            "name": "Bernoulli",
            "args": {},
            "parent": "p",
        },
        "link": "logit",
        "family": Bernoulli,
    },
    "beta": {
        "likelihood": {
            "name": "Beta",
            "args": {
                "kappa": "HalfCauchy"
            },
            "parent": "mu",
        },
        "link": "logit",
        "family": Beta,
    },
    "binomial": {
        "likelihood": {
            "name": "Binomial",
            "args": {},
            "parent": "p",
        },
        "link": "logit",
        "family": Binomial,
    },
    "categorical": {
        "likelihood": {
            "name": "Categorical",
            "args": {},
            "parent": "p",
        },
        "link": "softmax",
        "family": Categorical,
    },
    "gamma": {
        "likelihood": {
            "name": "Gamma",
            "args": {
                "alpha": "HalfCauchy"
            },
            "parent": "mu",
        },
        "link": "inverse",
        "family": Gamma,
    },
    "gaussian": {
        "likelihood": {
            "name": "Normal",
            "args": {
                "sigma": "HalfNormal"
            },
            "parent": "mu",
        },
        "link": "identity",
        "family": Gaussian,
    },
    "negativebinomial": {
        "likelihood": {
            "name": "NegativeBinomial",
            "args": {
                "alpha": "HalfCauchy"
            },
            "parent": "mu",
        },
        "link": "log",
        "family": NegativeBinomial,
    },
    "poisson": {
        "likelihood": {
            "name": "Poisson",
            "args": {},
            "parent": "mu",
        },
        "link": "log",

        "family": Poisson,