"""Comprehensive tests for cpz.risk modules.

Tests cover: cointegration, portfolio decomposition, scenario engine, fragility.
All tests use synthetic data — no external dependencies.
"""

import math
import pytest

# ── Test data generators ─────────────────────────────────────────────

def _random_walk(n: int = 200, drift: float = 0.001, vol: float = 0.02, seed: int = 42) -> list[float]:
    """Generate a random walk price series."""
    import numpy as np
    rng = np.random.default_rng(seed)
    returns = rng.normal(drift, vol, n)
    prices = [100.0]
    for r in returns:
        prices.append(prices[-1] * (1 + r))
    return prices


def _cointegrated_pair(n: int = 200, seed: int = 42) -> tuple[list[float], list[float]]:
    """Generate a cointegrated pair of price series."""
    import numpy as np
    rng = np.random.default_rng(seed)
    # Common stochastic trend
    trend = np.cumsum(rng.normal(0.001, 0.01, n))
    noise_a = rng.normal(0, 0.5, n)
    noise_b = rng.normal(0, 0.5, n)
    series_a = (100 + trend + noise_a).tolist()
    series_b = (50 + 0.5 * trend + noise_b).tolist()
    return series_a, series_b


def _daily_returns(n: int = 100, mu: float = 0.0005, sigma: float = 0.015, seed: int = 42) -> list[float]:
    """Generate synthetic daily returns."""
    import numpy as np
    rng = np.random.default_rng(seed)
    return rng.normal(mu, sigma, n).tolist()


def _strategy_returns(n_strategies: int = 4, n_days: int = 100, seed: int = 42) -> dict[str, list[float]]:
    """Generate returns for multiple strategies."""
    import numpy as np
    rng = np.random.default_rng(seed)
    result = {}
    for i in range(n_strategies):
        mu = 0.0003 + i * 0.0001
        sigma = 0.01 + i * 0.005
        result[f"strat_{i}"] = rng.normal(mu, sigma, n_days).tolist()
    return result


# ══════════════════════════════════════════════════════════════════════
#  COINTEGRATION TESTS
# ══════════════════════════════════════════════════════════════════════

class TestCointegration:

    def test_cointegration_test_basic(self):
        from cpz.risk.cointegration import cointegration_test
        a, b = _cointegrated_pair(300)
        result = cointegration_test(a, b)
        assert result.num_observations == 300
        assert result.hedge_ratio != 0
        assert result.engle_granger.test_statistic != 0
        assert result.engle_granger.p_value >= 0
        assert result.spread_std > 0

    def test_cointegration_test_returns_johansen(self):
        from cpz.risk.cointegration import cointegration_test
        a, b = _cointegrated_pair(200)
        result = cointegration_test(a, b)
        assert result.johansen is not None
        assert len(result.johansen.trace_stat) == 2
        assert len(result.johansen.critical_values_95) == 2

    def test_cointegration_test_half_life(self):
        from cpz.risk.cointegration import cointegration_test
        a, b = _cointegrated_pair(300)
        result = cointegration_test(a, b)
        # Half-life should be positive for a mean-reverting spread
        if result.half_life is not None:
            assert result.half_life > 0

    def test_cointegration_test_too_few_observations(self):
        from cpz.risk.cointegration import cointegration_test
        with pytest.raises(ValueError, match="at least 20"):
            cointegration_test([1, 2, 3], [4, 5, 6])

    def test_cointegration_rolling_basic(self):
        from cpz.risk.cointegration import cointegration_rolling
        a, b = _cointegrated_pair(200)
        result = cointegration_rolling(a, b, window=60)
        assert result.num_points == 200
        assert result.window == 60
        assert len(result.p_values) == 200
        assert len(result.hedge_ratios) == 200
        assert len(result.z_scores) == 200
        # First 59 should be None (window - 1)
        assert result.p_values[0] is None
        assert result.p_values[58] is None
        # 60th should have a value
        assert result.p_values[59] is not None

    def test_cointegration_rolling_too_few(self):
        from cpz.risk.cointegration import cointegration_rolling
        with pytest.raises(ValueError, match="at least"):
            cointegration_rolling([1] * 30, [2] * 30, window=60)

    def test_cointegration_result_is_pydantic(self):
        from cpz.risk.cointegration import cointegration_test, CointegrationResult
        a, b = _cointegrated_pair(100)
        result = cointegration_test(a, b)
        assert isinstance(result, CointegrationResult)
        d = result.model_dump()
        assert "engle_granger" in d
        assert "hedge_ratio" in d


