from typing import Iterable, Optional, Literal, List
from matplotlib.figure import Figure
from pandera.typing import DataFrame
from .config import QUINK_COLORS
import matplotlib.pyplot as plt
from .colors import gen_colors
from .core import quink_setup
from scipy import stats
import numpy as np


def qq_plot(
    y: Iterable[float],
    dist: Literal["normal", "t"] = "normal",
    df: Optional[float] = None,
    figsize: tuple[int, int] = (6, 5),
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    title: Optional[str] = None,
    scientific: bool = True,
    show: bool = True,
    dpi: int = 100,
) -> Optional[Figure]:
    """
    Generate a QQ plot against a Normal or Student-t distribution.

    Parameters
    ----------
    y : Iterable[float]
        Sample data.
    dist : {"normal", "t"}, default="normal"
        Theoretical distribution to compare against.
    df : float, optional
        Degrees of freedom for Student-t distribution.
        Required if dist="t".
    figsize : tuple[int, int], default=(6, 5)
        Figure size.
    scatter_color : str
        Color of scatter points.
    line_color : str
        Color of reference line.
    xlabel : str, optional
        Custom x-axis label.
    ylabel : str, optional
        Custom y-axis label.
    title : str, optional
        Custom plot title.
    scientific : bool, default=False
        If True, use scientific notation for axis tick labels.
    show : bool, default=True
        If True, display the plot. Otherwise return the Figure.
    dpi : int, default=100
        Figure resolution.

    Returns
    -------
    Optional[Figure]
        Figure object if show=False, otherwise None.
    """
    values: np.ndarray = np.asarray(y, dtype=float)
    values = values[~np.isnan(values)]

    n: int = values.size
    if n == 0:
        raise ValueError("Input array is empty after removing NaNs.")

    sorted_values: np.ndarray = np.sort(values)
    ecdf: np.ndarray = (np.arange(1, n + 1) - 0.5) / n

    mu: float = float(np.mean(values))
    sigma: float = float(np.std(values, ddof=1))

    if sigma == 0:
        raise ValueError("Standard deviation is zero. QQ plot undefined.")

    if dist == "normal":
        theoretical_std: np.ndarray = stats.norm.ppf(ecdf)
        default_title: str = "QQ Plot vs Normal"

    elif dist == "t":
        if df is None:
            raise ValueError("df must be provided for t distribution.")
        if df <= 0:
            raise ValueError("df must be positive.")

        theoretical_std = stats.t.ppf(ecdf, df=df)
        default_title = f"QQ Plot vs t (df={df})"

    else:
        raise ValueError(f"Unsupported distribution: {dist}")

    theoretical_q: np.ndarray = mu + sigma * theoretical_std

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

    ax.scatter(
        theoretical_q,
        sorted_values,
        color=QUINK_COLORS["dark_green"],
        linewidths=0.3,
        alpha=0.7,
        zorder=3,
    )

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

    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 scientific:
        ax.ticklabel_format(
            style="sci",
            axis="both",
            scilimits=(0, 0),
        )

    if show:
        plt.show()
        return None

    plt.close(fig)
    return fig


