import numpy as np
import pandas as pd
import pytest

from btcmodel.experiments import EXPERIMENTS
from btcmodel.halving import GENESIS
from btcmodel.models.drift import CycleDrift, PowerLawDrift, ShrunkDrift
from btcmodel.models.shape import StudentT
from btcmodel.models.volatility import ReversionVol, TrailingVol, TrendReversionVol

from .test_models import synthetic_prices

ALL_VARIANTS = {m.name: m for e in EXPERIMENTS.values() for m in e.models}


@pytest.mark.parametrize("name", list(ALL_VARIANTS))
def test_experiment_models_produce_valid_forecasts(name):
    prices = synthetic_prices(lambda day: 0.001 + 0 * day, noise=0.03)
    horizons = np.array([1, 30, 365, 1460])
    f = ALL_VARIANTS[name].forecast(prices, horizons)
    assert np.isfinite(f.log_quantiles).all()
    assert np.all(np.diff(f.log_quantiles, axis=1) >= 0)
    lo, hi = f.interval(0.8)
    assert np.all(np.diff(hi - lo) > 0)


def test_shrunk_drift_scales_linearly():
    prices = synthetic_prices(lambda day: np.where(day < 700, 0.002, -0.001))
    horizons = np.array([100, 1000])
    full = CycleDrift().expected_log_return(prices, horizons)
    np.testing.assert_allclose(
        ShrunkDrift(CycleDrift(), 0.5).expected_log_return(prices, horizons), full / 2
    )
    np.testing.assert_allclose(
        ShrunkDrift(CycleDrift(), 0.0).expected_log_return(prices, horizons), 0
    )


def test_power_law_recovers_its_exponent():
    dates = pd.date_range("2011-01-01", "2024-11-26", freq="D", name="date")
    t = (dates - GENESIS).days.to_numpy()
    prices = pd.DataFrame({"close": 1e-17 * t**5.8}, index=dates)
    _, slope, _ = PowerLawDrift().fit(prices)
    assert slope == pytest.approx(5.8, rel=1e-6)
    expected = 5.8 * np.log((t[-1] + 365) / t[-1])
    assert PowerLawDrift().expected_log_return(prices, np.array([365]))[0] == pytest.approx(
        expected
    )


def test_reversion_vol_matches_trailing_when_already_at_long_run():
    rng = np.random.default_rng(0)
    dates = pd.date_range("2011-01-01", "2024-11-26", freq="D", name="date")
    prices = pd.DataFrame(
        {"close": 100 * np.exp(np.cumsum(0.03 * rng.standard_normal(len(dates))))}, index=dates
    )
    horizons = np.array([30, 365, 1460])
    reverting = ReversionVol(now_span=1460, long_window=len(dates)).sd(prices, horizons)
    flat = TrailingVol(window=len(dates)).sd(prices, horizons)
    np.testing.assert_allclose(reverting, flat, rtol=0.05)


def test_student_t_shape_has_unit_variance_and_fatter_tails():
    q = StudentT(4).standard_quantiles(None, np.array([1]))[0]
    assert q[-1] > 2.576  # beyond the normal 99.5% quantile
    assert np.interp(0.8413, np.linspace(0.005, 0.995, 100), q) < 1.0  # thinner shoulders


def test_trend_reversion_vol_levels_off():
    # A power-law trend plus AR(1) deviations with a ~70-day half-life.
    rng = np.random.default_rng(0)
    dates = pd.date_range("2011-01-01", "2024-11-26", freq="D", name="date")
    t = (dates - GENESIS).days.to_numpy()
    gap = np.zeros(len(t))
    for i in range(1, len(t)):
        gap[i] = 0.99 * gap[i - 1] + 0.03 * rng.standard_normal()
    prices = pd.DataFrame({"close": np.exp(-40 + 5.8 * np.log(t) + gap)}, index=dates)

    horizons = np.array([30, 365, 1460, 5000])
    plain = TrendReversionVol().sd(prices, horizons)
    with_params = TrendReversionVol(parameter_uncertainty=True).sd(prices, horizons)
    assert np.all(np.diff(plain) >= 0)
    stationary = 0.03 / np.sqrt(1 - 0.99**2)
    assert plain[-1] == pytest.approx(stationary, rel=0.15)
    assert np.all(with_params >= plain)