# ===================================================================
# MODULE 5, LESSON 3: STRESS TESTING AND SCENARIO ANALYSIS
# ===================================================================

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score, brier_score_loss
from sklearn.model_selection import train_test_split
from scipy.stats import norm
import warnings
warnings.filterwarnings('ignore')

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

print("="*70)
print("MODEL VALIDATION AND STRESS TESTING")
print("="*70)

# ----------------------------------------------------------------
# PART A: GENERATE AND VALIDATE A CREDIT MODEL
# ----------------------------------------------------------------

# Generate synthetic data
n = 5000
income = np.random.gamma(5, 15, n) + 20
dti = np.random.beta(2, 5, n) * 60
credit_score = np.random.normal(700, 50, n).clip(550, 850)
loan_amt = np.random.gamma(4, 50, n) + 50

log_odds = -4.5 + 0.04*dti - 0.005*credit_score + 0.01*(loan_amt/1000)
prob_default = 1/(1+np.exp(-log_odds))
default = np.random.binomial(1, prob_default)

X = pd.DataFrame({'income': income, 'dti': dti, 'credit_score': credit_score, 'loan_amt': loan_amt})
y = default

# Train-test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# Fit logistic regression
model = LogisticRegression(random_state=42)
model.fit(X_train, y_train)

# Predictions
y_pred_prob = model.predict_proba(X_test)[:, 1]
y_pred_class = model.predict(X_test)

# ----------------------------------------------------------------
# PART B: QUANTITATIVE VALIDATION
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART B: QUANTITATIVE VALIDATION")
print("-"*60)

# AUC
auc = roc_auc_score(y_test, y_pred_prob)
print(f"AUC: {auc:.4f}")

# KS Statistic
from scipy.stats import ks_2samp
scores_good = y_pred_prob[y_test == 0]
scores_bad = y_pred_prob[y_test == 1]
ks_stat, ks_p = ks_2samp(scores_good, scores_bad)
print(f"KS Statistic: {ks_stat:.4f} (p={ks_p:.4f})")

# Gini Coefficient
gini = 2 * auc - 1
print(f"Gini Coefficient: {gini:.4f}")

# Brier Score
brier = brier_score_loss(y_test, y_pred_prob)
print(f"Brier Score: {brier:.4f}")

# Accuracy
accuracy = (y_pred_class == y_test).mean()
print(f"Accuracy: {accuracy:.4f}")

# Hosmer-Lemeshow Test (Calibration)
def hosmer_lemeshow(y_true, y_pred, n_groups=10):
    """Hosmer-Lemeshow goodness-of-fit test."""
    df = pd.DataFrame({'y_true': y_true, 'y_pred': y_pred})
    df['decile'] = pd.qcut(df['y_pred'], q=n_groups, labels=False, duplicates='drop')
    
    observed = df.groupby('decile')['y_true'].sum().values
    expected = df.groupby('decile')['y_pred'].sum().values
    n_obs = df.groupby('decile').size().values
    
    # Calculate H-L statistic
    hl_stat = np.sum((observed - expected)**2 / (expected * (1 - expected/n_obs)))
    # Degrees of freedom: n_groups - 2
    from scipy.stats import chi2
    p_value = 1 - chi2.cdf(hl_stat, df=n_groups-2)
    return hl_stat, p_value

hl_stat, hl_p = hosmer_lemeshow(y_test, y_pred_prob)
print(f"Hosmer-Lemeshow statistic: {hl_stat:.4f}")
print(f"Hosmer-Lemeshow p-value: {hl_p:.4f}")
if hl_p > 0.05:
    print("  ✓ Model is well-calibrated (p > 0.05)")
else:
    print("  ⚠ Model may be miscalibrated (p < 0.05)")

# ----------------------------------------------------------------
# PART C: CALIBRATION PLOT
# ----------------------------------------------------------------

def calibration_plot(y_true, y_pred, n_bins=10):
    """Create calibration plot."""
    df = pd.DataFrame({'y_true': y_true, 'y_pred': y_pred})
    df['bin'] = pd.qcut(df['y_pred'], q=n_bins, labels=False)
    
    cal_data = df.groupby('bin').agg(
        observed=('y_true', 'mean'),
        predicted=('y_pred', 'mean'),
        count=('y_true', 'count')
    ).reset_index()
    
    return cal_data

cal_data = calibration_plot(y_test, y_pred_prob)

