from typing import Optional, Tuple
import numpy as np
import scipy.stats as st

from ..base import Distribution, LinkFunction, ScipyMixin
from ..links import Identity, Log


class ReverseGumbel(ScipyMixin, Distribution):
    """
    The Reverse Gumbel (Type I minimum extreme value) distribution with location (mu) and scale (sigma) parameters.

    The probability density function is defined as:
    $$
    f(y | \\mu, \\sigma) = \\frac{1}{\\sigma} \\exp\\left( \\frac{y - \\mu}{\\sigma} - \\exp\\left( \\frac{y - \\mu}{\\sigma} \\right) \\right)
    $$

    This distribution corresponds to the RG() distribution in GAMLSS.

    Notes:
        - Mean = mu - digamma(1) * sigma ≈ mu - 0.5772157 * sigma
        - Variance = (pi^2 * sigma^2) / 6 ≈ 1.64493 * sigma^2
    """

    corresponding_gamlss: str = "RG"
    parameter_names = {0: "mu", 1: "sigma"}
    parameter_support = {0: (-np.inf, np.inf), 1: (np.nextafter(0, 1), np.inf)}
    distribution_support = (-np.inf, np.inf)
    scipy_dist = st.gumbel_r
    scipy_names = {"mu": "loc", "sigma": "scale"}

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

    def theta_to_scipy_params(self, theta: np.ndarray) -> dict:
        mu = theta[:, 0]
        sigma = theta[:, 1]
        return {
            "loc": mu,
            "scale": sigma,
        }

    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)
        z = (y - mu) / sigma
        ez = np.exp(-z)

        if param == 0:
            return (1 - ez) / sigma

        if param == 1:
            return -(1 / sigma) * (1 - z) - (z / sigma) * ez

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

        if param == 0:
            return -1 / sigma**2

        if param == 1:
            return -1.82368 / sigma**2  # approximation from GAMLSS

    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)
        _, sigma = self.theta_to_params(theta)
        return np.full_like(y, -0.422784 / sigma**2)

    def initial_values(self, y: np.ndarray) -> np.ndarray:
        gamma_const = 0.5772157
        sigma_init = (np.sqrt(6) * np.std(y)) / np.pi
        mu_init = np.mean(y) + gamma_const * sigma_init
        initial_params = [mu_init, sigma_init]
        return np.tile(initial_params, (y.shape[0], 1))
