from fairlearn.reductions import ExponentiatedGradient, DemographicParity
from fairlearn.metrics import demographic_parity_difference

# Example: Demographic Parity constraint
constraint = DemographicParity()
mitigator = ExponentiatedGradient(model, constraint, eps=0.01)
mitigator.fit(X_train, y_train, sensitive_features=A_train)
y_pred_fair = mitigator.predict(X_test)

# ===================================================================
# MODULE 6, LESSON 5: AI GOVERNANCE, ETHICS, AND RESPONSIBLE AI
# ===================================================================

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.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, roc_auc_score, confusion_matrix
import warnings
warnings.filterwarnings('ignore')

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

print("="*70)
print("AI GOVERNANCE, ETHICS, AND RESPONSIBLE AI IN BANKING")
print("="*70)

# ----------------------------------------------------------------
# PART A: GENERATE SYNTHETIC DATA WITH PROTECTED ATTRIBUTES
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART A: Synthetic Credit Data with Protected Attributes")
print("-"*60)

# Generate data with two groups: Group A (favoured) and Group B (protected)
n_samples = 5000
n_group_a = int(0.7 * n_samples)
n_group_b = n_samples - n_group_a

# Group A: Higher income, higher credit score
income_a = np.random.gamma(5, 20, n_group_a) + 40
credit_score_a = np.random.normal(720, 40, n_group_a).clip(600, 850)
dti_a = np.random.beta(2, 5, n_group_a) * 40

# Group B: Lower income, lower credit score
income_b = np.random.gamma(4, 15, n_group_b) + 25
credit_score_b = np.random.normal(680, 50, n_group_b).clip(550, 800)
dti_b = np.random.beta(2, 4, n_group_b) * 50

# Combine
income = np.concatenate([income_a, income_b])
credit_score = np.concatenate([credit_score_a, credit_score_b])
dti = np.concatenate([dti_a, dti_b])
group = np.array(['A'] * n_group_a + ['B'] * n_group_b)

# Generate default labels (with slight bias toward Group A)
# Group A has lower default probability for same characteristics
log_odds_a = -4.5 + 0.04 * dti - 0.005 * credit_score + 0.01 * (income/50)
log_odds_b = -4.0 + 0.04 * dti - 0.005 * credit_score + 0.01 * (income/50)

# Slight bias: Group B has higher default probability
log_odds = np.where(group == 'A', log_odds_a, log_odds_b)
prob_default = 1 / (1 + np.exp(-log_odds))
default = np.random.binomial(1, prob_default)

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

print("Data Summary:")
print(df.groupby('group').agg({
    'income': 'mean',
    'credit_score': 'mean',
    'dti': 'mean',
    'default': 'mean'
}).round(2))

print(f"\nDefault Rate by Group:")
print(df.groupby('group')['default'].mean().round(4))

# ----------------------------------------------------------------
# PART B: TRAIN MODEL (WITHOUT FAIRNESS CONSIDERATIONS)
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART B: Baseline Model (No Fairness Consideration)")
print("-"*60)

# Features (excluding the protected attribute)
feature_cols = ['income', 'credit_score', 'dti']
X = df[feature_cols]
y = df['default']
sensitive = df['group']

X_train, X_test, y_train, y_test, s_train, s_test = train_test_split(
    X, y, sensitive, test_size=0.3, random_state=42
)

# Train Random Forest
rf_model = RandomForestClassifier(n_estimators=100, max_depth=8, random_state=42)
rf_model.fit(X_train, y_train)

# Predictions
y_pred_proba = rf_model.predict_proba(X_test)[:, 1]
y_pred = (y_pred_proba >= 0.5).astype(int)

# Overall performance
auc = roc_auc_score(y_test, y_pred_proba)
accuracy = accuracy_score(y_test, y_pred)
print(f"Overall AUC: {auc:.4f}")
print(f"Overall Accuracy: {accuracy:.4f}")

