from typing import Optional, Tuple
from .config import QUINK_COLORS
import matplotlib.pyplot as plt
from .core import quink_setup
from scipy import stats
import numpy as np


def qq_plot(
    y: np.ndarray,
    dist: str = "normal",
    df: Optional[float] = None,
    figsize: Tuple[int, int] = (6, 5),
    scatter_color: str = "#1B5E63",
    line_color: str = "#E5A4C2",
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    title: Optional[str] = None,
    show: bool = True,
    dpi: int =100,
) -> plt.Figure | None:
    """
    Generate a QQ plot against a theoretical distribution.
    """
    y = np.asarray(y)
    y = np.sort(y)
    n = len(y)

    probs = (np.arange(1, n + 1) - 0.5) / n

    if dist == "normal":
        theoretical_q = stats.norm.ppf(probs)
        default_title = "QQ Plot vs Normal"

    elif dist == "t":
        if df is None:
            raise ValueError("df must be provided for t distribution.")
        theoretical_q = stats.t.ppf(probs, df=df)
        default_title = f"QQ Plot vs t (df={df})"

    else:
        raise ValueError("dist must be 'normal' or 't'.")

    fig, ax = quink_setup(figsize=figsize, dpi=dpi)

    ax.plot(
        theoretical_q,
        theoretical_q,
        color=QUINK_COLORS["fancy_pink"],
        zorder=3,
    )

    ax.scatter(
        theoretical_q,
        y,
        color=QUINK_COLORS["dark_green"],
        linewidths=0.3,
        alpha=0.6,
        zorder=4,
    )

    ax.set_xlabel(
        xlabel if xlabel is not None
        else "Theoretical Quantiles"
    )

    ax.set_ylabel(
        ylabel if ylabel is not None
        else "Empirical Quantiles"
    )

    ax.set_title(
        title if title is not None
        else default_title
    )

    if show:
        plt.show()
        return None

    plt.close(fig)
    return fig
