from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import numpy as np
from scipy import special, stats
from sklearn.mixture import GaussianMixture


EPS = 1e-10
MODEL_NAMES = (
    "gaussian",
    "student_t",
    "skew_normal",
    "generalized_normal",
    "gaussian_mixture_2",
    "kde",
    "gaussian_mixture_3",
    "recent_60_student_t",
    "recent_120_student_t",
)
UNIVERSE_SIZES = {"U3": 3, "U5": 5, "U9": 9}


@dataclass
class DistributionModel:
    name: str
    pdf: Callable[[np.ndarray], np.ndarray]
    cdf: Callable[[np.ndarray], np.ndarray]


def causal_ewma_adjust(
    values: np.ndarray, decay: float = 0.94
) -> tuple[np.ndarray, np.ndarray]:
    values = np.asarray(values, dtype=float)
    adjusted = np.zeros_like(values)
    scales = np.zeros_like(values)
    initial_n = min(60, len(values))
    variance = max(float(np.var(values[:initial_n], ddof=1)), EPS)
    for index, value in enumerate(values):
        scales[index] = np.sqrt(max(variance, EPS))
        adjusted[index] = value / scales[index]
        variance = decay * variance + (1 - decay) * value**2
    return adjusted, scales


def _student_parameters(values: np.ndarray) -> tuple[float, float, float]:
    values = np.asarray(values, dtype=float)
    loc = float(np.mean(values))
    sd = max(float(np.std(values, ddof=1)), EPS)
    excess = max(float(stats.kurtosis(values, fisher=True, bias=False)), 0.0)
    df = 100.0 if excess < 0.06 else float(np.clip(6 / excess + 4, 2.2, 100))
    scale = sd * np.sqrt(max((df - 2) / df, EPS))
    return df, loc, max(float(scale), EPS)


def _student_model(name: str, values: np.ndarray) -> DistributionModel:
    df, loc, scale = _student_parameters(values)
    return DistributionModel(
        name,
        lambda z, df=df, loc=loc, scale=scale: stats.t.pdf(
            np.asarray(z), df, loc=loc, scale=scale
        ),
        lambda z, df=df, loc=loc, scale=scale: stats.t.cdf(
            np.asarray(z), df, loc=loc, scale=scale
        ),
    )


def _skew_parameters(values: np.ndarray) -> tuple[float, float, float]:
    values = np.asarray(values, dtype=float)
    mean = float(np.mean(values))
    sd = max(float(np.std(values, ddof=1)), EPS)
    target = float(np.clip(stats.skew(values, bias=False), -0.98, 0.98))
    alpha_grid = np.linspace(-20, 20, 2001)
    delta = alpha_grid / np.sqrt(1 + alpha_grid**2)
    theoretical = (
        ((4 - np.pi) / 2)
        * (delta * np.sqrt(2 / np.pi)) ** 3
        / np.maximum((1 - 2 * delta**2 / np.pi) ** 1.5, EPS)
    )
    alpha = float(alpha_grid[np.argmin(np.abs(theoretical - target))])
    delta_value = alpha / np.sqrt(1 + alpha**2)
    scale = sd / np.sqrt(max(1 - 2 * delta_value**2 / np.pi, EPS))
    loc = mean - scale * delta_value * np.sqrt(2 / np.pi)
    return alpha, loc, max(float(scale), EPS)


def _generalized_parameters(
    values: np.ndarray,
) -> tuple[float, float, float]:
    values = np.asarray(values, dtype=float)
    loc = float(np.mean(values))
    variance = max(float(np.var(values, ddof=1)), EPS)
    target = float(
        np.clip(stats.kurtosis(values, fisher=False, bias=False), 1.82, 30)
    )
    beta_grid = np.linspace(0.35, 10.0, 2500)
    theoretical = np.exp(
        special.gammaln(5 / beta_grid)
        + special.gammaln(1 / beta_grid)
        - 2 * special.gammaln(3 / beta_grid)
    )
    beta = float(beta_grid[np.argmin(np.abs(theoretical - target))])
    scale = np.sqrt(
        variance
        * np.exp(special.gammaln(1 / beta) - special.gammaln(3 / beta))
    )
    return beta, loc, max(float(scale), EPS)


