import sys
import os

# Bulletproof pathing: Force Python to look in both the current folder AND the parent folder
# This ensures it finds models.py regardless of whether this file is in a /tests subfolder or flat.
_this_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _this_dir)
sys.path.insert(0, os.path.dirname(_this_dir))

import pandas as pd
import numpy as np

from models import model_capm, garch_scale_covariance, regime_stress_covariance, get_conditional_erp
from core_types import CovarianceResult
from config import DEFAULT_CONFIG

# ─────────────────────────────────────────────
# 1. CAPM MODEL TESTS
# ─────────────────────────────────────────────
def test_model_capm():
    """Tests if the Capital Asset Pricing Model calculates Beta correctly."""
    
    dates = pd.date_range("2023-01-01", periods=5)
    
    # Give the market actual variance (1%, 2%, 3%, 4%, 5%)
    # If this is flat (e.g. 1% every day), variance is 0 and Beta defaults to 1.0
    spy_rets = pd.Series([0.01, 0.02, 0.03, 0.04, 0.05], index=dates)
    
    # Stock A moves exactly twice as fast as the market (Beta = 2.0)
    # Stock B moves exactly half as fast as the market (Beta = 0.5)
    returns_df = pd.DataFrame({
        'STOCK_A': spy_rets * 2.0,
        'STOCK_B': spy_rets * 0.5
    }, index=dates)
    
    rfr = 0.04 # 4% Risk Free Rate
    
    # Act: Run your CAPM function (explicitly setting periods to 252 for daily)
    forecast = model_capm(returns_df, spy_rets, rfr, periods=252, silent=True)
    exp_rets, betas = forecast.expected_returns, forecast.betas
    
    # Assert: Check the Betas using np.isclose to avoid floating-point drift
    assert np.isclose(betas['STOCK_A'], 2.0)
    assert np.isclose(betas['STOCK_B'], 0.5)
    
    # Assert: Check    # CAPM Expected Return = RFR + Beta * ERP
    expected_a = rfr + (2.0 * get_conditional_erp(rfr))
    expected_b = rfr + (0.5 * get_conditional_erp(rfr))
    
    assert np.isclose(exp_rets['STOCK_A'], expected_a)
    assert np.isclose(exp_rets['STOCK_B'], expected_b)


def test_capm_monthly_vs_daily():
    """Monthly betas should be approximately equal to daily betas for the same asset."""
    rng = np.random.default_rng(42)
    dates = pd.date_range("2020-01-01", periods=756, freq="B")  # ~3 years
    spy_daily = pd.Series(rng.normal(0.0005, 0.01, 756), index=dates)
    stock_daily = spy_daily * 1.5 + rng.normal(0, 0.005, 756)
    
    # CAPM on daily
    daily_df = pd.DataFrame({'STOCK': stock_daily}, index=dates)
    daily_forecast = model_capm(daily_df, spy_daily, 0.04, periods=252, silent=True)
    daily_betas = daily_forecast.betas
    
    # CAPM on monthly
    monthly_spy = spy_daily.resample('ME').apply(lambda x: (1+x).prod()-1)
    monthly_stock = stock_daily.resample('ME').apply(lambda x: (1+x).prod()-1)
    monthly_df = pd.DataFrame({'STOCK': monthly_stock}, index=monthly_spy.index)
    monthly_forecast = model_capm(monthly_df, monthly_spy, 0.04, periods=12, silent=True)
    monthly_betas = monthly_forecast.betas
    
    # Betas should be within 0.20 of each other (monthly has less data = more noise)
    assert abs(daily_betas['STOCK'] - monthly_betas['STOCK']) < 0.20


