# ===================================================================
# MODULE 7, LESSON 3: FEDERATED LEARNING FOR BANKING
# ===================================================================

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
from sklearn.datasets import make_classification
import warnings
warnings.filterwarnings('ignore')

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

print("="*70)
print("FEDERATED LEARNING FOR FINANCIAL APPLICATIONS")
print("="*70)

# ----------------------------------------------------------------
# PART A: SIMULATE FEDERATED DATA (NON-IID)
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART A: Simulating Federated Banking Data")
print("-"*60)

def generate_bank_data(n_samples=1000, n_features=10, n_informative=8, 
                       n_classes=2, n_clients=5, client_dist_shift=0.5):
    """
    Generate data for multiple clients (banks) with non-IID distributions.
    """
    clients_data = []
    for i in range(n_clients):
        # Shift the mean and covariance for each client
        mean_shift = np.random.normal(0, client_dist_shift, n_features)
        # Generate data
        X, y = make_classification(
            n_samples=n_samples,
            n_features=n_features,
            n_informative=n_informative,
            n_redundant=0,
            n_classes=n_classes,
            n_clusters_per_class=1,
            random_state=i + 42,
            shift=0.0,  # We'll add shift manually
            scale=1.0
        )
        # Add shift to features
        X = X + mean_shift
        # Add some label shift (some clients have higher default rates)
        if i < n_clients // 2:
            # Higher default rate for some clients
            y = (y + np.random.binomial(1, 0.1, n_samples)) % 2
        clients_data.append({
            'X': X,
            'y': y,
            'n_samples': n_samples,
            'client_id': i
        })
    return clients_data

# Generate data for 5 banks
n_clients = 5
n_samples_per_client = 2000
clients_data = generate_bank_data(
    n_samples=n_samples_per_client,
    n_features=10,
    n_informative=8,
    n_clients=n_clients,
    client_dist_shift=0.3
)

print(f"Generated data for {n_clients} clients (banks).")
print(f"Samples per client: {n_samples_per_client}")
print(f"Total samples: {n_clients * n_samples_per_client}")

# Show data distribution per client
for i, data in enumerate(clients_data):
    default_rate = data['y'].mean()
    print(f"Client {i}: default rate = {default_rate:.4f}")

# ----------------------------------------------------------------
# PART B: CENTRALISED BASELINE
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART B: Centralised Baseline")
print("-"*60)

# Combine all data
X_combined = np.vstack([data['X'] for data in clients_data])
y_combined = np.hstack([data['y'] for data in clients_data])

# Train centralised model
model_central = LogisticRegression(max_iter=1000, random_state=42)
model_central.fit(X_combined, y_combined)
y_pred_central = model_central.predict_proba(X_combined)[:, 1]
auc_central = roc_auc_score(y_combined, y_pred_central)

print(f"Centralised AUC: {auc_central:.4f}")

# Split into train/test (80/20)
from sklearn.model_selection import train_test_split
X_train_central, X_test_central, y_train_central, y_test_central = train_test_split(
    X_combined, y_combined, test_size=0.2, random_state=42
)
model_central_split = LogisticRegression(max_iter=1000, random_state=42)
model_central_split.fit(X_train_central, y_train_central)
y_pred_test = model_central_split.predict_proba(X_test_central)[:, 1]
auc_test_central = roc_auc_score(y_test_central, y_pred_test)
print(f"Centralised Test AUC: {auc_test_central:.4f}")

# ----------------------------------------------------------------
# PART C: FEDERATED AVERAGING (FEDAVG) IMPLEMENTATION
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART C: Federated Averaging (FedAvg)")
print("-"*60)