def _mixture_model(
    values: np.ndarray, components: int, seed: int
) -> DistributionModel:
    gm = GaussianMixture(
        n_components=components,
        covariance_type="full",
        random_state=seed,
        reg_covar=1e-5,
        n_init=1,
        max_iter=150,
    )
    gm.fit(np.asarray(values).reshape(-1, 1))

    def pdf(z: np.ndarray) -> np.ndarray:
        array = np.atleast_1d(np.asarray(z, dtype=float))
        return np.exp(gm.score_samples(array.reshape(-1, 1)))

    def cdf(z: np.ndarray) -> np.ndarray:
        array = np.atleast_1d(np.asarray(z, dtype=float))
        result = np.zeros_like(array)
        for weight, mean, covariance in zip(
            gm.weights_, gm.means_.ravel(), gm.covariances_.ravel()
        ):
            result += weight * stats.norm.cdf(
                array, mean, np.sqrt(max(covariance, EPS))
            )
        return result

    return DistributionModel(f"gaussian_mixture_{components}", pdf, cdf)


def _kde_model(values: np.ndarray, anchor_n: int = 81) -> DistributionModel:
    values = np.asarray(values, dtype=float)
    probabilities = (np.arange(anchor_n) + 0.5) / anchor_n
    anchors = np.quantile(values, probabilities)
    spread = min(
        float(np.std(values, ddof=1)),
        float(stats.iqr(values) / 1.34),
    )
    bandwidth = max(0.9 * max(spread, EPS) * len(values) ** (-1 / 5), 0.05)

    def pdf(z: np.ndarray) -> np.ndarray:
        array = np.atleast_1d(np.asarray(z, dtype=float))
        scaled = (array[:, None] - anchors[None, :]) / bandwidth
        return np.mean(stats.norm.pdf(scaled), axis=1) / bandwidth

    def cdf(z: np.ndarray) -> np.ndarray:
        array = np.atleast_1d(np.asarray(z, dtype=float))
        scaled = (array[:, None] - anchors[None, :]) / bandwidth
        return np.mean(stats.norm.cdf(scaled), axis=1)

    return DistributionModel("kde", pdf, cdf)


def fit_models(values: np.ndarray, seed: int) -> list[DistributionModel]:
    values = np.asarray(values, dtype=float)
    mean = float(np.mean(values))
    sd = max(float(np.std(values, ddof=1)), EPS)
    gaussian = DistributionModel(
        "gaussian",
        lambda z, mean=mean, sd=sd: stats.norm.pdf(
            np.asarray(z), loc=mean, scale=sd
        ),
        lambda z, mean=mean, sd=sd: stats.norm.cdf(
            np.asarray(z), loc=mean, scale=sd
        ),
    )
    student = _student_model("student_t", values)
    alpha, skew_loc, skew_scale = _skew_parameters(values)
    skew = DistributionModel(
        "skew_normal",
        lambda z, a=alpha, loc=skew_loc, scale=skew_scale: stats.skewnorm.pdf(
            np.asarray(z), a, loc=loc, scale=scale
        ),
        lambda z, a=alpha, loc=skew_loc, scale=skew_scale: stats.skewnorm.cdf(
            np.asarray(z), a, loc=loc, scale=scale
        ),
    )
    beta, general_loc, general_scale = _generalized_parameters(values)
    generalized = DistributionModel(
        "generalized_normal",
        lambda z, b=beta, loc=general_loc, scale=general_scale: stats.gennorm.pdf(
            np.asarray(z), b, loc=loc, scale=scale
        ),
        lambda z, b=beta, loc=general_loc, scale=general_scale: stats.gennorm.cdf(
            np.asarray(z), b, loc=loc, scale=scale
        ),
    )
    return [
        gaussian,
        student,
        skew,
        generalized,
        _mixture_model(values, 2, seed),
        _kde_model(values),
        _mixture_model(values, 3, seed + 100_000),
        _student_model("recent_60_student_t", values[-60:]),
        _student_model("recent_120_student_t", values[-120:]),
    ]


def observation_log_losses(
    models: list[DistributionModel], values: np.ndarray
) -> np.ndarray:
    return np.column_stack(
        [
            -np.log(np.maximum(model.pdf(np.asarray(values)), EPS))
            for model in models
        ]
    )


