# ===================================================================
# MODULE 6, LESSON 3: EXPLAINABLE AI (XAI) FOR REGULATORY COMPLIANCE
# ===================================================================

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_auc_score
import shap
import lime
import lime.lime_tabular
from sklearn.inspection import partial_dependence
import warnings
warnings.filterwarnings('ignore')

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

print("="*70)
print("EXPLAINABLE AI (XAI) FOR FINANCIAL MODELS")
print("="*70)

# ----------------------------------------------------------------
# PART A: GENERATE CREDIT DATA
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART A: Credit Data Preparation")
print("-"*60)

# Generate synthetic credit data
n_samples = 5000
income = np.random.gamma(5, 15, n_samples) + 20
dti = np.random.beta(2, 5, n_samples) * 60
credit_score = np.random.normal(700, 50, n_samples).clip(550, 850)
loan_amount = np.random.gamma(4, 50, n_samples) + 50
employment_years = np.random.gamma(3, 5, n_samples).clip(0, 30)
age = np.random.normal(45, 12, n_samples).clip(22, 75)

# Create non-linear relationships
log_odds = (-4.5 + 
            0.04 * dti - 
            0.005 * credit_score + 
            0.01 * (loan_amount/1000) + 
            0.02 * employment_years -
            0.01 * age +
            # Non-linear term: DTI squared (higher DTI has increasing effect)
            0.0003 * dti**2 +
            # Interaction: credit_score * dti
            -0.00001 * credit_score * dti)

prob_default = 1 / (1 + np.exp(-log_odds))
default = np.random.binomial(1, prob_default)

# Create DataFrame
df = pd.DataFrame({
    'income': income,
    'dti': dti,
    'credit_score': credit_score,
    'loan_amount': loan_amount,
    'employment_years': employment_years,
    'age': age,
    'default': default
})

feature_names = ['income', 'dti', 'credit_score', 'loan_amount', 'employment_years', 'age']

# Train-test split
X = df[feature_names]
y = df['default']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# Standardise (for LIME and linear models)
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

print(f"Training samples: {len(X_train)}")
print(f"Test samples: {len(X_test)}")
print(f"Default rate: {y_train.mean():.2%}")

# ----------------------------------------------------------------
# PART B: TRAIN BLACK-BOX MODEL (RANDOM FOREST)
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART B: Training Random Forest Model")
print("-"*60)

rf_model = RandomForestClassifier(
    n_estimators=100,
    max_depth=8,
    min_samples_split=50,
    min_samples_leaf=20,
    random_state=42,
    n_jobs=-1
)
rf_model.fit(X_train, y_train)

# Evaluate
train_auc = roc_auc_score(y_train, rf_model.predict_proba(X_train)[:, 1])
test_auc = roc_auc_score(y_test, rf_model.predict_proba(X_test)[:, 1])

print(f"Training AUC: {train_auc:.4f}")
print(f"Test AUC: {test_auc:.4f}")

# ----------------------------------------------------------------
# PART C: GLOBAL INTERPRETABILITY – FEATURE IMPORTANCE
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART C: Global Interpretability – Feature Importance")
print("-"*60)

# Built-in feature importance (Gini importance)
feature_importance = pd.DataFrame({
    'feature': feature_names,
    'importance': rf_model.feature_importances_
}).sort_values('importance', ascending=False)

print("Feature Importance (Gini):")
print(feature_importance.to_string(index=False))

