"""Uncertainty components."""

from dataclasses import dataclass

import numpy as np
import pandas as pd

from ..data import log_returns
from ..halving import GENESIS, cycle_position
from .drift import PowerLawDrift, mean_by_cycle_day


@dataclass(frozen=True)
class TrailingVol:
    """Standard deviation of daily log returns over the trailing `window` days, scaled by √h."""

    window: int = 365

    def sd(self, history: pd.DataFrame, horizons: np.ndarray) -> np.ndarray:
        return log_returns(history).iloc[-self.window :].std() * np.sqrt(horizons)


@dataclass(frozen=True)
class EwmaVol:
    """Weighted blend of exponentially weighted standard deviations, scaled by √h."""

    spans: tuple[int, ...] = (30, 90, 180)
    weights: tuple[float, ...] = (0.2, 0.5, 0.3)

    def sd(self, history: pd.DataFrame, horizons: np.ndarray) -> np.ndarray:
        returns = log_returns(history)
        sigma = sum(
            w * returns.ewm(span=s).std().iloc[-1]
            for s, w in zip(self.spans, self.weights, strict=True)
        )
        return sigma * np.sqrt(horizons)


@dataclass(frozen=True)
class ReversionVol:
    """
    Volatility starts at its current (short EWMA) level and decays toward a
    long-run level with a `half_life_days` half-life.

    `long_run="level"` uses the trailing `long_window` standard deviation;
    `long_run="trend"` extrapolates a log-linear trend in 90-day realised
    volatility over the same window, since volatility has fallen for years.
    """

    now_span: int = 30
    half_life_days: float = 90.0
    long_run: str = "level"
    long_window: int = 1460

    def sd(self, history: pd.DataFrame, horizons: np.ndarray) -> np.ndarray:
        returns = log_returns(history)
        now = returns.ewm(span=self.now_span).std().iloc[-1]
        days = np.arange(1, horizons.max() + 1)
        trailing = returns.iloc[-self.long_window :]
        if self.long_run == "level":
            long_run = np.full(len(days), trailing.std())
        elif self.long_run == "trend":
            # Non-overlapping 90-day windows, anchored at the origin.
            realised = trailing.rolling(90).std().iloc[::-90].dropna()
            age = (realised.index - history.index[-1]).days.to_numpy()
            slope, intercept = np.polyfit(age, np.log(realised.to_numpy()), 1)
            long_run = np.exp(intercept + slope * days)
        else:
            raise ValueError(f"unknown long_run {self.long_run!r}")
        phi = 0.5 ** (1 / self.half_life_days)
        daily_var = long_run**2 + (now**2 - long_run**2) * phi**days
        return np.sqrt(np.cumsum(daily_var)[horizons - 1])


@dataclass(frozen=True)
class CycleVol:
    """
    Trailing volatility, modulated by how volatile each day of the halving
    cycle has been relative to the rest of the cycle.

    The per-day ratio is estimated like CycleDrift's drift (kernel-smoothed,
    recency-weighted), then pulled `shrink` of the way back toward 1. The
    trailing estimate is first divided by the ratio it was measured under, so
    the cycle effect isn't counted twice.
    """

    window: int = 365
    bandwidth_days: float = 30.0
    recency_half_life: float = 1.0
    shrink: float = 0.5

    def sd(self, history: pd.DataFrame, horizons: np.ndarray) -> np.ndarray:
        returns = log_returns(history)
        cycle, day = cycle_position(returns.index)
        current_cycle = cycle_position(history.index[-1:])[0][0]
        weight = 0.5 ** ((current_cycle - cycle) / self.recency_half_life)
        squared = returns.to_numpy() ** 2
        variance = mean_by_cycle_day(squared, day, weight, self.bandwidth_days, prior_days=10.0)
        overall = (weight * squared).sum() / weight.sum()
        ratio = 1 + (1 - self.shrink) * (np.sqrt(variance / overall) - 1)

        past_ratio = ratio[day[-self.window :]]
        base = returns.iloc[-self.window :].std() / np.sqrt(np.mean(past_ratio**2))
        future = history.index[-1] + pd.to_timedelta(np.arange(1, horizons.max() + 1), unit="D")
        _, future_day = cycle_position(future)
        return base * np.sqrt(np.cumsum(ratio[future_day] ** 2)[horizons - 1])


@dataclass(frozen=True)
class TrendReversionVol:
    """
    Uncertainty for a price that reverts to the power-law trend.

    Deviations from the trend follow a daily AR(1) with coefficient φ (fitted
    by PowerLawDrift), so their variance levels off: after h days it is
    σ²(1 − φ^2h) / (1 − φ²), with σ the trailing `window`-day volatility.

    With `parameter_uncertainty`, the uncertainty of the fitted trend line is
    added. The residuals are so autocorrelated that ~5000 days carry the
    information of only n(1 − φ)/(1 + φ) independent points, and the
    coefficient covariance is inflated to match.
    """

    window: int = 365
    parameter_uncertainty: bool = False

    def sd(self, history: pd.DataFrame, horizons: np.ndarray) -> np.ndarray:
        intercept, slope, phi = PowerLawDrift().fit(history)
        phi = min(phi, 0.9999)
        sigma = log_returns(history).iloc[-self.window :].std()
        variance = sigma**2 * (1 - phi ** (2 * horizons)) / (1 - phi**2)
        if self.parameter_uncertainty:
            variance = variance + self._trend_variance(history, horizons, intercept, slope, phi)
        return np.sqrt(variance)

    @staticmethod
    def _trend_variance(history, horizons, intercept, slope, phi) -> np.ndarray:
        t = (history.index - GENESIS).days.to_numpy()
        x = np.column_stack([np.ones(len(t)), np.log(t)])
        resid = np.log(history["close"].to_numpy()) - x @ [intercept, slope]
        n_eff = len(t) * (1 - phi) / (1 + phi)
        cov = resid.var() * np.linalg.inv(x.T @ x) * len(t) / n_eff
        # The forecast mean is a(1 − φ^h) + b(ln t_h − φ^h ln t_0) + φ^h ln P_0.
        decay = phi**horizons
        g = np.column_stack([1 - decay, np.log(t[-1] + horizons) - decay * np.log(t[-1])])
        return np.einsum("hi,ij,hj->h", g, cov, g)