# ===================================================================
# MODULE 5, LESSON 7: MARKET RISK – FACTOR MODELS AND STRESS TESTING
# ===================================================================

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import norm, t
from sklearn.decomposition import PCA
from sklearn.covariance import LedoitWolf
import warnings
warnings.filterwarnings('ignore')

# Set style
sns.set_style("whitegrid")
np.random.seed(42)

print("="*70)
print("MARKET RISK – FACTOR MODELS, VaR, AND STRESS TESTING")
print("="*70)

# ----------------------------------------------------------------
# PART A: GENERATE MULTI-ASSET PORTFOLIO WITH FACTOR STRUCTURE
# ----------------------------------------------------------------

# Define factors: Market (S&P 500), Interest Rates (10Y yield), FX (USD/EUR)
n_days = 500
n_assets = 20

# Factor returns (daily)
factor_means = [0.0005, -0.0001, 0.0002]  # mean returns: market, rates, FX
factor_vols = [0.015, 0.008, 0.012]       # daily vols
factor_corr = np.array([[1.0, -0.3, 0.2],
                        [-0.3, 1.0, -0.1],
                        [0.2, -0.1, 1.0]])

# Generate factor returns
factor_cov = np.diag(factor_vols) @ factor_corr @ np.diag(factor_vols)
factor_returns = np.random.multivariate_normal(factor_means, factor_cov, n_days)

# Asset betas (exposures to factors)
betas = np.random.uniform(-0.5, 1.5, (n_assets, 3))
# Add some structure: first 10 assets are equity-like (positive market beta)
betas[:10, 0] = np.random.uniform(0.8, 1.2, 10)
betas[10:15, 1] = np.random.uniform(0.5, 1.0, 5)  # rate-sensitive
betas[15:, 2] = np.random.uniform(0.6, 1.4, 5)    # FX-sensitive

# Idiosyncratic returns
idio_vol = np.random.uniform(0.005, 0.02, n_assets)
idio_returns = np.random.normal(0, idio_vol, (n_days, n_assets))

# Total asset returns: factor returns * betas + idiosyncratic
asset_returns = factor_returns @ betas.T + idio_returns

# Create DataFrame
asset_names = [f'Asset_{i+1}' for i in range(n_assets)]
returns_df = pd.DataFrame(asset_returns, columns=asset_names)
factor_df = pd.DataFrame(factor_returns, columns=['Market', 'Rates', 'FX'])

print("Portfolio Summary:")
print(f"  {n_assets} assets, {n_days} days")
print(f"  Factor correlation matrix:\n{factor_corr}")

# Portfolio weights (equal-weighted)
weights = np.ones(n_assets) / n_assets
port_returns = asset_returns @ weights

# ----------------------------------------------------------------
# PART B: FACTOR MODEL ESTIMATION
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART B: Factor Model Estimation (OLS)")
print("-"*60)

# Regress each asset return on the factors
from sklearn.linear_model import LinearRegression

betas_est = []
idios_est = []
r2_values = []

for i in range(n_assets):
    model = LinearRegression()
    model.fit(factor_returns, asset_returns[:, i])
    betas_est.append(model.coef_)
    idios_est.append(np.std(asset_returns[:, i] - model.predict(factor_returns)))
    r2_values.append(model.score(factor_returns, asset_returns[:, i]))

betas_est = np.array(betas_est)
idios_est = np.array(idios_est)

print("Estimated Betas (first 5 assets):")
betas_df = pd.DataFrame(betas_est[:5], columns=['Market', 'Rates', 'FX'])
betas_df.index = asset_names[:5]
print(betas_df.round(3))

print(f"\nAverage R²: {np.mean(r2_values):.4f}")
print(f"Average idiosyncratic volatility: {np.mean(idios_est):.4f}")

# Portfolio betas
port_betas = weights @ betas_est
print(f"\nPortfolio Betas: Market={port_betas[0]:.4f}, Rates={port_betas[1]:.4f}, FX={port_betas[2]:.4f}")

# ----------------------------------------------------------------
# PART C: PORTFOLIO RISK USING FACTOR COVARIANCE
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART C: Portfolio Risk – Factor Covariance Approach")
print("-"*60)

# Factor covariance matrix
factor_cov_est = np.cov(factor_returns.T)

# Idiosyncratic variance (diagonal)
idio_var = idios_est ** 2
idio_cov = np.diag(idio_var)

# Portfolio variance: beta_p^T * Sigma_F * beta_p + w^T * Sigma_idio * w
port_var_factor = port_betas @ factor_cov_est @ port_betas + weights @ idio_cov @ weights
port_vol_factor = np.sqrt(port_var_factor)