# Visualise
fig, ax = plt.subplots(figsize=(10, 6))
ax.barh(feature_importance['feature'], feature_importance['importance'], color='blue', alpha=0.7)
ax.set_xlabel('Importance')
ax.set_title('Global Feature Importance – Random Forest')
ax.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('global_feature_importance.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART D: SHAP – GLOBAL AND LOCAL EXPLANATIONS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART D: SHAP Explanations")
print("-"*60)

# Create SHAP explainer (TreeExplainer for Random Forest)
explainer_shap = shap.TreeExplainer(rf_model)
shap_values = explainer_shap.shap_values(X_test)

# For binary classification, use the positive class SHAP values
shap_values_pos = shap_values[1] if isinstance(shap_values, list) else shap_values

# Global SHAP summary
fig, ax = plt.subplots(figsize=(12, 7))
shap.summary_plot(shap_values_pos, X_test, feature_names=feature_names, show=False)
plt.title('SHAP Summary Plot – Feature Impact on Default Probability', fontsize=14)
plt.tight_layout()
plt.savefig('shap_summary_plot.png', dpi=300)
plt.show()

# SHAP bar plot (mean absolute SHAP values)
fig, ax = plt.subplots(figsize=(10, 6))
shap.summary_plot(shap_values_pos, X_test, feature_names=feature_names, plot_type='bar', show=False)
plt.title('SHAP Feature Importance (Mean |SHAP Value|)', fontsize=14)
plt.tight_layout()
plt.savefig('shap_bar_plot.png', dpi=300)
plt.show()

# SHAP dependence plot for the most important feature (credit_score)
# Find the most important feature
mean_shap = np.abs(shap_values_pos).mean(axis=0)
most_important_idx = np.argmax(mean_shap)
most_important_feature = feature_names[most_important_idx]

fig, ax = plt.subplots(figsize=(10, 6))
shap.dependence_plot(most_important_idx, shap_values_pos, X_test, 
                     feature_names=feature_names, show=False)
plt.title(f'SHAP Dependence Plot – {most_important_feature}', fontsize=14)
plt.tight_layout()
plt.savefig('shap_dependence_plot.png', dpi=300)
plt.show()

print(f"Most important feature: {most_important_feature}")

# Local SHAP explanation (force plot for a specific prediction)
# Choose a specific test sample
sample_idx = 5
sample = X_test.iloc[sample_idx:sample_idx+1]
sample_shap = explainer_shap.shap_values(sample)

fig, ax = plt.subplots(figsize=(14, 3))
shap.force_plot(explainer_shap.expected_value[1] if isinstance(explainer_shap.expected_value, list) else explainer_shap.expected_value, 
                sample_shap[1] if isinstance(sample_shap, list) else sample_shap, 
                sample, feature_names=feature_names, matplotlib=True, show=False)
plt.title('SHAP Force Plot – Individual Prediction Explanation', fontsize=12)
plt.tight_layout()
plt.savefig('shap_force_plot.png', dpi=300)
plt.show()

print(f"Prediction for sample {sample_idx}: {rf_model.predict_proba(sample)[0, 1]:.3f}")

# ----------------------------------------------------------------
# PART E: LIME – LOCAL EXPLANATIONS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART E: LIME Local Explanations")
print("-"*60)

# Create LIME explainer (use scaled data for LIME)
explainer_lime = lime.lime_tabular.LimeTabularExplainer(
    X_train_scaled,
    feature_names=feature_names,
    class_names=['No Default', 'Default'],
    mode='classification',
    discretize_continuous=True,
    discretize_continuous_method='quantile',
    n_jobs=1
)

# Explain the same sample as above
sample_scaled = X_test_scaled[sample_idx:sample_idx+1].flatten()
exp = explainer_lime.explain_instance(
    sample_scaled,
    lambda x: rf_model.predict_proba(scaler.inverse_transform(x)),
    num_features=len(feature_names),
    top_labels=1
)

# Visualise LIME explanation
fig = exp.as_pyplot_figure()
plt.title('LIME Explanation – Individual Prediction', fontsize=12)
plt.tight_layout()
plt.savefig('lime_explanation.png', dpi=300)
plt.show()

print("LIME Explanation (top features):")
for feature, weight in exp.local_exp[1]:  # class 1 (default)
    print(f"  {feature_names[feature]}: {weight:.4f}")

# ----------------------------------------------------------------
# PART F: PARTIAL DEPENDENCE PLOTS (PDP)
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART F: Partial Dependence Plots (PDP)")
print("-"*60)

# Create PDP for credit_score
fig, ax = plt.subplots(figsize=(12, 8))

# PDP for credit_score
from sklearn.inspection import PartialDependenceDisplay
PartialDependenceDisplay.from_estimator(
    rf_model,
    X_test,
    features=['credit_score', 'dti'],
    kind='average',
    grid_resolution=20,
    ax=ax,
    feature_names=feature_names
)
plt.suptitle('Partial Dependence Plots – Credit Score and DTI', fontsize=14)
plt.tight_layout()
plt.savefig('pdp_plots.png', dpi=300)
plt.show()

# PDP with ICE (Individual Conditional Expectation)
fig, ax = plt.subplots(figsize=(12, 6))
PartialDependenceDisplay.from_estimator(
    rf_model,
    X_test,
    features=['credit_score'],
    kind='both',  # Both PDP and ICE
    grid_resolution=20,
    ax=ax,
    feature_names=feature_names,
    ice_lines_kw={'alpha': 0.1, 'color': 'gray'},
    pd_line_kw={'color': 'blue', 'linewidth': 3}
)
plt.title('PDP and ICE for Credit Score', fontsize=14)
plt.tight_layout()
plt.savefig('pdp_with_ice.png', dpi=300)
plt.show()

# ----------------------------------------------------------------
# PART G: COUNTERFACTUAL EXPLANATIONS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART G: Counterfactual Explanations")
print("-"*60)

print("""
Counterfactual explanations answer: "What would need to change for the prediction to be different?"

For example, if a loan application is rejected, a counterfactual explanation might say:
"If your credit score were 720 instead of 680, your application would have been approved."

This is powerful for:
  - Providing actionable feedback to rejected applicants.
  - Demonstrating compliance with fair lending regulations.
  - Building trust with customers.

Simple counterfactual search algorithm:
  1. Start with the original instance.
  2. Iteratively adjust features (within plausible bounds).
  3. Check if the prediction changes.
  4. Return the minimal set of changes.

Algorithm for credit_score counterfactual:
""")

def find_counterfactual(model, instance, feature_names, target_class=0, max_iter=100):
    """
    Simple counterfactual search.
    Adjust credit_score until prediction changes from default to no-default.
    """
    # Copy the instance
    cf = instance.copy()
    original_prob = model.predict_proba(cf.values.reshape(1, -1))[0, 1]
    print(f"Original instance: Default probability = {original_prob:.3f}")
    
    # Try to find a counterfactual by increasing credit_score
    for i in range(max_iter):
        cf['credit_score'] = instance['credit_score'] + i * 2  # Increase by 2 each step
        prob = model.predict_proba(cf.values.reshape(1, -1))[0, 1]
        if prob < 0.5:  # No-default
            print(f"Counterfactual found! credit_score: {instance['credit_score']:.0f} → {cf['credit_score']:.0f}")
            print(f"New default probability: {prob:.3f}")
            return cf, i
    
    print("Counterfactual not found within max iterations.")
    return None, None

# Use the same sample instance
sample_df = pd.DataFrame([X_test.iloc[sample_idx].values], columns=feature_names)
cf, iterations = find_counterfactual(rf_model, sample_df, feature_names)

# ----------------------------------------------------------------
# PART H: COMPREHENSIVE XAI DASHBOARD
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART H: XAI Dashboard – Summary for a Single Prediction")
print("-"*60)

def xai_dashboard(model, X_data, instance_idx, feature_names):
    """
    Generate a comprehensive XAI dashboard for a single prediction.
    """
    instance = X_data.iloc[instance_idx:instance_idx+1]
    pred_prob = model.predict_proba(instance)[0, 1]
    pred_class = "Default" if pred_prob >= 0.5 else "No Default"
    
    # SHAP explanation
    shap_values = explainer_shap.shap_values(instance)
    shap_vals = shap_values[1] if isinstance(shap_values, list) else shap_values
    
    # Create dashboard
    print("\n" + "="*70)
    print(f"XAI DASHBOARD – Instance {instance_idx}")
    print("="*70)
    print(f"\nPrediction: {pred_class} (Probability: {pred_prob:.3f})")
    
    print("\nFeature Values:")
    for feature in feature_names:
        print(f"  {feature}: {instance[feature].values[0]:.2f}")
    
    print("\nSHAP Contributions:")
    # Sort by absolute SHAP value
    shap_df = pd.DataFrame({
        'feature': feature_names,
        'shap_value': shap_vals.flatten(),
        'feature_value': instance.values.flatten()
    }).sort_values('shap_value', ascending=False)
    
    for _, row in shap_df.iterrows():
        direction = "↑ (increases risk)" if row['shap_value'] > 0 else "↓ (decreases risk)"
        print(f"  {row['feature']:>15}: {row['shap_value']:+.4f} → {direction}")
    
    print("\nExplanation Summary:")
    # Identify the top positive and negative contributors
    top_positive = shap_df[shap_df['shap_value'] > 0].nlargest(2, 'shap_value')
    top_negative = shap_df[shap_df['shap_value'] < 0].nsmallest(2, 'shap_value')
    
    if not top_positive.empty:
        print("  Risk-increasing factors:")
        for _, row in top_positive.iterrows():
            print(f"    • {row['feature']}: {row['feature_value']:.2f} (contributes +{row['shap_value']:.3f})")
    
    if not top_negative.empty:
        print("  Risk-decreasing factors:")
        for _, row in top_negative.iterrows():
            print(f"    • {row['feature']}: {row['feature_value']:.2f} (contributes {row['shap_value']:.3f})")
    
    return shap_df

# Run dashboard for a sample instance
dashboard = xai_dashboard(rf_model, X_test, sample_idx, feature_names)

# ----------------------------------------------------------------
# PART I: REGULATORY COMPLIANCE CHECKLIST
# ----------------------------------------------------------------

print("\n" + "="*70)
print("PART I: Regulatory Compliance Checklist for XAI")
print("="*70)

checklist = [
    "✅ Model documentation includes feature importance (global interpretability).",
    "✅ Individual predictions can be explained (local interpretability).",
    "✅ SHAP or LIME is used for complex models.",
    "✅ Calibration plots show model probabilities match actual outcomes.",
    "✅ Disparate impact analysis has been conducted across protected groups.",
    "✅ Counterfactual explanations are available for rejected applications.",
    "✅ Model monitoring includes drift detection (PSI, CSI).",
    "✅ Regular model validation includes XAI review.",
    "✅ Explanation reports are stored for audit purposes.",
    "✅ Stakeholders have been trained on interpreting explanations."
]

for item in checklist:
    print(item)

print("\nRecommended XAI Framework for Banking:")
print("""
1. For all models: Provide global feature importance and PDPs.
2. For high-risk models (credit scoring): Provide SHAP for all predictions.
3. For customer-facing decisions: Provide LIME or counterfactual explanations.
4. For regulatory submissions: Include SHAP summary plots and feature importance.
5. For ongoing monitoring: Track SHAP distributions over time (drift in explanations).
""")