class FedAvg:
    """
    Simple Federated Averaging implementation for logistic regression.
    """
    def __init__(self, n_clients, n_features, n_classes=2, local_epochs=5, learning_rate=0.1):
        self.n_clients = n_clients
        self.n_features = n_features
        self.n_classes = n_classes
        self.local_epochs = local_epochs
        self.lr = learning_rate
        self.global_weights = np.random.randn(n_features, n_classes - 1) * 0.01
        self.global_bias = np.zeros(n_classes - 1)
        
        # Store history
        self.history = {'global_weights': [], 'auc': []}
    
    def _sigmoid(self, X, weights, bias):
        """Sigmoid function for logistic regression."""
        z = X @ weights + bias
        return 1 / (1 + np.exp(-z))
    
    def _logistic_loss(self, X, y, weights, bias):
        """Binary cross-entropy loss."""
        y_pred = self._sigmoid(X, weights, bias)
        y_pred = np.clip(y_pred, 1e-10, 1 - 1e-10)
        return -np.mean(y * np.log(y_pred) + (1 - y) * np.log(1 - y_pred))
    
    def _logistic_gradient(self, X, y, weights, bias):
        """Gradient of logistic loss."""
        y_pred = self._sigmoid(X, weights, bias)
        grad_w = (X.T @ (y_pred - y)) / len(X)
        grad_b = np.mean(y_pred - y)
        return grad_w, grad_b
    
    def local_train(self, X, y, weights, bias, epochs=5):
        """Train locally on client data."""
        for _ in range(epochs):
            grad_w, grad_b = self._logistic_gradient(X, y, weights, bias)
            weights = weights - self.lr * grad_w
            bias = bias - self.lr * grad_b
        return weights, bias
    
    def aggregate(self, client_weights, client_biases, n_samples):
        """Aggregate client updates using weighted averaging."""
        total_samples = sum(n_samples)
        new_weights = np.zeros_like(self.global_weights)
        new_bias = np.zeros_like(self.global_bias)
        
        for w, b, n in zip(client_weights, client_biases, n_samples):
            new_weights += (n / total_samples) * w
            new_bias += (n / total_samples) * b
        
        return new_weights, new_bias
    
    def fit(self, clients_data, n_rounds=20, client_fraction=1.0):
        """
        Train the federated model.
        """
        n_clients = len(clients_data)
        n_samples = [data['n_samples'] for data in clients_data]
        
        for round_idx in range(n_rounds):
            # Select clients (fraction)
            n_selected = max(1, int(client_fraction * n_clients))
            selected_indices = np.random.choice(n_clients, n_selected, replace=False)
            
            client_weights = []
            client_biases = []
            selected_n_samples = []
            
            # Local training on selected clients
            for idx in selected_indices:
                X = clients_data[idx]['X']
                y = clients_data[idx]['y']
                
                # Start from global weights
                local_w = self.global_weights.copy()
                local_b = self.global_bias.copy()
                
                # Local training
                local_w, local_b = self.local_train(
                    X, y, local_w, local_b, epochs=self.local_epochs
                )
                
                client_weights.append(local_w)
                client_biases.append(local_b)
                selected_n_samples.append(n_samples[idx])
            
            # Aggregate
            self.global_weights, self.global_bias = self.aggregate(
                client_weights, client_biases, selected_n_samples
            )
            
            # Evaluate on held-out test set (if available)
            # For demonstration, we evaluate on a global test set (if available)
            # In practice, use a central validation set.
            if hasattr(self, 'X_test') and hasattr(self, 'y_test'):
                auc = self.evaluate(self.X_test, self.y_test)
                self.history['auc'].append(auc)
            
            self.history['global_weights'].append(self.global_weights.copy())
            
            if round_idx % 5 == 0:
                print(f"Round {round_idx+1}/{n_rounds}: AUC = {self.history['auc'][-1]:.4f}" if self.history['auc'] else "")
    
    def evaluate(self, X, y):
        """Evaluate the global model."""
        y_pred = self._sigmoid(X, self.global_weights, self.global_bias)
        return roc_auc_score(y, y_pred)
    
    def predict_proba(self, X):
        """Predict probabilities."""
        return self._sigmoid(X, self.global_weights, self.global_bias)

# Split combined data for central validation
X_train_cv, X_test_cv, y_train_cv, y_test_cv = train_test_split(
    X_combined, y_combined, test_size=0.2, random_state=42
)

# Initialise FedAvg
fedavg = FedAvg(
    n_clients=n_clients,
    n_features=X_combined.shape[1],
    n_classes=2,
    local_epochs=3,
    learning_rate=0.1
)

# Add test set for evaluation
fedavg.X_test = X_test_cv
fedavg.y_test = y_test_cv