fig, ax = plt.subplots(figsize=(8, 6))
ax.plot(cal_data['predicted'], cal_data['observed'], 'bo-', linewidth=2, markersize=8, label='Model')
ax.plot([0, 1], [0, 1], 'r--', linewidth=1, label='Perfect Calibration')
ax.set_xlabel('Mean Predicted Probability')
ax.set_ylabel('Observed Default Rate')
ax.set_title('Calibration Plot', fontsize=12)
ax.legend()
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('calibration_plot.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART D: SENSITIVITY ANALYSIS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART D: SENSITIVITY ANALYSIS")
print("-"*60)

# Baseline scenario: average customer
baseline = pd.DataFrame({
    'income': [50],
    'dti': [30],
    'credit_score': [700],
    'loan_amt': [150]
})

# Define scenarios
scenarios = {
    'Baseline': {'income': 50, 'dti': 30, 'credit_score': 700, 'loan_amt': 150},
    'DTI +10%': {'income': 50, 'dti': 33, 'credit_score': 700, 'loan_amt': 150},
    'Credit Score -50': {'income': 50, 'dti': 30, 'credit_score': 650, 'loan_amt': 150},
    'Income -20%': {'income': 40, 'dti': 30, 'credit_score': 700, 'loan_amt': 150},
    'Severe Stress': {'income': 35, 'dti': 40, 'credit_score': 600, 'loan_amt': 200},
}

results = []
for name, params in scenarios.items():
    X_scenario = pd.DataFrame([params])
    prob = model.predict_proba(X_scenario)[0, 1]
    results.append({'Scenario': name, 'Default Probability': prob})

sensitivity_df = pd.DataFrame(results)
print("\nSensitivity Analysis Results:")
print(sensitivity_df.to_string(index=False))

# Visualise sensitivity
fig, ax = plt.subplots(figsize=(10, 6))
colors = ['blue' if s == 'Baseline' else 'orange' if 'Severe' in s else 'green' for s in sensitivity_df['Scenario']]
ax.barh(sensitivity_df['Scenario'], sensitivity_df['Default Probability'], color=colors)
ax.set_xlabel('Predicted Default Probability')
ax.set_title('Sensitivity Analysis – Scenario Impact on Default Probability', fontsize=12)
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('sensitivity_analysis.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART E: STRESS TESTING – MACROECONOMIC SCENARIOS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART E: STRESS TESTING – MACROECONOMIC SCENARIOS")
print("-"*60)

# Simulate stress factors: GDP growth, unemployment, interest rates
# We'll map these to model features (income, dti, credit_score)

def apply_macro_stress(X, gdp_shock, unemp_shock, rate_shock):
    """
    Apply macroeconomic shocks to the dataset.
    In practice, we would use a more sophisticated mapping (e.g., through a macro model).
    """
    X_stressed = X.copy()
    # Income decreases with unemployment
    X_stressed['income'] = X_stressed['income'] * (1 - 0.5 * abs(unemp_shock))
    # DTI increases with interest rates (higher debt service)
    X_stressed['dti'] = X_stressed['dti'] * (1 + 0.3 * abs(rate_shock))
    # Credit scores decline with unemployment (financial stress)
    X_stressed['credit_score'] = X_stressed['credit_score'] - 20 * abs(unemp_shock)
    # Loan amounts may increase (more borrowing)
    X_stressed['loan_amt'] = X_stressed['loan_amt'] * (1 + 0.2 * abs(gdp_shock))
    return X_stressed

# Define CCAR-like scenarios
scenarios_macro = {
    'Baseline': {'gdp': 0, 'unemp': 0, 'rates': 0},
    'Adverse': {'gdp': -0.02, 'unemp': 0.03, 'rates': 0.01},
    'Severely Adverse': {'gdp': -0.05, 'unemp': 0.06, 'rates': 0.025},
}

stress_results = []
for name, shocks in scenarios_macro.items():
    X_stressed = apply_macro_stress(X_test, shocks['gdp'], shocks['unemp'], shocks['rates'])
    y_pred_stressed = model.predict_proba(X_stressed)[:, 1]
    avg_default = y_pred_stressed.mean()
    stress_results.append({
        'Scenario': name,
        'Avg Default Rate': avg_default,
        'Increase from Baseline (%)': (avg_default - y_pred_prob.mean()) * 100
    })

stress_df = pd.DataFrame(stress_results)
print("\nStress Testing Results:")
print(stress_df.to_string(index=False))

# Visualise
fig, ax = plt.subplots(figsize=(10, 6))
x = np.arange(len(stress_df))
width = 0.6
bars = ax.bar(x, stress_df['Avg Default Rate'], width)
ax.set_xticks(x)
ax.set_xticklabels(stress_df['Scenario'])
ax.set_ylabel('Average Default Probability')
ax.set_title('Stress Testing – Impact on Default Rates', fontsize=12)
for bar, val in zip(bars, stress_df['Avg Default Rate']):
    ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01, f'{val:.2%}', ha='center', va='bottom')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('stress_testing_results.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART F: REVERSE STRESS TESTING
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART F: REVERSE STRESS TESTING")
print("-"*60)

# Find the conditions that would cause the default rate to exceed a threshold
target_default_rate = 0.15
print(f"Target: Find scenarios that would push default rate to {target_default_rate:.1%}")

# We'll simulate by iteratively adjusting DTI and credit score
dti_range = np.linspace(20, 70, 50)
cs_range = np.linspace(550, 800, 50)

# Create a grid of scenarios
grid_results = []
for dti_val in dti_range:
    for cs_val in cs_range:
        X_scenario = pd.DataFrame({
            'income': [50],
            'dti': [dti_val],
            'credit_score': [cs_val],
            'loan_amt': [150]
        })
        prob = model.predict_proba(X_scenario)[0, 1]
        grid_results.append({'dti': dti_val, 'credit_score': cs_val, 'prob': prob})

grid_df = pd.DataFrame(grid_results)

# Find combinations that exceed the target
threshold_df = grid_df[grid_df['prob'] >= target_default_rate]

if len(threshold_df) > 0:
    print(f"Found {len(threshold_df)} combinations that exceed {target_default_rate:.1%}")
    print("\nExample combinations that cause failure:")
    print(threshold_df.sort_values('prob', ascending=False).head(5).to_string(index=False))
else:
    print("No combinations found. The model is robust.")

# Visualise reverse stress testing
fig, ax = plt.subplots(figsize=(10, 7))
pivot = grid_df.pivot(index='dti', columns='credit_score', values='prob')
im = ax.contourf(pivot.columns, pivot.index, pivot.values, levels=20, cmap='RdYlBu_r')
ax.contour(pivot.columns, pivot.index, pivot.values, levels=[target_default_rate], colors='red', linewidths=2, linestyles='--')
ax.scatter(700, 30, color='black', s=100, marker='*', label='Baseline')
ax.set_xlabel('Credit Score')
ax.set_ylabel('DTI (%)')
ax.set_title('Reverse Stress Testing – Combinations That Exceed Target Default Rate', fontsize=12)
cbar = plt.colorbar(im, ax=ax)
cbar.set_label('Default Probability')
ax.legend()
plt.tight_layout()
plt.savefig('reverse_stress_testing.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART G: DOCUMENTATION TEMPLATE
# ----------------------------------------------------------------

print("\n" + "="*70)
print("PART G: MODEL VALIDATION DOCUMENTATION TEMPLATE")
print("="*70)

documentation = """
--- MODEL VALIDATION REPORT ---

1. MODEL IDENTIFICATION
   Model Name: Credit Default Predictor v3.2
   Model Type: Logistic Regression
   Development Date: 2026-01-15
   Validation Date: {validation_date}
   Model Owner: Credit Risk Analytics

2. DATA QUALITY
   - Data Source: Internal loan origination system
   - Sample Period: {sample_start} to {sample_end}
   - Sample Size: {sample_size} loans
   - Default Rate: {default_rate:.2%}
   - Data Completeness: {data_completeness:.1%}

3. MODEL PERFORMANCE (Validation Sample)
   - AUC: {auc:.4f}
   - KS Statistic: {ks_stat:.4f}
   - Gini Coefficient: {gini:.4f}
   - Brier Score: {brier:.4f}
   - Accuracy: {accuracy:.4f}
   - Calibration (H-L test p-value): {hl_p:.4f}

4. STRESS TESTING RESULTS
   - Adverse Scenario Default Rate: {adverse_rate:.2%}
   - Severely Adverse Default Rate: {severe_rate:.2%}
   - Capital Impact Estimate: ${capital_impact:,.2f}

5. VALIDATION CONCLUSION
   - Model is {status} for use in {business_use}.
   - Recommended next validation date: {next_validation}

6. APPROVALS
   - Model Developer: __________________
   - Validator: __________________
   - Model Risk Committee: __________________
"""

# Fill with example values
validation_date = "2026-04-15"
sample_start = "2023-01-01"
sample_end = "2025-12-31"
sample_size = len(X)
default_rate = y.mean()
data_completeness = 98.7
capital_impact = 1500000
adverse_rate = stress_df[stress_df['Scenario'] == 'Adverse']['Avg Default Rate'].iloc[0]
severe_rate = stress_df[stress_df['Scenario'] == 'Severely Adverse']['Avg Default Rate'].iloc[0]

status = "approved" if hl_p > 0.05 and auc > 0.7 else "conditional"
business_use = "credit underwriting and portfolio management"
next_validation = "2027-04-15"

print(documentation.format(
    validation_date=validation_date,
    sample_start=sample_start,
    sample_end=sample_end,
    sample_size=sample_size,
    default_rate=default_rate,
    data_completeness=data_completeness,
    auc=auc,
    ks_stat=ks_stat,
    gini=gini,
    brier=brier,
    accuracy=accuracy,
    hl_p=hl_p,
    adverse_rate=adverse_rate,
    severe_rate=severe_rate,
    capital_impact=capital_impact,
    status=status,
    business_use=business_use,
    next_validation=next_validation
))