"""Unit tests for macro_gen.py."""
import datetime as dt
import math

import numpy as np
import pytest

from libs.backtest.scenarios.macro_gen import (
    HYSpreadConfig,
    VIXConfig,
    _compute_rolling_indicators,
    _generate_hy_series,
    _generate_vix_series,
    generate_macro_data,
)
from libs.backtest.scenarios.price_gen import PriceRegime, generate_market_etf_paths

_DATES_RAW = [dt.date(2024, 1, 2) + dt.timedelta(days=i) for i in range(400)]
_TRADING_DATES = [d for d in _DATES_RAW if d.weekday() < 5][:252]
_RNG = np.random.default_rng(99)


def _make_simple_bars(dates: list[dt.date], start_price: float = 100.0) -> dict:
    """Create simple ascending price bars for testing."""
    bars = {}
    price = start_price
    for d in dates:
        price *= 1.0003  # tiny daily drift
        bars[d] = {"open": price * 0.999, "high": price * 1.005, "low": price * 0.995, "close": price, "volume": 1_000_000}
    return bars


@pytest.mark.unit
class TestComputeRollingIndicators:
    def test_sma_values(self):
        dates = _TRADING_DATES[:50]
        bars = _make_simple_bars(dates, 100.0)
        result = {}
        _compute_rolling_indicators("spy", bars, result)

        # SMA-20 at day 20 should be the average of first 20 prices
        d20 = dates[19]
        closes = [bars[dates[i]]["close"] for i in range(20)]
        expected_sma20 = sum(closes) / 20
        assert result[d20]["spy_sma_20"] == pytest.approx(expected_sma20, rel=1e-6)

    def test_sma_none_before_warmup(self):
        dates = _TRADING_DATES[:50]
        bars = _make_simple_bars(dates)
        result = {}
        _compute_rolling_indicators("spy", bars, result)
        # SMA-20 should be None for first 19 days
        for d in dates[:19]:
            assert result[d]["spy_sma_20"] is None

    def test_momentum_calculation(self):
        dates = _TRADING_DATES[:30]
        bars = _make_simple_bars(dates, 100.0)
        result = {}
        _compute_rolling_indicators("spy", bars, result)

        d10 = dates[10]
        c10 = bars[d10]["close"]
        c0 = bars[dates[0]]["close"]
        expected_mom_10 = (c10 - c0) / c0
        assert result[d10]["spy_mom_10"] == pytest.approx(expected_mom_10, rel=1e-5)

    def test_all_required_keys_present(self):
        dates = _TRADING_DATES[:100]
        bars = _make_simple_bars(dates)
        result = {}
        _compute_rolling_indicators("spy", bars, result)

        last_date = dates[-1]
        required_keys = [
            "spy_close", "spy_open", "spy_high", "spy_low", "spy_volume",
            "spy_sma_20", "spy_sma_50",
            "spy_high_20", "spy_high_50",
            "spy_mom_5", "spy_mom_10", "spy_mom_20",
            "spy_vol_20",
            "spy_efficiency_10", "spy_efficiency_20",
            "spy_downside_vol_10",
            "spy_entropy_10", "spy_entropy_20",
            "spy_ulcer_10", "spy_ulcer_20",
            "spy_drawdown_10", "spy_drawdown_20",
            "spy_autocorr_20",
        ]
        for key in required_keys:
            assert key in result[last_date], f"Missing key: {key}"


@pytest.mark.unit
class TestGenerateVIXSeries:
    def test_length_matches_dates(self):
        rng = np.random.default_rng(42)
        mkt_rets = [0.0] * len(_TRADING_DATES)
        vix_cfg = VIXConfig(base_level=18.0)
        vix = _generate_vix_series(mkt_rets, _TRADING_DATES, vix_cfg, rng)
        assert len(vix) == len(_TRADING_DATES)
        assert set(vix.keys()) == set(_TRADING_DATES)

    def test_vix_within_bounds(self):
        rng = np.random.default_rng(42)
        mkt_rets = [rng.normal(0, 0.01) for _ in _TRADING_DATES]
        vix_cfg = VIXConfig(base_level=18.0)
        vix = _generate_vix_series(mkt_rets, _TRADING_DATES, vix_cfg, rng)
        for v in vix.values():
            assert 5.0 <= v <= 90.0, f"VIX out of bounds: {v}"

    def test_higher_base_level_produces_higher_average(self):
        low_cfg = VIXConfig(base_level=12.0, mean_reversion=10.0)
        high_cfg = VIXConfig(base_level=35.0, mean_reversion=10.0)
        mkt_rets = [0.0] * len(_TRADING_DATES)
        vix_low = _generate_vix_series(mkt_rets, _TRADING_DATES, low_cfg, np.random.default_rng(1))
        vix_high = _generate_vix_series(mkt_rets, _TRADING_DATES, high_cfg, np.random.default_rng(1))
        assert sum(vix_low.values()) / len(vix_low) < sum(vix_high.values()) / len(vix_high)


