from typing import Tuple

import numpy as np
import scipy.special as spc
import scipy.stats as st

from ..base import Distribution, LinkFunction, ScipyMixin
from ..links import Logit


class Beta(ScipyMixin, Distribution):
    """The Beta Distribution for GAMLSS.

    The distribution function is defined as in GAMLSS as:
    $$
    f(y|\\mu,\\sigma)=\\frac{\\Gamma(\\frac{1 - \sigma^2}{\sigma^2})}
        {
        \\Gamma(\\frac{\\mu (1 - \\sigma^2)}{\\sigma^2})
        \Gamma(\\frac{(1 - \\mu) (1 - \sigma^2)}{\\sigma^2})}
        y^{\\frac{\mu (1 - \\sigma^2)}{\\sigma^2} - 1}
        (1-y)^{\\frac{(1 - \\mu) (1 - \\sigma^2)}{\\sigma^2} - 1}
    $$



    with the location and shape parameters $\mu, \sigma > 0$.

    !!! Note
        The function is parameterized as GAMLSS' BE() distribution.

        This parameterization is different to the `scipy.stats.beta(alpha, beta, loc, scale)` parameterization.

        We can use `Beta().gamlss_to_scipy(mu, sigma)` to map the distribution parameters to scipy.

    The `scipy.stats.beta()` distribution is defined as:
    $$
    f(x, \\alpha, \\beta) = \\frac{\\Gamma(\\alpha + \\beta) x^{\\alpha - 1} {(1 - x)}^{\\beta - 1}}{\Gamma(\\alpha) \\Gamma(\\beta)}
    $$

    with the paramters $\\alpha, \\beta >0$. The parameters can be mapped as follows:
    $$
    \\alpha = \\mu (1 - \\sigma^2) / \\sigma^2 \\Leftrightarrow \\mu = \\alpha / (\\alpha + \\beta)
    $$
    and
    $$
    \\beta = (1 - \\mu) (1 - \\sigma^2)/ \\sigma^2 \\Leftrightarrow \\sigma = \\sqrt{((\\alpha + \\beta + 1) )}
    $$


    Args:
        loc_link (LinkFunction, optional): The link function for $\mu$. Defaults to  LOGIT
        scale_link (LinkFunction, optional): The link function for $\sigma$. Defaults to LOGIT
    """

    corresponding_gamlss: str = "BE"

    parameter_names = {0: "mu", 1: "sigma"}
    parameter_support = {
        0: (np.nextafter(0, 1), np.nextafter(1, 0)),
        1: (np.nextafter(0, 1), np.nextafter(1, 0)),
    }
    distribution_support = (np.nextafter(0, 1), np.nextafter(1, 0))
    # Scipy equivalent and parameter mapping ondil -> scipy
    scipy_dist = st.beta
    # Theta columns do not map 1:1 to scipy parameters for beta
    # So we have to overload theta_to_scipy_params
    scipy_names = {}

    def __init__(
        self,
        loc_link: LinkFunction = Logit(),
        scale_link: LinkFunction = Logit(),
    ) -> None:
        super().__init__(links={0: loc_link, 1: scale_link})

    def theta_to_scipy_params(self, theta: np.ndarray) -> dict:
        """Map GAMLSS Parameters to scipy parameters.

        Args:
            theta (np.ndarray): parameters

        Returns:
            dict: Dict of (a, b, loc, scale) for scipy.stats.beta(a, b, loc, scale)
        """
        mu = theta[:, 0]
        sigma = theta[:, 1]
        alpha = mu * (1 - sigma**2) / sigma**2
        beta = (1 - mu) * (1 - sigma**2) / sigma**2
        params = {"a": alpha, "b": beta, "loc": 0, "scale": 1}
        return params

    def dl1_dp1(self, y: np.ndarray, theta: np.ndarray, param: int = 0) -> np.ndarray:
        self._validate_dln_dpn_inputs(y, theta, param)
        mu, sigma = self.theta_to_params(theta)

        if param == 0:
            alpha = mu * (1 - sigma**2) / sigma**2
            beta = (1 - mu) * (1 - sigma**2) / sigma**2

            return ((1 - sigma**2) / sigma**2) * (
                -spc.digamma(alpha) + spc.digamma(beta) + np.log(y) - np.log(1 - y)
            )

        if param == 1:
            alpha = mu * (1 - sigma**2) / sigma**2
            beta = (1 - mu) * (1 - sigma**2) / sigma**2

            return -(2 / sigma**3) * (
                mu * (-spc.digamma(alpha) + spc.digamma(alpha + beta) + np.log(y))
                + (1 - mu)
                * (-spc.digamma(beta) + spc.digamma(alpha + beta) + np.log(1 - y))
            )

    def dl2_dp2(self, y: np.ndarray, theta: np.ndarray, param: int = 0) -> np.ndarray:
        self._validate_dln_dpn_inputs(y, theta, param)
        mu, sigma = self.theta_to_params(theta)
        if param == 0:
            # MU
            alpha = mu * (1 - sigma**2) / sigma**2
            beta = (1 - mu) * (1 - sigma**2) / sigma**2

            return -(((1 - sigma**2) ** 2) / sigma**4) * (
                spc.polygamma(1, alpha) + spc.polygamma(1, beta)
            )

        if param == 1:
            # SIGMA
            alpha = mu * (1 - sigma**2) / sigma**2
            beta = (1 - mu) * (1 - sigma**2) / sigma**2

            return -(4 / (sigma**6)) * (
                (mu**2) * spc.polygamma(1, alpha)
                + ((1 - mu) ** 2) * spc.polygamma(1, beta)
                - spc.polygamma(1, alpha + beta)
            )

    def dl2_dpp(
        self, y: np.ndarray, theta: np.ndarray, params: Tuple[int, int] = (0, 1)
    ) -> np.ndarray:
        self._validate_dl2_dpp_inputs(y, theta, params)
        mu, sigma = self.theta_to_params(theta)

        if sorted(params) == [0, 1]:
            alpha = mu * (1 - sigma**2) / sigma**2
            beta = (1 - mu) * (1 - sigma**2) / sigma**2

            return (2 * (1 - sigma**2) / sigma**5) * (
                mu * spc.polygamma(1, alpha) - (1 - mu) * spc.polygamma(1, beta)
            )

    def initial_values(self, y: np.ndarray) -> np.ndarray:
        initial_params = [np.mean(y, axis=0), 0.5]
        return np.tile(initial_params, (y.shape[0], 1))