# Performance by group
print("\nPerformance by Group:")
for group_val in ['A', 'B']:
    mask = s_test == group_val
    group_auc = roc_auc_score(y_test[mask], y_pred_proba[mask])
    group_accuracy = accuracy_score(y_test[mask], y_pred[mask])
    group_approval_rate = y_pred[mask].mean()
    print(f"  Group {group_val}:")
    print(f"    AUC: {group_auc:.4f}")
    print(f"    Accuracy: {group_accuracy:.4f}")
    print(f"    Approval Rate (Predicted Positive): {group_approval_rate:.4f}")

# ----------------------------------------------------------------
# PART C: FAIRNESS METRICS – BASELINE
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART C: Fairness Metrics (Baseline Model)")
print("-"*60)

# Calculate fairness metrics
def fairness_metrics(y_true, y_pred, sensitive):
    """
    Calculate key fairness metrics.
    """
    # Overall metrics
    metrics = {}
    
    # Group-specific metrics
    groups = np.unique(sensitive)
    group_stats = {}
    
    for g in groups:
        mask = sensitive == g
        group_stats[g] = {
            'approval_rate': y_pred[mask].mean(),
            'tpr': np.mean(y_pred[mask] & y_true[mask]) / np.mean(y_true[mask]) if np.mean(y_true[mask]) > 0 else 0,
            'fpr': np.mean(y_pred[mask] & ~y_true[mask]) / np.mean(~y_true[mask]) if np.mean(~y_true[mask]) > 0 else 0,
            'count': mask.sum()
        }
    
    # Demographic parity difference
    rates = [group_stats[g]['approval_rate'] for g in groups]
    metrics['demographic_parity_difference'] = max(rates) - min(rates)
    metrics['demographic_parity_ratio'] = min(rates) / max(rates) if max(rates) > 0 else 0
    
    # Equal opportunity difference (TPR)
    tprs = [group_stats[g]['tpr'] for g in groups]
    metrics['equal_opportunity_difference'] = max(tprs) - min(tprs)
    
    # Equal odds difference (max of TPR and FPR differences)
    fprs = [group_stats[g]['fpr'] for g in groups]
    metrics['equal_odds_difference'] = max(
        max(tprs) - min(tprs),
        max(fprs) - min(fprs)
    )
    
    return metrics, group_stats

metrics_base, group_stats_base = fairness_metrics(y_test, y_pred, s_test)

print("Fairness Metrics (Baseline):")
print(f"  Demographic Parity Difference: {metrics_base['demographic_parity_difference']:.4f}")
print(f"  Demographic Parity Ratio: {metrics_base['demographic_parity_ratio']:.4f} (Threshold: 0.8)")
print(f"  Equal Opportunity Difference: {metrics_base['equal_opportunity_difference']:.4f}")
print(f"  Equal Odds Difference: {metrics_base['equal_odds_difference']:.4f}")

print("\nGroup Statistics:")
for g, stats in group_stats_base.items():
    print(f"  Group {g}:")
    print(f"    Approval Rate: {stats['approval_rate']:.4f}")
    print(f"    TPR: {stats['tpr']:.4f}")
    print(f"    FPR: {stats['fpr']:.4f}")
    print(f"    Count: {stats['count']}")

# Check if fairness is violated
if metrics_base['demographic_parity_ratio'] < 0.8:
    print("\n⚠ DISPARATE IMPACT: The 4/5 rule is violated (ratio < 0.8).")
    print(f"   Ratio: {metrics_base['demographic_parity_ratio']:.4f} < 0.8")
else:
    print("\n✓ No disparate impact detected (ratio ≥ 0.8).")

# ----------------------------------------------------------------
# PART D: FAIRNESS-AWARE MITIGATION – RE-WEIGHTING
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART D: Fairness Mitigation – Re-weighting Training Data")
print("-"*60)

# Calculate sample weights to balance groups
def compute_sample_weights(y_train, s_train):
    """
    Compute sample weights to balance groups.
    Weight = 1 / (group_proportion)
    """
    group_counts = s_train.value_counts()
    total = len(s_train)
    weights = np.zeros(len(s_train))
    for g in group_counts.index:
        mask = s_train == g
        weights[mask] = total / group_counts[g] / len(group_counts)
    return weights