def acf_plot(
    acf_values: Iterable[float],
    n_obs: int,
    figsize: tuple[int, int] = (6, 5),
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    title: Optional[str] = None,
    scientific: bool = True,
    show: bool = True,
    dpi: int = 100,
    gamma: float = 0.95
) -> Optional[Figure]:
    """
    Plot the autocorrelation function (excluding lag 0).

    Parameters
    ----------
    acf_values : Iterable[float]
        Autocorrelation values including lag 0.
    n_obs : int
        Number of observations used to compute the ACF.
    figsize : tuple[int, int], default=(6, 5)
        Figure size.
    xlabel : str, optional
        Custom x-axis label.
    ylabel : str, optional
        Custom y-axis label.
    title : str, optional
        Custom plot title.
    scientific : bool, default=True
        Use scientific notation for tick labels.
    show : bool, default=True
        If True display the plot, otherwise return the Figure.
    dpi : int, default=100
        Figure resolution.
    gamma : float, default=0.95
        Confidence level for ACF bands.

    Returns
    -------
    Optional[Figure]
        Figure object if show=False, otherwise None.
    """

    values: np.ndarray = np.asarray(acf_values, dtype=float)
    values = values[~np.isnan(values)]

    if values.size <= 1:
        raise ValueError("ACF array must contain at least two lags.")

    # remove lag 0
    values = values[1:]
    lags = np.arange(1, len(values) + 1)

    z: float = stats.norm.ppf((1 + gamma) / 2)
    conf: float = z / np.sqrt(n_obs)

    default_title: str = f"Autocorrelation (CI={gamma:.2f})"

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

    ax.vlines(
        lags,
        0,
        values,
        color=QUINK_COLORS["lollo_green"],
        zorder=2,
    )

    ax.scatter(
        lags,
        values,
        color=QUINK_COLORS["dark_green"],
        linewidths=0.3,
        alpha=0.9,
        zorder=3,
    )

    ax.axhline(
        0,
        linewidth=0.8,
        color="white",
        alpha=0.4,
    )

    ax.axhline(
        conf,
        color=QUINK_COLORS["fancy_pink"],
        linestyle="--",
        linewidth=0.8,
        alpha=0.5,
    )

    ax.axhline(
        -conf,
        color=QUINK_COLORS["fancy_pink"],
        linestyle="--",
        linewidth=0.8,
        alpha=0.5,
    )

    step: int = max(1, len(lags) // 5)
    ax.set_xticks(lags[::step])

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

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

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

    if scientific:
        ax.ticklabel_format(
            style="sci",
            axis="both",
            scilimits=(0, 0),
        )

    if show:
        plt.show()
        return None

    plt.close(fig)
    return fig


def stackplot(
    weights: DataFrame,
    figsize: tuple[int, int] = (6, 5),
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    title: Optional[str] = None,
    show: bool = True,
    dpi: int = 100,
    colors: str | List[str] | None = None,
    step: Optional[float] = 0.05
) -> Optional[Figure]:
    """
    Plot a stacked area chart of portfolio weights.

    Parameters
    ----------
    weights : pd.DataFrame
        DataFrame with datetime index and assets as columns.
        Values represent portfolio weights.

    figsize : tuple[int, int], default=(6, 5)
        Figure size.

    xlabel : str, optional
        Custom x-axis label.

    ylabel : str, optional
        Custom y-axis label.

    title : str, optional
        Custom plot title.

    show : bool, default=True
        If True display the plot, otherwise return the Figure.

    dpi : int, default=100
        Figure resolution.

    colors : str | list[str] | None, optional
        Color specification:
        - list[str]: explicit list of hex colors, length must match columns
        - str:
            * hex color (e.g. "#2E5F73") -> generate palette around base color
            * "default" -> use matplotlib default colors
        - None: same as "default"

    step : float, default=0.05
        Hue increment used when generating colors from a base color.

    Returns
    -------
    Optional[Figure]
        Figure object if show=False, otherwise None.
    """
    if isinstance(colors, list):
        if len(colors) != len(weights.columns):
            raise ValueError("colors dont match columns length.")
        c = colors

    elif colors == "default" or colors is None:
        c = None

    elif isinstance(colors, str):
        c = gen_colors(colors, n=len(weights.columns), step=step)

    else:
        raise ValueError("Invalid colors argument.")

    weights = weights.sort_values(by=weights.index[-1], axis=1)
    weights = weights.fillna(0)

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

    ax.stackplot(
        weights.index,
        weights.T.values,
        labels=weights.columns,
        colors=c,
    )

    ax.legend(
        loc="center left",
        bbox_to_anchor=(1, 0.5),
        frameon=False
    )

    ax.set_xlabel(xlabel or "Date")
    ax.set_ylabel(ylabel or "Weight")
    ax.set_title(title or "Weights Plot")

    if show:
        plt.show()
        return None

    plt.close(fig)
    return fig