# ══════════════════════════════════════════════════════════════════════
#  PORTFOLIO TESTS
# ══════════════════════════════════════════════════════════════════════

class TestPortfolio:

    def _make_cov(self) -> tuple[dict[str, float], dict[str, dict[str, float]]]:
        """Create simple weights and covariance matrix."""
        weights = {"A": 0.4, "B": 0.3, "C": 0.3}
        cov = {
            "A": {"A": 0.04, "B": 0.01, "C": 0.005},
            "B": {"A": 0.01, "B": 0.09, "C": 0.02},
            "C": {"A": 0.005, "B": 0.02, "C": 0.0625},
        }
        return weights, cov

    def test_euler_decomposition_sums_to_portfolio_vol(self):
        from cpz.risk.portfolio import euler_decomposition
        weights, cov = self._make_cov()
        result = euler_decomposition(weights, cov)
        assert result.portfolio_volatility > 0
        total_rc = sum(c.risk_contribution for c in result.contributions)
        assert abs(total_rc - result.portfolio_volatility) < 1e-6

    def test_euler_decomposition_pct_sums_to_100(self):
        from cpz.risk.portfolio import euler_decomposition
        weights, cov = self._make_cov()
        result = euler_decomposition(weights, cov)
        total_pct = sum(c.risk_contribution_pct for c in result.contributions)
        assert abs(total_pct - 100.0) < 0.5  # Allow small rounding

    def test_mctr_returns_per_strategy(self):
        from cpz.risk.portfolio import mctr
        weights, cov = self._make_cov()
        result = mctr(weights, cov)
        assert result.portfolio_volatility > 0
        assert len(result.mctr) == 3
        assert "A" in result.mctr
        assert "B" in result.mctr

    def test_optimize_min_variance(self):
        from cpz.risk.portfolio import optimize_portfolio
        returns = _strategy_returns(3, 100)
        result = optimize_portfolio(returns, method="min_variance")
        assert result.method == "min_variance"
        assert len(result.weights) == 3
        total_w = sum(result.weights.values())
        assert abs(total_w - 1.0) < 0.1  # Weights should roughly sum to 1

    def test_optimize_risk_parity(self):
        from cpz.risk.portfolio import optimize_portfolio
        returns = _strategy_returns(3, 100)
        result = optimize_portfolio(returns, method="risk_parity")
        assert result.method == "risk_parity"
        assert result.volatility > 0
        assert len(result.weights) == 3

    def test_optimize_max_sharpe(self):
        from cpz.risk.portfolio import optimize_portfolio
        returns = _strategy_returns(3, 100)
        result = optimize_portfolio(returns, method="max_sharpe")
        assert result.method == "max_sharpe"
        assert len(result.weights) == 3

    def test_optimize_equal_weight(self):
        from cpz.risk.portfolio import optimize_portfolio
        returns = _strategy_returns(4, 100)
        result = optimize_portfolio(returns, method="equal_weight")
        for w in result.weights.values():
            assert abs(w - 0.25) < 0.01

    def test_what_if_analysis(self):
        from cpz.risk.portfolio import what_if_analysis
        weights, cov = self._make_cov()
        new_returns = _daily_returns(100, mu=0.001, sigma=0.02, seed=99)
        result = what_if_analysis(weights, cov, new_returns, "D", allocation_pct=10.0)
        assert result.new_volatility > 0
        assert "D" in result.new_weights
        assert abs(sum(result.new_weights.values()) - 1.0) < 0.01

    def test_risk_budget_check_within(self):
        from cpz.risk.portfolio import risk_budget_check
        weights, cov = self._make_cov()
        budgets = {"A": 60.0, "B": 60.0, "C": 60.0}  # Generous budgets
        result = risk_budget_check(weights, cov, budgets)
        assert result.within_budget is True
        assert len(result.breaches) == 0

    def test_risk_budget_check_breach(self):
        from cpz.risk.portfolio import risk_budget_check
        weights, cov = self._make_cov()
        budgets = {"A": 5.0, "B": 5.0, "C": 5.0}  # Impossible budgets
        result = risk_budget_check(weights, cov, budgets)
        assert result.within_budget is False
        assert len(result.breaches) > 0

    def test_empty_portfolio(self):
        from cpz.risk.portfolio import euler_decomposition, mctr
        result_e = euler_decomposition({}, {})
        assert result_e.portfolio_volatility == 0
        result_m = mctr({}, {})
        assert result_m.portfolio_volatility == 0