@pytest.mark.unit
class TestGenerateHYSeries:
    def test_length_and_bounds(self):
        rng = np.random.default_rng(42)
        vix_cfg = VIXConfig(base_level=18.0)
        mkt_rets = [0.0] * len(_TRADING_DATES)
        vix = _generate_vix_series(mkt_rets, _TRADING_DATES, vix_cfg, rng)
        hy_cfg = HYSpreadConfig(base_level=4.0)
        hy = _generate_hy_series(vix, _TRADING_DATES, hy_cfg, rng)
        assert len(hy) == len(_TRADING_DATES)
        for v in hy.values():
            assert 1.5 <= v <= 30.0, f"HY spread out of bounds: {v}"


@pytest.mark.unit
class TestGenerateMacroData:
    def setup_method(self):
        rng = np.random.default_rng(77)
        regimes = [PriceRegime(0.10, 0.15, 252)]
        self.spy_bars, self.qqq_bars, self.mkt_log_rets = generate_market_etf_paths(
            regimes, _TRADING_DATES, rng
        )
        self.rng = np.random.default_rng(77)

    def test_all_trading_dates_present(self):
        macro = generate_macro_data(
            self.spy_bars, self.qqq_bars,
            VIXConfig(), HYSpreadConfig(),
            _TRADING_DATES, self.rng, self.mkt_log_rets,
        )
        for d in _TRADING_DATES:
            assert d in macro, f"Missing date {d} in macro"

    def test_vixcls_and_macro_vix_present(self):
        macro = generate_macro_data(
            self.spy_bars, self.qqq_bars,
            VIXConfig(), HYSpreadConfig(),
            _TRADING_DATES, self.rng, self.mkt_log_rets,
        )
        last_date = _TRADING_DATES[-1]
        assert "VIXCLS" in macro[last_date]
        assert "macro_vix" in macro[last_date]
        assert "macro_hy_spread" in macro[last_date]

    def test_spy_qqq_sma_present(self):
        macro = generate_macro_data(
            self.spy_bars, self.qqq_bars,
            VIXConfig(), HYSpreadConfig(),
            _TRADING_DATES, self.rng,
        )
        last_date = _TRADING_DATES[-1]
        for prefix in ("spy", "qqq"):
            assert f"{prefix}_close" in macro[last_date]
            assert f"{prefix}_sma_20" in macro[last_date]
            assert f"{prefix}_sma_50" in macro[last_date]
            assert f"{prefix}_mom_5" in macro[last_date]

    def test_parking_etf_proxies_present(self):
        macro = generate_macro_data(
            self.spy_bars, self.qqq_bars,
            VIXConfig(), HYSpreadConfig(),
            _TRADING_DATES, self.rng,
        )
        last_date = _TRADING_DATES[-1]
        assert "tqqq_close" in macro[last_date]
        assert "qqqm_close" in macro[last_date]
        assert "sgov_close" in macro[last_date]

    def test_bear_market_higher_vix(self):
        """Bear market should produce higher average VIX than bull."""
        rng_bull = np.random.default_rng(42)
        rng_bear = np.random.default_rng(42)
        regimes_bull = [PriceRegime(0.20, 0.14, 252)]
        regimes_bear = [PriceRegime(-0.20, 0.30, 252)]

        spy_bull, qqq_bull, rets_bull = generate_market_etf_paths(regimes_bull, _TRADING_DATES, rng_bull)
        spy_bear, qqq_bear, rets_bear = generate_market_etf_paths(regimes_bear, _TRADING_DATES, rng_bear)

        macro_bull = generate_macro_data(
            spy_bull, qqq_bull,
            VIXConfig(base_level=14.0, market_corr=-0.8), HYSpreadConfig(),
            _TRADING_DATES, np.random.default_rng(1), rets_bull,
        )
        macro_bear = generate_macro_data(
            spy_bear, qqq_bear,
            VIXConfig(base_level=28.0, market_corr=-0.8), HYSpreadConfig(),
            _TRADING_DATES, np.random.default_rng(1), rets_bear,
        )

        avg_vix_bull = sum(macro_bull[d].get("VIXCLS", 18) for d in _TRADING_DATES) / len(_TRADING_DATES)
        avg_vix_bear = sum(macro_bear[d].get("VIXCLS", 18) for d in _TRADING_DATES) / len(_TRADING_DATES)
        assert avg_vix_bear > avg_vix_bull