# Compute weights
weights = compute_sample_weights(y_train, s_train)

print("Sample Weights:")
for g in ['A', 'B']:
    mask = s_train == g
    print(f"  Group {g}: mean weight = {weights[mask].mean():.3f}")

# Train weighted model
rf_weighted = RandomForestClassifier(n_estimators=100, max_depth=8, random_state=42)
rf_weighted.fit(X_train, y_train, sample_weight=weights)

# Predictions
y_pred_proba_w = rf_weighted.predict_proba(X_test)[:, 1]
y_pred_w = (y_pred_proba_w >= 0.5).astype(int)

# Evaluate performance
auc_w = roc_auc_score(y_test, y_pred_proba_w)
accuracy_w = accuracy_score(y_test, y_pred_w)
print(f"\nWeighted Model AUC: {auc_w:.4f}")
print(f"Weighted Model Accuracy: {accuracy_w:.4f}")

# Fairness metrics for weighted model
metrics_w, group_stats_w = fairness_metrics(y_test, y_pred_w, s_test)

print("\nFairness Metrics (Weighted Model):")
print(f"  Demographic Parity Difference: {metrics_w['demographic_parity_difference']:.4f}")
print(f"  Demographic Parity Ratio: {metrics_w['demographic_parity_ratio']:.4f}")

print("\nGroup Statistics (Weighted Model):")
for g, stats in group_stats_w.items():
    print(f"  Group {g}: Approval Rate = {stats['approval_rate']:.4f}, TPR = {stats['tpr']:.4f}")

# ----------------------------------------------------------------
# PART E: FAIRNESS-AWARE MITIGATION – POST-PROCESSING
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART E: Fairness Mitigation – Post-Processing (Threshold Adjustment)")
print("-"*60)

# Find optimal thresholds for each group to equalise approval rates
def find_group_thresholds(y_true, y_pred_proba, sensitive, target_ratio=0.05):
    """
    Find thresholds for each group to achieve demographic parity.
    """
    groups = np.unique(sensitive)
    thresholds = {}
    
    # Overall threshold
    overall_threshold = 0.5
    
    # Adjust thresholds to equalise approval rates
    # This is a simplified approach
    for g in groups:
        mask = sensitive == g
        group_proba = y_pred_proba[mask]
        # Find threshold that gives desired approval rate
        # For simplicity, we use the overall mean as target
        target_approval = np.mean(y_pred_proba) 
        # Find threshold closest to target
        sorted_proba = np.sort(group_proba)
        target_idx = int(len(sorted_proba) * (1 - target_approval))
        target_idx = max(0, min(target_idx, len(sorted_proba) - 1))
        thresholds[g] = sorted_proba[target_idx]
    
    return thresholds

# Compute group thresholds
thresholds = find_group_thresholds(y_test, y_pred_proba, s_test)

print("Thresholds by Group:")
for g in thresholds:
    print(f"  Group {g}: {thresholds[g]:.4f}")

# Apply thresholds
y_pred_pp = np.zeros(len(y_test))
for i in range(len(y_test)):
    g = s_test.iloc[i]
    y_pred_pp[i] = 1 if y_pred_proba[i] >= thresholds[g] else 0

# Evaluate
metrics_pp, group_stats_pp = fairness_metrics(y_test, y_pred_pp, s_test)

print("\nFairness Metrics (Post-Processing):")
print(f"  Demographic Parity Difference: {metrics_pp['demographic_parity_difference']:.4f}")
print(f"  Demographic Parity Ratio: {metrics_pp['demographic_parity_ratio']:.4f}")

print("\nGroup Statistics (Post-Processing):")
for g, stats in group_stats_pp.items():
    print(f"  Group {g}: Approval Rate = {stats['approval_rate']:.4f}, TPR = {stats['tpr']:.4f}")