# Train
print("\nTraining Federated Model...")
fedavg.fit(clients_data, n_rounds=20, client_fraction=0.8)

# Final evaluation on test set
final_auc = fedavg.evaluate(X_test_cv, y_test_cv)
print(f"\nFederated Test AUC: {final_auc:.4f}")

# ----------------------------------------------------------------
# PART D: COMPARISON OF APPROACHES
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART D: Comparison of Approaches")
print("-"*60)

# Individual client performance
print("\nIndividual Client Performance:")
client_aucs = []
for i, data in enumerate(clients_data):
    X = data['X']
    y = data['y']
    X_train_c, X_test_c, y_train_c, y_test_c = train_test_split(
        X, y, test_size=0.3, random_state=42
    )
    model_local = LogisticRegression(max_iter=1000, random_state=42)
    model_local.fit(X_train_c, y_train_c)
    y_pred_c = model_local.predict_proba(X_test_c)[:, 1]
    auc_c = roc_auc_score(y_test_c, y_pred_c)
    client_aucs.append(auc_c)
    print(f"  Client {i}: AUC = {auc_c:.4f}")

# Comparison table
comparison = pd.DataFrame({
    'Method': ['Centralised', 'Federated (FedAvg)', 'Local (Average)'],
    'Test AUC': [auc_test_central, final_auc, np.mean(client_aucs)],
    'Privacy': ['Low', 'High', 'High'],
    'Data Sharing Required': ['Yes', 'No', 'No'],
    'Communication Cost': ['None', 'Low', 'None']
})
print("\n" + "="*70)
print("COMPARISON:")
print(comparison.to_string(index=False))

print("\nInterpretation:")
print("  • Centralised: Best performance but requires data sharing (privacy risk).")
print("  • Federated: Good performance without data sharing (privacy-preserving).")
print("  • Local: Worst performance (limited data).")
print("  • Federated closes the gap to centralised while preserving privacy.")

# ----------------------------------------------------------------
# PART E: VISUALISATION – FEDERATED LEARNING CONVERGENCE
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART E: Visualising Federated Learning Convergence")
print("-"*60)

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

# Convergence of global model
ax = axes[0]
if fedavg.history['auc']:
    ax.plot(fedavg.history['auc'], 'b-', linewidth=2, label='Federated')
    ax.axhline(y=auc_test_central, color='red', linestyle='--', label=f'Centralised: {auc_test_central:.4f}')
    ax.axhline(y=np.mean(client_aucs), color='green', linestyle=':', label=f'Local Average: {np.mean(client_aucs):.4f}')
    ax.set_xlabel('Communication Round')
    ax.set_ylabel('AUC')
    ax.set_title('Federated Learning Convergence')
    ax.legend()
    ax.grid(True, alpha=0.3)

# Client vs Global Performance
ax = axes[1]
x = np.arange(n_clients)
client_bars = ax.bar(x - 0.2, client_aucs, 0.4, label='Local Model', color='blue', alpha=0.7)
ax.bar(x + 0.2, [final_auc] * n_clients, 0.4, label='Federated Model', color='green', alpha=0.7)
ax.set_xticks(x)
ax.set_xticklabels([f'Client {i}' for i in range(n_clients)])
ax.set_ylabel('AUC')
ax.set_title('Local vs Federated Performance')
ax.legend()
ax.grid(True, alpha=0.3)

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

# ----------------------------------------------------------------
# PART F: PRIVACY AND SECURITY CONSIDERATIONS
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART F: Privacy and Security in Federated Learning")
print("-"*60)

print("""
Privacy Enhancements:

1. Differential Privacy (DP):
   - Add calibrated noise to gradients before transmission.
   - Trade-off: Privacy vs accuracy.

2. Secure Aggregation (SecAgg):
   - Encrypt client updates so the server cannot see individual contributions.
   - Uses cryptographic techniques (secret sharing, homomorphic encryption).

3. Homomorphic Encryption (HE):
   - Perform computations on encrypted data.
   - Server can aggregate without decryption.
   - High computational overhead.

Security Threats:
  - Model Poisoning: Malicious clients send bad updates.
  - Inference Attacks: Attackers infer client data from model updates.
  - Free-riding: Clients benefit without contributing.

Mitigations:
  - Robust aggregation (median, trimmed mean).
  - Anomaly detection on client updates.
  - Client validation and authentication.
  - Differential privacy to limit information leakage.
""")

