from dataclasses import dataclass
from typing import Literal

import numpy as np

from pymgcv.families import SupportsCDF
from pymgcv.gam import AbstractGAM
from pymgcv.rlibs import rbase, rstats
from pymgcv.rpy_utils import is_null, to_py, to_rpy


@dataclass
class QQResult:
    """Results required for a qq plot."""

    theoretical: np.ndarray
    residuals: np.ndarray
    interval: tuple[np.ndarray, np.ndarray]  # lower, upper


def qq_simulate(
    gam: AbstractGAM,
    *,
    n_sim: int = 50,
    level: float = 0.9,
    type: Literal["deviance", "response", "pearson"] = "deviance",
) -> QQResult:
    """Generate data for qq-plot via simulation from the family.

    Args:
        gam: The fitted GAM object.
        n_sim: The number of simulations to perform.
        level: The level (i.e. 0.9 means 90% interval).
        type: The type of residuals to use.
    """
    if gam.fit_state is None:
        raise ValueError("GAM must be fitted before simulating quantiles.")
    if n_sim < 2:
        raise ValueError("n must be at least 2.")

    model = gam.fit_state.rgam
    fit = rstats.fitted(model)
    weights = rstats.weights(model, type="prior")
    sigma2 = model.rx2["sig2"]

    if is_null(sigma2):
        sigma2 = rbase.summary(model, re_test=False).rx2["dispersion"]

    fit, weights, sigma2 = to_py(fit), to_py(weights), to_py(sigma2)

    if len(gam.predictors) > 1:
        raise NotImplementedError(
            "Multivariate response families are not supported for qq_simulate.",
        )

    sims = []
    for _ in range(n_sim):
        ysim = gam.family.sample(mu=fit, wt=weights, scale=sigma2)
        res = gam.residuals_from_y_and_fit(
            y=ysim,
            fit=fit,
            weights=weights,
            type=type,
        )
        sim = np.sort(res)
        sims.append(sim)

    sims = np.stack(sims, axis=1)
    n_obs = len(fit)
    theoretical = np.quantile(
        sims,
        q=(np.arange(n_obs) + 0.5) / n_obs,
    )
    alpha = (1 - level) / 2
    interval = np.quantile(sims, q=(alpha, 1 - alpha), axis=1)
    residuals = gam.residuals(type=type)
    residuals = np.sort(residuals)

    return QQResult(
        theoretical=theoretical,
        residuals=residuals,
        interval=(interval[0], interval[1]),
    )


# TODO support normal. For this logp in cdf matters!
def qq_cdf(
    gam: AbstractGAM,
    *,
    level: float = 0.9,
) -> QQResult:
    """Generate a QQ-plot by transforming the data to uniform using the families CDF.

    In this case:

       - The theoretical values are evenly spaced between 0 and 1,
       - The residuals are the result of passing the data through the family's CDF,
         parameterized using the fitted values (and weights and scale as appropriate).

    !!! note

        The Q-Q plot generated by this function likely will not be good at for detecting deviations
        in tail behaviour/due to outliers. For example, a few outliers would be transformed to
        values close to zero or one, which will often **not** be visable in the plot!


        The Q-Q curve when using [`qq_cdf`][pymgcv.qq.qq_cdf] will naturally be a
        different shape to those produced by [`qq_simulate`][pymgcv.qq.qq_simulate]!

    Args:
        gam: The fitted GAM object. The family should support the CDF method,
            which can be checked with `isinstance(family, SupportsCDF)`.
        level: The confidence level for the interval.

    Returns:
        QQResult: The results required for a qq plot.
    """
    if gam.fit_state is None:
        raise ValueError("GAM has not been fit")

    if not isinstance(gam.family, SupportsCDF):
        raise TypeError("Family must support CDF method to use qq_cdf.")

    model = gam.fit_state.rgam

    fit = rstats.fitted(model)
    weights = rstats.weights(model, type="prior")
    sigma2 = model.rx2["sig2"]

    if is_null(sigma2):
        sigma2 = rbase.summary(model, re_test=False).rx2["dispersion"]

    fit, weights, sigma2 = to_py(fit), to_py(weights), to_py(sigma2)

    # Transform the data to uniform (CDF)
    residuals = gam.family.cdf(
        to_py(gam.fit_state.rgam.rx2["y"]),
        mu=fit,
        wt=weights,
        scale=sigma2,
    )
    n_obs = len(residuals)
    theoretical = (np.arange(n_obs) + 0.5) / n_obs
    residuals = np.sort(residuals)
    tmp = np.arange(n_obs)
    alpha = (1 - level) / 2
    lower = rstats.qbeta(alpha, to_rpy(tmp), to_rpy(n_obs - tmp + 1))
    upper = rstats.qbeta(1 - alpha, to_rpy(tmp), to_rpy(n_obs - tmp + 1))

    return QQResult(
        theoretical=theoretical,
        residuals=residuals,
        interval=(lower, upper),
    )