# ----------------------------------------------------------------
# PART F: COMPARISON OF MODELS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART F: Model Comparison – Fairness vs Performance Trade-off")
print("-"*60)

comparison = pd.DataFrame({
    'Model': ['Baseline', 'Weighted', 'Post-Processing'],
    'AUC': [auc, auc_w, roc_auc_score(y_test, y_pred_pp)],
    'Accuracy': [accuracy, accuracy_w, accuracy_score(y_test, y_pred_pp)],
    'Demographic Parity Ratio': [metrics_base['demographic_parity_ratio'], 
                                  metrics_w['demographic_parity_ratio'],
                                  metrics_pp['demographic_parity_ratio']],
    'Demographic Parity Diff': [metrics_base['demographic_parity_difference'],
                                 metrics_w['demographic_parity_difference'],
                                 metrics_pp['demographic_parity_difference']]
})
print(comparison.round(4).to_string(index=False))

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

# Fairness vs Performance
ax = axes[0]
ax.scatter(comparison['Demographic Parity Ratio'], comparison['AUC'], s=100)
for i, row in comparison.iterrows():
    ax.annotate(row['Model'], (row['Demographic Parity Ratio'], row['AUC']))
ax.axvline(x=0.8, color='red', linestyle='--', label='4/5 Rule Threshold')
ax.set_xlabel('Demographic Parity Ratio')
ax.set_ylabel('AUC')
ax.set_title('Fairness vs Performance Trade-off')
ax.legend()
ax.grid(True, alpha=0.3)

# Approval rates by group
ax = axes[1]
models = comparison['Model'].values
group_approval_base = [group_stats_base['A']['approval_rate'], group_stats_base['B']['approval_rate']]
group_approval_w = [group_stats_w['A']['approval_rate'], group_stats_w['B']['approval_rate']]
group_approval_pp = [group_stats_pp['A']['approval_rate'], group_stats_pp['B']['approval_rate']]

x = np.arange(2)
width = 0.25
ax.bar(x - width, group_approval_base, width, label='Baseline', color='blue', alpha=0.7)
ax.bar(x, group_approval_w, width, label='Weighted', color='green', alpha=0.7)
ax.bar(x + width, group_approval_pp, width, label='Post-Processing', color='orange', alpha=0.7)
ax.set_xticks(x)
ax.set_xticklabels(['Group A', 'Group B'])
ax.set_ylabel('Approval Rate')
ax.set_title('Approval Rates by Group and Model')
ax.legend()
ax.grid(True, alpha=0.3)

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

# ----------------------------------------------------------------
# PART G: AI GOVERNANCE FRAMEWORK
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART G: AI Governance Framework for Banking")
print("-"*60)

governance = {
    "1. Governance Structure": {
        "Board Oversight": "Board-level AI committee with risk and compliance representation.",
        "Chief AI Officer": "Executive responsible for AI strategy and governance.",
        "AI Ethics Committee": "Cross-functional team reviewing AI use cases.",
        "Model Governance": "Formal model risk management framework (SR 11-7)."
    },
    "2. Risk Management": {
        "Risk Assessment": "AI risk classification (low, medium, high) based on impact.",
        "Model Validation": "Independent validation of all material AI models.",
        "Monitoring": "Ongoing performance monitoring and drift detection.",
        "Incident Response": "Process for handling AI failures and complaints."
    },
    "3. Fairness and Bias": {
        "Fairness Testing": "Regular testing for disparate impact and bias.",
        "Mitigation": "Remediation of identified bias (re-weighting, constraints).",
        "Transparency": "Documentation of fairness metrics and decisions."
    },
    "4. Transparency and Explainability": {
        "Model Documentation": "Model cards and datasheets for all models.",
        "Explainability": "SHAP/LIME explanations for high-risk decisions.",
        "Disclosure": "Customer-facing explanations in plain language."
    },
    "5. Data Governance": {
        "Data Privacy": "Compliance with GDPR, CCPA, and data protection laws.",
        "Data Quality": "Data lineage, quality monitoring, and handling of missing data.",
        "Sensitive Data": "Special handling of protected attributes."
    },
    "6. Human Oversight": {
        "Human-in-the-loop": "Manual review for high-risk decisions.",
        "Override Mechanisms": "Ability to override AI decisions.",
        "Training": "Staff training on AI ethics and governance."
    }
}