# ══════════════════════════════════════════════════════════════════════
#  SCENARIO TESTS
# ══════════════════════════════════════════════════════════════════════

class TestScenario:

    def _weights(self) -> dict[str, float]:
        return {"AAPL": 0.3, "MSFT": 0.2, "BTC/USD": 0.15, "TLT": 0.2, "GLD": 0.15}

    def test_custom_scenario(self):
        from cpz.risk.scenario import custom_scenario
        shocks = {"equities": -0.20, "crypto": -0.50, "bonds": 0.05, "commodities": -0.10}
        result = custom_scenario(shocks, self._weights())
        assert result.total_impact_pct != 0
        assert len(result.per_position) == 5
        assert "AAPL" in result.per_position
        assert "BTC/USD" in result.per_position

    def test_custom_scenario_crypto_gets_crypto_shock(self):
        from cpz.risk.scenario import custom_scenario
        shocks = {"equities": -0.10, "crypto": -0.50}
        weights = {"BTC/USD": 1.0}
        result = custom_scenario(shocks, weights)
        # BTC should use the crypto shock (-50%), not equity (-10%)
        assert result.total_impact_pct < -40  # Should be near -50%

    def test_cascade_effects_amplifies(self):
        from cpz.risk.scenario import cascade_effects
        primary = {"equities": -0.30}
        result = cascade_effects(primary)
        assert result.amplification_factor >= 1.0
        assert len(result.total_shock) > 1  # Should propagate to other classes
        assert "equities" in result.total_shock

    def test_cascade_custom_transmission(self):
        from cpz.risk.scenario import cascade_effects
        primary = {"equities": -0.20}
        transmission = {"equities": {"bonds": -0.5, "crypto": 1.0}}
        result = cascade_effects(primary, transmission)
        # Bonds should get positive effect (negative * negative coefficient)
        # Crypto should get negative effect
        assert result.secondary_effects.get("crypto", 0) < 0 or result.secondary_effects.get("bonds", 0) != 0

    def test_historical_scenario(self):
        from cpz.risk.scenario import historical_scenario
        result = historical_scenario("covid_2020", self._weights())
        assert result.total_impact_pct < 0  # COVID was bad

    def test_historical_scenario_unknown(self):
        from cpz.risk.scenario import historical_scenario
        with pytest.raises(ValueError, match="Unknown scenario"):
            historical_scenario("nonexistent", self._weights())

    def test_scenario_comparison(self):
        from cpz.risk.scenario import scenario_comparison
        scenarios = {
            "mild": {"equities": -0.05},
            "severe": {"equities": -0.40, "crypto": -0.60},
        }
        result = scenario_comparison(scenarios, self._weights())
        assert len(result.scenarios) == 2
        assert result.worst_case == "severe"
        assert result.best_case == "mild"


# ══════════════════════════════════════════════════════════════════════
#  FRAGILITY TESTS
# ══════════════════════════════════════════════════════════════════════