# ─────────────────────────────────────────────
# 2. RISK MODEL COVARIANCE & STRESS TESTS
# ─────────────────────────────────────────────
def test_garch_stationarity():
    """Ensures alpha + beta < 1 for all assets after fitting the GARCH(1,1) model."""
    rng = np.random.default_rng(42)
    
    # Create a stationary return series with a volatility spike to give GARCH something to fit
    rets = rng.normal(0, 0.01, size=1000)
    rets[500:550] *= 5 
    
    returns_df = pd.DataFrame({'ASSET_X': rets})
    base_cov_df = returns_df.cov() * 252
    vols = np.sqrt(np.diag(base_cov_df))
    base_corr_df = pd.DataFrame(base_cov_df.values / np.outer(vols, vols), index=base_cov_df.index, columns=base_cov_df.columns)
    base_cov = CovarianceResult(covariance=base_cov_df, correlation=base_corr_df, volatility=pd.Series(vols, index=base_cov_df.index))
    
    scaled_cov_res, garch_info = garch_scale_covariance(returns_df, base_cov, silent=True)
    
    assert 'ASSET_X' in garch_info
    alpha = garch_info['ASSET_X']['alpha']
    beta = garch_info['ASSET_X']['beta']
    
    # Strict stationarity requirement for mean-reverting volatility (or exactly 1.0 for EWMA fallback)
    assert alpha + beta <= 1.0


def test_garch_dynamic_correlation_structure():
    """Verifies that Factor-GARCH scaling generates valid correlation matrices that dynamically adapt."""
    rng = np.random.default_rng(42)
    
    # Create highly correlated synthetic returns
    dates = pd.date_range("2020-01-01", periods=500, freq="B")
    spy = rng.normal(0, 0.01, size=500)
    aapl = spy * 1.2 + rng.normal(0, 0.005, size=500)
    tlt = spy * -0.3 + rng.normal(0, 0.005, size=500)
    
    returns_df = pd.DataFrame({'SPY': spy, 'AAPL': aapl, 'TLT': tlt}, index=dates)
    base_cov_df = returns_df.cov() * 252
    
    # Extract base correlation matrix
    vols = np.sqrt(np.diag(base_cov_df))
    outer_vols = np.outer(vols, vols)
    base_corr_df = pd.DataFrame(base_cov_df.values / outer_vols, index=base_cov_df.index, columns=base_cov_df.columns)
    base_cov = CovarianceResult(covariance=base_cov_df, correlation=base_corr_df, volatility=pd.Series(vols, index=base_cov_df.index))
    
    # Run GARCH scaler
    scaled_cov_res, _ = garch_scale_covariance(returns_df, base_cov, silent=True)
    
    # Extract scaled correlation matrix
    scaled_cov = scaled_cov_res.covariance
    
    # Extract scaled correlation matrix
    scaled_vols = np.sqrt(np.diag(scaled_cov))
    scaled_outer = np.outer(scaled_vols, scaled_vols)
    scaled_corr = scaled_cov.values / scaled_outer
    
    # Assert correlations are valid [-1, 1]
    assert np.all(scaled_corr >= -1.0 - 1e-5)
    assert np.all(scaled_corr <= 1.0 + 1e-5)
    
    # Assert Factor-GARCH allows correlation structure to diverge from uniform CCC
    assert not np.allclose(base_corr_df, scaled_corr, atol=1e-8)


def test_regime_stress_cov_is_psd():
    """Verifies that the regime stressor maintains Positive Semi-Definiteness even under extreme skew."""
    rng = np.random.default_rng(42)
    
    # Create a random valid covariance matrix
    A = rng.random((5, 5))
    base_cov_values = np.dot(A, A.transpose()) 
    tickers = ["A", "B", "C", "D", "E"]
    base_cov_df = pd.DataFrame(base_cov_values, index=tickers, columns=tickers)
    
    vols = np.sqrt(np.diag(base_cov_df))
    base_corr_df = pd.DataFrame(base_cov_df.values / np.outer(vols, vols), index=tickers, columns=tickers)
    base_cov = CovarianceResult(covariance=base_cov_df, correlation=base_corr_df, volatility=pd.Series(vols, index=tickers))
    
    # Apply extreme regime stress (3.0 = severe crash multiplier)
    stressed_cov_res = regime_stress_covariance(base_cov, regime_severity=3.0)
    stressed_cov = stressed_cov_res.covariance
    
    # Check eigenvalues of the resulting matrix
    eigvals = np.linalg.eigvalsh(stressed_cov.values)
    
    # All eigenvalues must be >= 0 (allowing for tiny IEEE floating point errors)
    assert np.all(eigvals >= -1e-8)