for pillar, details in governance.items():
    print(f"\n{pillar}:")
    for key, value in details.items():
        print(f"  {key}: {value}")

# ----------------------------------------------------------------
# PART H: MODEL CARD TEMPLATE
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART H: Model Card Template (Transparency Documentation)")
print("-"*60)

model_card = """
--- MODEL CARD ---

1. MODEL DETAILS
   - Model Name: Credit Default Predictor v3.2
   - Model Type: Random Forest Classifier
   - Developer: Credit Risk Analytics Team
   - Date: 2026-06-15
   - Version: 3.2

2. INTENDED USE
   - Primary Use: Credit underwriting and risk assessment
   - Secondary Use: Portfolio monitoring and stress testing
   - Out-of-Scope: Automated lending decisions without human review

3. DATA
   - Training Data: 2018-2025 loan origination data
   - Features: income, credit_score, dti, loan_amount, employment_years, age
   - Protected Attributes: None used directly (but proxy risk exists)
   - Sample Size: 35,000 loans
   - Default Rate: 4.2%

4. PERFORMANCE
   - AUC: 0.82
   - Accuracy: 0.79
   - Calibration: Hosmer-Lemeshow p-value = 0.12
   - Fairness: Demographic parity ratio = 0.85 (within 4/5 rule)

5. FAIRNESS ANALYSIS
   - Protected Attributes Considered: Gender, Age, Race (proxy testing)
   - Disparate Impact: Ratio = 0.85 (acceptable)
   - Mitigation: Re-weighting applied to training data
   - Monitoring: Quarterly fairness review

6. LIMITATIONS
   - Model may not perform well during severe economic downturns.
   - Limited data on certain customer segments.
   - Not validated for use with alternative data (e.g., social media).

7. GOVERNANCE
   - Model Owner: Chief Credit Officer
   - Validator: Model Validation Team
   - Review Frequency: Annual (or triggered by performance drift)

8. ETHICAL CONSIDERATIONS
   - Potential for bias through proxy variables.
   - Regular monitoring for disparate impact.
   - Human review for all rejected applications.
"""
print(model_card)

# ----------------------------------------------------------------
# PART I: RECOMMENDATIONS AND BEST PRACTICES
# ----------------------------------------------------------------

print("\n" + "="*70)
print("PART I: Recommendations and Best Practices")
print("="*70)

print("""
Key Recommendations for Responsible AI in Banking:

1. Adopt a risk-based approach:
   - Classify AI use cases by risk level (low, medium, high).
   - Apply proportional governance (more oversight for high-risk).

2. Implement fairness testing:
   - Test all models for disparate impact (4/5 rule).
   - Use multiple fairness metrics (demographic parity, equal opportunity).
   - Document fairness results and mitigation actions.

3. Ensure explainability:
   - Use SHAP/LIME for all high-risk models.
   - Provide customer-friendly explanations for adverse decisions.
   - Train staff on interpreting explanations.

4. Establish governance structures:
   - Board-level oversight.
   - Chief AI Officer role.
   - AI Ethics Committee with diverse membership.

5. Monitor continuously:
   - Track model performance over time.
   - Monitor for drift (data drift, concept drift, fairness drift).
   - Trigger revalidation when thresholds are breached.

6. Maintain transparency:
   - Publish model cards for all models.
   - Document data sources, features, and limitations.
   - Be open about AI use with customers and regulators.

7. Invest in talent:
   - Train staff on AI ethics and governance.
   - Hire diverse teams to reduce blind spots.
   - Engage external experts for independent review.
""")

print("\n" + "="*70)
print("END OF LESSON 5 – MODULE 6")
print("="*70)