from typing import Iterable, Optional, Literal
import numpy as np
from .base import (
    estimate_sigma,
    weighted_sorted_dist,
    tail_quantile,
    time_scaling,
    estimate_student_df,
)

ArrayLike = Iterable[float]
Method = Literal["empirical", "parametric"]
Dist = Literal["normal", "t"]

def var(
    returns: ArrayLike,
    n_days: int = 1,
    confidence: float = 0.99,
    method: Method = "empirical",
    weights: Optional[ArrayLike] = None,
    distribution: Dist = "normal",
    df: Optional[float] = None,
    lamb: Optional[float] = None,
    mean: Literal["zero", "sample"] = "zero",
) -> float:
    """
    Compute Value at Risk (VaR) for a return series or portfolio.

    Supports empirical and parametric methods with normal or t-Student
    distributions. Returns a positive VaR.

    Notes
    -----
    The function is scale-agnostic with respect to the input returns:

    - If `returns` represent PnL values, the VaR is expressed in monetary units
    - If `returns` represent simple returns, the VaR is expressed in percentage
      terms (assuming unit notional).

    Implementation details:
    - For the unweighted empirical method, the quantile is computed with
      `method="higher"` to match the formal VaR definition based on the
      infimum of the set where the CDF exceeds `confidence` (i.e. the
      conservative choice when the quantile falls between observations).

    Parameters
    ----------
    returns : ArrayLike
        Return/PnL observations.
        For `method="empirical"` this must be a 1D time series of
        already-aggregated (already-weighted) portfolio returns/losses.
        For `method="parametric"`, 2D input (T, N) is interpreted as
        asset returns.
    n_days : int, default=1
        Forecast horizon in days. Must be positive.
    confidence : float, default=0.99
        Confidence level in (0, 1).
    method : {"empirical", "parametric"}, default="empirical"
        Estimation method. Empirical method supports only 1D inputs.
    weights : Optional[ArrayLike], default=None
        Portfolio weights. Used for 2D inputs in the parametric method.
    distribution : {"normal", "t"}, default="normal"
        Parametric distribution: "normal" or "t".
    df : Optional[float], default=None
        Degrees of freedom for Student-t. If None and
        `distribution="t"`, estimated from sample via MLE.
    lamb : Optional[float], default=None
        Exponential decay parameter for weighted empirical VaR. If provided,
        must be in (0, 1). If None, equal weights are used.

    Returns
    -------
    float
        VaR as a positive number (loss magnitude).
    """
    if distribution not in {"normal", "t"}:
        raise ValueError("distribution must be 'normal' or 't'")

    if mean not in {"zero", "sample"}:
        raise ValueError("mean must be 'zero' or 'sample'")
    
    losses = -np.asarray(returns, dtype=float)

    if not 0.0 < confidence < 1.0:
        raise ValueError("confidence must be in (0, 1).")

    if method == "parametric":
        return _parametric_var(
            losses=losses,
            gamma=confidence,
            n_days=n_days,
            weights=weights,
            distribution=distribution,
            df=df,
            mean=mean,
        )
    
    if method == "empirical":
        return _empirical_var(
            losses=losses,
            gamma=confidence,
            n_days=n_days,
            lamb=lamb,
            
        )

    raise ValueError(f"Unsupported VaR method: {method}.")


def _empirical_var(
    losses: np.ndarray,
    gamma: float,
    n_days: int,
    lamb: Optional[float],
) -> float:
    """
    Empirical (historical) VaR.
    """
    if losses.ndim != 1:
        raise ValueError("Empirical VaR requires 1D returns.")

    if lamb is None:
        # NOTE: The formal definition of VaR requires finding the infimum of
        # the set of losses where the cumulative distribution exceeds gamma.
        # In practical terms: always choose the greater value when estimating
        # the quantile.
        var_value = np.quantile(losses, gamma, method="higher")
        
    else:
        sorted_pnl, _, cum_w = weighted_sorted_dist(losses, lamb)
        idx = int(np.searchsorted(cum_w, gamma, side="right"))
        idx = min(idx, len(sorted_pnl) - 1)
        var_value = float(sorted_pnl[idx])

    return time_scaling(var_value, n_days)


def _parametric_var(
    losses: np.ndarray,
    gamma: float,
    n_days: int,
    distribution: Dist,
    df: Optional[float],
    weights: Optional[Iterable[float]] = None,
    mean: Literal["zero", "sample"] = "zero",
) -> float:
    """
    Parametric Value-at-Risk under i.i.d. assumption.

    The model assumes:

        L_t = mu + sigma * Z_t

    where Z_t follows a standardized distribution (normal or Student-t).

    If mean="zero", mu is assumed equal to 0.
    If mean="sample", mu is estimated from the sample.

    Multi-day VaR is computed as:

        VaR_N = N * mu + sqrt(N) * sigma *  z_gamma

    where z_gamma is the gamma-quantile of the standardized distribution.
    """
    sigma = estimate_sigma(
        values=losses,
        weights=weights,
    )

    portfolio_losses: Optional[np.ndarray] = None

    if mean == "sample":
        if weights is None:
            mu = np.mean(losses)
        else:
            weights_arr = np.asarray(weights, dtype=float)
            if losses.ndim == 1:
                mu = np.mean(losses)
            else:
                portfolio_losses = np.sum(losses * weights_arr, axis=1)
                mu = np.mean(portfolio_losses)
    else:
        mu = 0

    if distribution == "t":
        if df is None:
            if losses.ndim == 1:
                sample_for_df = losses
            elif weights is not None:
                if portfolio_losses is None:
                    weights_arr = np.asarray(weights, dtype=float)
                    portfolio_losses = np.sum(losses * weights_arr, axis=1)
                sample_for_df = portfolio_losses
            else:
                sample_for_df = losses.ravel()

            df = estimate_student_df(
                sample_for_df,
                fit_mean=(mean == "sample"),
            )

        if df <= 2:
            raise ValueError("df must be > 2 for finite variance.")

    z = tail_quantile(
        gamma=gamma,
        distribution=distribution,
        df=df,
    )

    if distribution == "t":
        # Standardize Student-t so that Var(Z) = 1.
        # The standard t has variance df / (df - 2),
        # so we rescale the quantile accordingly.
        z *= np.sqrt((df - 2) / df)

    var_value = mu * n_days + time_scaling(z * sigma, n_days)

    return float(var_value)