# ----------------------------------------------------------------
# PART G: BANKING USE CASES
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART G: Banking Use Cases for Federated Learning")
print("-"*60)

use_cases = {
    "Cross-Bank Fraud Detection": {
        "Description": "Multiple banks collaborate to train a fraud detection model.",
        "Data": "Transaction data (local).",
        "Privacy": "Transactions never leave each bank.",
        "Benefit": "Improves model coverage and accuracy."
    },
    "Credit Scoring Consortium": {
        "Description": "Banks share model updates to improve credit scoring.",
        "Data": "Customer profiles and repayment history.",
        "Privacy": "Customer data stays with the bank.",
        "Benefit": "Better risk assessment for underserved populations."
    },
    "AML (Anti-Money Laundering)": {
        "Description": "Detect money laundering patterns across banks.",
        "Data": "Transaction patterns and suspicious activity reports.",
        "Privacy": "Sensitive transaction data is protected.",
        "Benefit": "More comprehensive AML coverage."
    },
    "Personalised Banking": {
        "Description": "Train personalised models on customer devices.",
        "Data": "Customer spending habits and preferences.",
        "Privacy": "Data stays on the customer's device.",
        "Benefit": "Better recommendations without privacy intrusion."
    },
    "Insurance Risk Assessment": {
        "Description": "Collaborate with insurers for better risk models.",
        "Data": "Claims data and policyholder information.",
        "Privacy": "Policyholder data is protected.",
        "Benefit": "More accurate risk pricing."
    }
}

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

# ----------------------------------------------------------------
# PART H: REGULATORY PERSPECTIVE
# ----------------------------------------------------------------

print("\n" + "-"*60)
print("PART H: Regulatory Perspective on Federated Learning")
print("-"*60)

print("""
Regulatory Acceptance:

1. GDPR (EU):
   - FL is compliant with GDPR (data minimisation, privacy by design).
   - Data never leaves the organisation (data controller remains in control).
   - However, model updates must be protected to prevent inference attacks.

2. CCPA (California):
   - Similar to GDPR: FL reduces the risk of data breaches.
   - Consumers' "right to delete" is preserved (data is not centralised).

3. SR 11-7 (US):
   - Model validation in FL is more complex.
   - Banks must validate the global model and understand local variations.
   - Explainability (SHAP) is required for regulatory submissions.

4. Basel III:
   - FL can be used for risk models while maintaining data sovereignty.
   - Must demonstrate model robustness across diverse client data.

5. EU AI Act:
   - High-risk AI (credit scoring, fraud detection) requires transparency.
   - FL models must be explainable and auditable.

Best Practices:
  - Document the FL training process (rounds, clients, aggregation).
  - Validate models on a centralised hold-out set.
  - Implement differential privacy for model updates.
  - Maintain audit trails of model versions and client contributions.
  - Engage with regulators early in FL adoption.
""")

# ----------------------------------------------------------------
# PART I: SUMMARY AND RECOMMENDATIONS
# ----------------------------------------------------------------

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

print("""
Federated Learning – Key Takeaways:

1. FL enables collaborative model training without sharing raw data.
2. FedAvg is the foundational algorithm (weighted averaging of client updates).
3. Privacy can be enhanced with differential privacy, secure aggregation, and homomorphic encryption.
4. Challenges: Non-IID data, communication efficiency, adversarial attacks, and model validation.
5. Banking use cases: Fraud detection, credit scoring, AML, personalised banking.
6. Regulatory acceptance is growing – FL aligns with GDPR/CCPA privacy requirements.
7. Trade-offs: Privacy vs accuracy, communication cost vs model quality.

Recommendations:
  - Start with a pilot involving 2-3 partner banks.
  - Use FedAvg as the baseline algorithm.
  - Implement differential privacy for model updates.
  - Validate models on a centralised test set.
  - Document FL training for regulatory compliance.
  - Explore secure aggregation for high-sensitivity applications.
  - Consider federated learning for cross-border data collaboration.
""")

print("="*70)
print("END OF LESSON 3 – MODULE 7")
print("="*70)