class TestFragility:

    def test_fragility_score_low(self):
        from cpz.risk.fragility import fragility_score
        result = fragility_score(10, 10, 10, 10, 10)
        assert result.score == pytest.approx(10.0, abs=1)
        assert result.level == "low"

    def test_fragility_score_extreme(self):
        from cpz.risk.fragility import fragility_score
        result = fragility_score(90, 90, 90, 90, 90)
        assert result.score >= 80
        assert result.level == "extreme"

    def test_fragility_score_clamped(self):
        from cpz.risk.fragility import fragility_score
        result = fragility_score(200, 200, 200, 200, 200)  # Over 100
        assert result.score <= 100

    def test_fragility_custom_weights(self):
        from cpz.risk.fragility import fragility_score
        weights = {"crowding": 1.0, "regime": 0.0, "liquidity": 0.0, "vol_surface": 0.0, "correlation": 0.0}
        result = fragility_score(80, 20, 20, 20, 20, weights=weights)
        assert result.score >= 70  # Should be driven by crowding only

    def test_correlation_clustering(self):
        from cpz.risk.fragility import correlation_clustering
        returns = _strategy_returns(5, 100)
        result = correlation_clustering(returns)
        assert 0 <= result.average_correlation <= 1
        assert 0 <= result.cluster_score <= 100
        assert result.regime in ("normal", "elevated", "crisis")

    def test_correlation_clustering_few_assets(self):
        from cpz.risk.fragility import correlation_clustering
        result = correlation_clustering({"A": [0.01] * 10})
        assert result.cluster_score == 0  # Need at least 2

    def test_regime_from_returns_expansion(self):
        from cpz.risk.fragility import regime_from_returns
        # Steady positive returns = expansion
        returns = _daily_returns(100, mu=0.001, sigma=0.008, seed=42)
        result = regime_from_returns(returns, window=60)
        assert result.regime in ("expansion", "recovery")
        assert result.confidence > 0

    def test_regime_from_returns_crisis(self):
        from cpz.risk.fragility import regime_from_returns
        import numpy as np
        # Large negative returns with high vol = crisis
        rng = np.random.default_rng(42)
        returns = rng.normal(-0.02, 0.05, 100).tolist()
        result = regime_from_returns(returns, window=60)
        assert result.regime in ("crisis", "contraction")
        assert result.volatility_regime in ("high", "extreme")

    def test_fragility_result_serializable(self):
        from cpz.risk.fragility import fragility_score
        result = fragility_score(50, 50, 50, 50, 50)
        d = result.model_dump()
        assert "score" in d
        assert "level" in d
        assert "components" in d
        assert "description" in d


# ══════════════════════════════════════════════════════════════════════
#  EXISTING ENGINE TESTS (basic smoke tests)
# ══════════════════════════════════════════════════════════════════════

class TestRiskEngine:

    def test_compute_returns_snapshot(self):
        from cpz.risk import RiskAnalytics
        engine = RiskAnalytics(risk_free_rate=0.04)
        returns = _daily_returns(100)
        spy = _daily_returns(100, seed=99)
        weights = {"AAPL": 0.5, "MSFT": 0.5}
        result = engine.compute(returns, spy, weights, total_exposure=100000)
        assert result.portfolio_var_95 > 0
        assert result.sharpe_ratio != 0
        assert result.risk_score >= 0

    def test_monte_carlo_var(self):
        from cpz.risk import RiskAnalytics
        engine = RiskAnalytics()
        returns = _daily_returns(100)
        result = engine.monte_carlo_var(returns, num_simulations=1000, seed=42)
        assert result.var_95 != 0
        assert result.expected_shortfall_95 != 0
        assert result.num_simulations == 1000

    def test_tail_risk(self):
        from cpz.risk import RiskAnalytics
        engine = RiskAnalytics()
        returns = _daily_returns(100)
        result = engine.tail_risk(returns)
        assert result.skewness != 0 or True  # Could be near zero
        assert result.cornish_fisher_var_95 != 0

    def test_sharpe_inference(self):
        from cpz.risk import RiskAnalytics
        engine = RiskAnalytics()
        returns = _daily_returns(100)
        result = engine.sharpe_inference(returns)
        assert result.num_observations == 100
        assert result.p_value >= 0

    def test_factor_exposure(self):
        from cpz.risk import RiskAnalytics
        engine = RiskAnalytics()
        returns = _daily_returns(100)
        factors = {"Market": _daily_returns(100, seed=99)}
        result = engine.factor_exposure(returns, factors)
        assert result.r_squared >= 0
        assert "Market" in result.factors

    def test_stress_test_all(self):
        from cpz.risk import RiskAnalytics
        engine = RiskAnalytics()
        weights = {"AAPL": 0.5, "BTC/USD": 0.3, "TLT": 0.2}
        result = engine.stress_test_all(weights)
        assert len(result) >= 7
        assert "COVID-19 Crisis" in result