# Direct portfolio volatility (using historical returns)
port_vol_direct = np.std(port_returns, ddof=1)

print(f"Portfolio volatility (factor model): {port_vol_factor*100:.4f}%")
print(f"Portfolio volatility (direct): {port_vol_direct*100:.4f}%")
print(f"Difference: {(port_vol_factor - port_vol_direct)*100:.4f}%")

# VaR and ES using the factor model
confidence = 0.975  # FRTB uses 97.5% for ES
alpha = 1 - confidence
z_alpha = norm.ppf(alpha)

# 10-day VaR (scaling factor: sqrt(10))
scaling = np.sqrt(10)
var_10d_factor = -(0 + z_alpha * port_vol_factor * scaling)  # assuming zero mean
es_10d_factor = -(0 + port_vol_factor * scaling * norm.pdf(z_alpha) / alpha)

print(f"\n10-day VaR (97.5%): {var_10d_factor*100:.4f}%")
print(f"10-day ES (97.5%): {es_10d_factor*100:.4f}%")

# ----------------------------------------------------------------
# PART D: STRESS TESTING – HISTORICAL SCENARIOS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART D: Stress Testing – Historical Scenarios")
print("-"*60)

# Define historical stress scenarios (as factor shocks)
stress_scenarios = {
    '2008 Financial Crisis': {'Market': -0.15, 'Rates': 0.02, 'FX': 0.05},
    'COVID-19 (Mar 2020)': {'Market': -0.12, 'Rates': -0.03, 'FX': 0.02},
    '2022 Inflation Shock': {'Market': -0.08, 'Rates': 0.04, 'FX': -0.03},
    'Tech Bubble (2000)': {'Market': -0.10, 'Rates': 0.01, 'FX': 0.01},
    'Extreme Hypothetical': {'Market': -0.25, 'Rates': 0.06, 'FX': 0.08}
}

# Function to compute portfolio impact
def portfolio_stress(factor_shocks, betas, weights):
    """Compute portfolio return under stress scenario."""
    asset_shocks = betas @ np.array([factor_shocks['Market'], 
                                     factor_shocks['Rates'], 
                                     factor_shocks['FX']])
    port_shock = weights @ asset_shocks
    return port_shock

stress_results = []
for name, shocks in stress_scenarios.items():
    port_loss = -portfolio_stress(shocks, betas_est, weights)  # positive loss
    stress_results.append({'Scenario': name, 'Portfolio Loss': port_loss})

stress_df = pd.DataFrame(stress_results)
stress_df['Portfolio Loss %'] = stress_df['Portfolio Loss'] * 100
print("\nStress Test Results (Portfolio Loss):")
print(stress_df.to_string(index=False))

# Visualise
fig, ax = plt.subplots(figsize=(10, 6))
colors = ['red' if loss > 0.10 else 'orange' if loss > 0.05 else 'blue' 
          for loss in stress_df['Portfolio Loss']]