def newey_west_mean_se(values: np.ndarray) -> float:
    values = np.asarray(values, dtype=float)
    n = len(values)
    if n < 2:
        return float("inf")
    centered = values - np.mean(values)
    lag = int(np.floor(4 * (n / 100) ** (2 / 9)))
    gamma0 = float(np.dot(centered, centered) / n)
    long_run = gamma0
    for ell in range(1, min(lag, n - 1) + 1):
        covariance = float(
            np.dot(centered[ell:], centered[:-ell]) / n
        )
        long_run += 2 * (1 - ell / (lag + 1)) * covariance
    return float(np.sqrt(max(long_run, EPS) / n))


def acceptable_indices(
    losses: np.ndarray,
    alpha: float = 0.05,
    negligible_margin: float = 0.02,
) -> tuple[np.ndarray, int, np.ndarray, np.ndarray]:
    losses = np.asarray(losses, dtype=float)
    means = losses.mean(axis=0)
    winner = int(np.argmin(means))
    count = losses.shape[1]
    critical = stats.norm.ppf(1 - alpha / max(count - 1, 1))
    accepted = np.zeros(count, dtype=bool)
    accepted[winner] = True
    lower_bounds = np.full(count, np.nan)
    standard_errors = np.full(count, np.nan)
    for index in range(count):
        if index == winner:
            lower_bounds[index] = 0.0
            standard_errors[index] = 0.0
            continue
        differences = losses[:, index] - losses[:, winner]
        standard_error = newey_west_mean_se(differences)
        lower_bound = float(np.mean(differences) - critical * standard_error)
        accepted[index] = lower_bound <= negligible_margin
        lower_bounds[index] = lower_bound
        standard_errors[index] = standard_error
    return np.flatnonzero(accepted), winner, lower_bounds, standard_errors


def score_weights(losses: np.ndarray) -> np.ndarray:
    means = np.asarray(losses, dtype=float).mean(axis=0)
    scaled = -len(losses) * (means - np.min(means))
    scaled = np.clip(scaled, -700, 0)
    weights = np.exp(scaled)
    return weights / max(float(weights.sum()), EPS)


def model_quantile_es(
    model: DistributionModel,
    training: np.ndarray,
    probability: float = 0.05,
    grid_n: int = 4001,
) -> tuple[float, float]:
    training = np.asarray(training, dtype=float)
    center = float(np.mean(training))
    sd = max(float(np.std(training, ddof=1)), EPS)
    bound = max(
        abs(float(np.min(training)) - center),
        abs(float(np.max(training)) - center),
        8 * sd,
    )
    grid = np.linspace(center - 1.5 * bound, center + 1.5 * bound, grid_n)
    cdf = np.maximum.accumulate(np.clip(model.cdf(grid), 0, 1))
    quantile = float(np.interp(probability, cdf, grid))
    density = np.maximum(model.pdf(grid), 0)
    mask = grid <= quantile
    mass = float(np.trapz(density[mask], grid[mask]))
    if mass <= EPS:
        expected_shortfall = quantile
    else:
        expected_shortfall = float(
            np.trapz(grid[mask] * density[mask], grid[mask]) / mass
        )
    return quantile, expected_shortfall


def ensemble_quantile_es(
    models: list[DistributionModel],
    weights: np.ndarray,
    training: np.ndarray,
    probability: float = 0.05,
    grid_n: int = 4001,
) -> tuple[float, float]:
    training = np.asarray(training, dtype=float)
    center = float(np.mean(training))
    sd = max(float(np.std(training, ddof=1)), EPS)
    bound = max(
        abs(float(np.min(training)) - center),
        abs(float(np.max(training)) - center),
        8 * sd,
    )
    grid = np.linspace(center - 1.5 * bound, center + 1.5 * bound, grid_n)
    cdf = np.zeros_like(grid)
    density = np.zeros_like(grid)
    for model, weight in zip(models, weights):
        cdf += weight * model.cdf(grid)
        density += weight * model.pdf(grid)
    cdf = np.maximum.accumulate(np.clip(cdf, 0, 1))
    quantile = float(np.interp(probability, cdf, grid))
    mask = grid <= quantile
    mass = float(np.trapz(density[mask], grid[mask]))
    expected_shortfall = (
        quantile
        if mass <= EPS
        else float(np.trapz(grid[mask] * density[mask], grid[mask]) / mass)
    )
    return quantile, expected_shortfall