ax.barh(stress_df['Scenario'], stress_df['Portfolio Loss %'], color=colors)
ax.set_xlabel('Portfolio Loss (%)')
ax.set_title('Stress Testing – Historical and Hypothetical Scenarios')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('market_stress_testing.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART E: HYPOTHETICAL SCENARIO GENERATION
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART E: Hypothetical Scenario Generation")
print("-"*60)

# Generate a range of hypothetical scenarios
def generate_hypothetical_scenarios(n_scenarios=10):
    """Generate random but plausible stress scenarios."""
    scenarios = []
    for i in range(n_scenarios):
        # Random shocks with correlations
        shocks = np.random.multivariate_normal(
            [-0.05, 0.01, 0.01],
            [[0.04, -0.01, 0.005],
             [-0.01, 0.01, -0.002],
             [0.005, -0.002, 0.02]]
        )
        # Clip to plausible ranges
        shocks = np.clip(shocks, [-0.20, -0.05, -0.10], [0.05, 0.08, 0.10])
        scenarios.append({
            'Market': shocks[0],
            'Rates': shocks[1],
            'FX': shocks[2]
        })
    return scenarios

hypothetical_scenarios = generate_hypothetical_scenarios(20)

# Evaluate portfolio loss under each scenario
hypo_losses = []
for shocks in hypothetical_scenarios:
    loss = -portfolio_stress(shocks, betas_est, weights)
    hypo_losses.append(loss)

hypo_losses = np.array(hypo_losses)

print(f"Hypothetical Scenario Losses:")
print(f"  Mean: {hypo_losses.mean()*100:.2f}%")
print(f"  Std: {hypo_losses.std()*100:.2f}%")
print(f"  Max: {hypo_losses.max()*100:.2f}%")
print(f"  Min: {hypo_losses.min()*100:.2f}%")

# 95th percentile of hypothetical losses
hypo_var_95 = np.percentile(hypo_losses, 95)
print(f"  95th percentile loss: {hypo_var_95*100:.2f}%")

# Visualise hypothetical losses
fig, ax = plt.subplots(figsize=(10, 5))
ax.hist(hypo_losses * 100, bins=15, edgecolor='black', alpha=0.7, color='purple')
ax.axvline(hypo_var_95 * 100, color='red', linestyle='--', 
           label=f'95th percentile: {hypo_var_95*100:.2f}%')
ax.set_xlabel('Portfolio Loss (%)')
ax.set_ylabel('Frequency')
ax.set_title('Hypothetical Scenario Loss Distribution')
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('hypothetical_scenarios.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART F: PCA FOR RISK FACTOR REDUCTION
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART F: PCA – Risk Factor Reduction")
print("-"*60)

# Perform PCA on asset returns
pca = PCA()
pca.fit(asset_returns)

explained_var = pca.explained_variance_ratio_
cumulative_var = np.cumsum(explained_var)

print(f"First 3 principal components explain {cumulative_var[2]*100:.2f}% of variance")
print(f"Number of PCs for 95% variance: {np.argmax(cumulative_var >= 0.95) + 1}")

# Visualise
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

ax = axes[0]
ax.bar(range(1, len(explained_var)+1), explained_var, alpha=0.7, color='blue')
ax.set_xlabel('Principal Component')
ax.set_ylabel('Explained Variance Ratio')
ax.set_title('Variance Explained by Each PC')
ax.grid(True, alpha=0.3)

ax = axes[1]
ax.plot(range(1, len(cumulative_var)+1), cumulative_var, 'bo-', linewidth=2)
ax.axhline(y=0.95, color='red', linestyle='--', label='95% threshold')
ax.set_xlabel('Number of Components')
ax.set_ylabel('Cumulative Explained Variance')
ax.set_title('Cumulative Variance Explained')
ax.legend()
ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('pca_risk_factors.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART G: FRTB – STANDARDISED APPROACH (CONCEPTUAL)
# ----------------------------------------------------------------

print("\n" + "="*70)
print("PART G: FRTB – Standardised Approach (Conceptual)")
print("="*70)

print("""
FRTB Standardised Approach (SBM) – Key Components:

1. Risk Classes:
   - GIRR (General Interest Rate Risk)
   - CSR (Credit Spread Risk)
   - Equity Risk
   - FX Risk
   - Commodity Risk

2. Sensitivity-Based Method (SBM):
   - Calculate risk sensitivities (delta, vega, curvature) for each risk factor.
   - Apply risk weights (prescribed by the regulator) to each sensitivity.
   - Aggregate using correlation assumptions (within and across risk classes).

3. Default Risk Charge:
   - Additional capital for jump-to-default risk in debt instruments.

4. Simplified Implementation (Illustrative):
   - For a portfolio of equities, capital = 0.15 × portfolio value (simplified risk weight).
   - For a portfolio of bonds, capital = sensitivity × risk weight × volatility.
   
The actual calculation involves thousands of risk factors and is computationally intensive.
""")

# Simplified SA capital
sa_capital = 0.12 * 1000000  # 12% of portfolio value
print(f"Simplified SA Capital (illustrative): ${sa_capital:,.2f}")

# ----------------------------------------------------------------
# PART H: REGULATORY CAPITAL COMPARISON
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART H: Regulatory Capital Comparison")
print("-"*60)

# Portfolio value
portfolio_value = 1000000  # $1M

# IMA capital (based on ES)
es_10d = es_10d_factor * portfolio_value
ima_capital = es_10d  # Simplified: ES is the capital

# SA capital (simplified)
sa_capital_simplified = 0.12 * portfolio_value

print(f"Portfolio Value: ${portfolio_value:,.2f}")
print(f"\nFRTB Internal Models Approach (IMA) Capital: ${ima_capital:,.2f}")
print(f"FRTB Standardised Approach (SA) Capital: ${sa_capital_simplified:,.2f}")
print(f"Difference: SA is {(sa_capital_simplified/ima_capital - 1)*100:+.1f}% vs IMA")

# Backtesting requirement
print("""
Backtesting Requirements (FRTB):
  - Banks must backtest their internal models daily.
  - Test: compare 1-day VaR (99%) with actual P&L.
  - Green zone: 0-4 exceptions per year (acceptable).
  - Yellow zone: 5-9 exceptions (requires review).
  - Red zone: 10+ exceptions (model invalidated).
""")