🏥 Healthcare Federated Learning with MATL¶

import torch
from matl import MATLClient, MATLMode
from matl.privacy import DifferentialPrivacy
from matl.audit import AuditLogger

# Initialize HIPAA-compliant MATL client
matl_client = MATLClient(
    mode=MATLMode.MODE2,  # Highest security (TEE-backed)
    node_id="hospital_a",
    private_key="/secure/hospital_a.pem",  # HSM-backed key

    # Differential privacy for PHI protection
    privacy=DifferentialPrivacy(
        epsilon=1.0,  # Privacy budget
        delta=1e-5,
        clip_norm=1.0,  # Gradient clipping
    ),

    # Comprehensive audit logging
    audit_logger=AuditLogger(
        backend="postgresql",  # HIPAA-compliant audit DB
        retention_years=7,  # HIPAA requirement
        encrypt=True,
    ),

    # Access control
    allowed_nodes=[
        "hospital_a",
        "hospital_b",
        "hospital_c",
        "hospital_d",
        "hospital_e",
    ],
)

print("✅ HIPAA-compliant MATL client initialized")

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import pandas as pd

class RetinalDataset(Dataset):
    """
    HIPAA-compliant retinal image dataset.
    All PHI stays local to the hospital.
    """
    def __init__(self, image_dir, labels_csv, transform=None):
        self.image_dir = image_dir
        self.labels = pd.read_csv(labels_csv)
        self.transform = transform or self.get_default_transform()

        # Verify no PHI in labels file
        assert "patient_id" not in self.labels.columns, "PHI detected!"
        assert "patient_name" not in self.labels.columns, "PHI detected!"

    def get_default_transform(self):
        return transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.RandomHorizontalFlip(),
            transforms.RandomRotation(10),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        # Load image (de-identified filename only)
        img_name = self.labels.iloc[idx]['image_id']  # Not patient_id!
        img_path = f"{self.image_dir}/{img_name}.jpg"
        image = Image.open(img_path).convert('RGB')

        # Load label (0=no DR, 1=mild, 2=moderate, 3=severe, 4=proliferative)
        label = self.labels.iloc[idx]['dr_level']

        if self.transform:
            image = self.transform(image)

        return image, label

# Load hospital-specific dataset
hospital_dataset = RetinalDataset(
    image_dir="/data/hospital_a/retinal_images",  # Local only!
    labels_csv="/data/hospital_a/labels_deidentified.csv"
)

print(f"Hospital A dataset: {len(hospital_dataset)} de-identified images")

import torch.nn as nn
from torchvision.models import resnet50, ResNet50_Weights

class DiabRetinoModel(nn.Module):
    """
    ResNet-50 adapted for 5-class diabetic retinopathy classification.
    """
    def __init__(self, pretrained=True):
        super(DiabRetinoModel, self).__init__()

        # Load pretrained ResNet-50 (ImageNet)
        if pretrained:
            weights = ResNet50_Weights.IMAGENET1K_V2
            self.backbone = resnet50(weights=weights)
        else:
            self.backbone = resnet50()

        # Replace final layer for 5-class classification
        num_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Sequential(
            nn.Linear(num_features, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 5)  # 5 DR severity levels
        )

    def forward(self, x):
        return self.backbone(x)

# Initialize model
model = DiabRetinoModel(pretrained=True)
print("✅ Medical AI model initialized")

from torch.optim import Adam
from torch.nn import CrossEntropyLoss

def train_federated_medical_ai(
    hospital_datasets,
    matl_client,
    num_rounds=50,
    local_epochs=3,
):
    """
    HIPAA-compliant federated learning for medical AI.

    Key features:
    - No data leaves hospitals
    - Byzantine resistance (handles Hospital E's bad data)
    - Differential privacy protection
    - Comprehensive audit logging
    """
    # Initialize global model
    global_model = DiabRetinoModel(pretrained=True)
    criterion = CrossEntropyLoss()

    # Training history
    history = {
        'rounds': [],
        'accuracy': [],
        'trust_scores': [],
        'privacy_budget': [],
    }

    for round_num in range(num_rounds):
        print(f"\n{'='*60}")
        print(f"Training Round {round_num + 1}/{num_rounds}")
        print(f"{'='*60}")

        round_gradients = []

        # Each hospital trains locally
        for hospital_id, hospital_data in enumerate(hospital_datasets):
            hospital_name = f"hospital_{chr(65 + hospital_id)}"  # A, B, C, D, E
            print(f"\n🏥 {hospital_name.upper()} training...")

            # Create local model (copy of global)
            local_model = DiabRetinoModel(pretrained=False)
            local_model.load_state_dict(global_model.state_dict())
            local_model.train()

            optimizer = Adam(local_model.parameters(), lr=0.0001)

            # Local training (PHI never leaves hospital)
            dataloader = DataLoader(
                hospital_data,
                batch_size=32,
                shuffle=True,
                num_workers=4
            )

            for epoch in range(local_epochs):
                total_loss = 0
                for images, labels in dataloader:
                    optimizer.zero_grad()
                    outputs = local_model(images)
                    loss = criterion(outputs, labels)
                    loss.backward()
                    optimizer.step()
                    total_loss += loss.item()

                avg_loss = total_loss / len(dataloader)
                print(f"  Epoch {epoch+1}/{local_epochs}: Loss = {avg_loss:.4f}")

            # Extract gradient (no PHI)
            gradient = [
                p.grad.clone() if p.grad is not None else torch.zeros_like(p)
                for p in local_model.parameters()
            ]

            # ============================================================
            # ⭐ MATL VALIDATION: Submit gradient for Byzantine resistance
            # ============================================================
            result = matl_client.submit_gradient(
                gradient=gradient,
                metadata={
                    "hospital": hospital_name,
                    "round": round_num,
                    "num_samples": len(hospital_data),
                    "local_epochs": local_epochs,
                }
            )

            print(f"  Trust Score: {result.trust_score:.3f}")
            print(f"  Privacy Budget Remaining: {result.privacy_budget:.2f}")

            # Audit log (HIPAA compliance)
            matl_client.audit_log(
                event="gradient_submission",
                hospital=hospital_name,
                round=round_num,
                trust_score=result.trust_score,
                privacy_budget=result.privacy_budget,
            )

            round_gradients.append({
                "gradient": gradient,
                "trust_score": result.trust_score,
                "hospital": hospital_name,
            })

        # ============================================================
        # ⭐ MATL AGGREGATION: Byzantine-resistant weighted average
        # ============================================================
        print(f"\n🔄 Aggregating gradients with Byzantine resistance...")

        aggregated_gradient = matl_client.aggregate(
            gradients=[g["gradient"] for g in round_gradients],
            trust_scores=[g["trust_score"] for g in round_gradients],
            method="reputation_weighted_with_privacy"
        )

        # Update global model
        for param, grad in zip(global_model.parameters(), aggregated_gradient):
            if param.grad is None:
                param.grad = grad
            else:
                param.grad.copy_(grad)

        optimizer = Adam(global_model.parameters(), lr=0.0001)
        optimizer.step()

        # Evaluate on holdout test set (aggregated, de-identified)
        accuracy = evaluate_medical_model(global_model, test_loader)

        # Update history
        history['rounds'].append(round_num)
        history['accuracy'].append(accuracy)
        history['trust_scores'].append([g["trust_score"] for g in round_gradients])
        history['privacy_budget'].append(matl_client.get_remaining_privacy_budget())

        print(f"\n📊 Round {round_num+1} Results:")
        print(f"  Global Model Accuracy: {accuracy:.2f}%")
        print(f"  Trust Scores: {[f'{ts:.2f}' for ts in history['trust_scores'][-1]]}")
        print(f"  Privacy Budget: {history['privacy_budget'][-1]:.2f}")

        # HIPAA audit checkpoint
        matl_client.audit_checkpoint(round_num, accuracy, history)

    return global_model, history

def evaluate_medical_model(model, test_loader):
    """Evaluate model on holdout test set"""
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in test_loader:
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    return accuracy

import matplotlib.pyplot as plt
import numpy as np

def plot_medical_fl_results(history):
    """
    Visualize federated learning results for medical AI.
    """
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))

    rounds = history['rounds']

    # Plot 1: Model Accuracy Over Time
    ax1.plot(rounds, history['accuracy'], 'b-', linewidth=2, marker='o')
    ax1.axhline(y=92, color='g', linestyle='--', label='Clinical Threshold (92%)')
    ax1.set_xlabel('Training Round')
    ax1.set_ylabel('Accuracy (%)')
    ax1.set_title('Diabetic Retinopathy Detection Accuracy')
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # Plot 2: Trust Scores by Hospital
    trust_matrix = np.array(history['trust_scores'])
    for hospital_idx in range(trust_matrix.shape[1]):
        hospital_name = f"Hospital {chr(65 + hospital_idx)}"
        ax2.plot(rounds, trust_matrix[:, hospital_idx], label=hospital_name, marker='o')
    ax2.set_xlabel('Training Round')
    ax2.set_ylabel('Trust Score')
    ax2.set_title('Hospital Trust Scores Over Time')
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    # Plot 3: Privacy Budget Consumption
    ax3.plot(rounds, history['privacy_budget'], 'r-', linewidth=2, marker='s')
    ax3.axhline(y=0.5, color='orange', linestyle='--', label='Warning Threshold')
    ax3.set_xlabel('Training Round')
    ax3.set_ylabel('Remaining Privacy Budget (ε)')
    ax3.set_title('Differential Privacy Budget')
    ax3.legend()
    ax3.grid(True, alpha=0.3)

    # Plot 4: Hospital Contribution Quality
    final_trust_scores = trust_matrix[-1, :]
    hospitals = [f"Hospital {chr(65 + i)}" for i in range(len(final_trust_scores))]
    colors = ['green' if ts > 0.7 else 'orange' if ts > 0.4 else 'red'
              for ts in final_trust_scores]
    ax4.bar(hospitals, final_trust_scores, color=colors)
    ax4.set_ylabel('Final Trust Score')
    ax4.set_title('Hospital Contribution Quality')
    ax4.axhline(y=0.7, color='g', linestyle='--', alpha=0.5)
    ax4.grid(True, alpha=0.3, axis='y')

    plt.tight_layout()
    plt.savefig('medical_fl_results.png', dpi=300, bbox_inches='tight')
    plt.show()

    # Summary statistics
    print("\n" + "="*60)
    print("MEDICAL FL TRAINING SUMMARY")
    print("="*60)
    print(f"Final Model Accuracy: {history['accuracy'][-1]:.2f}%")
    print(f"Clinical Threshold Met: {'✅ Yes' if history['accuracy'][-1] >= 92 else '❌ No'}")
    print(f"\nHospital Trust Scores:")
    for i, ts in enumerate(final_trust_scores):
        hospital = f"Hospital {chr(65 + i)}"
        status = "✅ Trusted" if ts > 0.7 else "⚠️  Marginal" if ts > 0.4 else "❌ Excluded"
        print(f"  {hospital}: {ts:.3f} {status}")
    print(f"\nPrivacy Budget Remaining: {history['privacy_budget'][-1]:.2f}")
    print(f"HIPAA Compliance: ✅ Maintained")
    print("="*60)

def verify_hipaa_compliance(matl_client, hospital_datasets):
    """
    Verify HIPAA compliance before production deployment.
    """
    print("\n🔍 HIPAA Compliance Verification")
    print("="*60)

    checks = []

    # Check 1: No PHI in datasets
    for i, dataset in enumerate(hospital_datasets):
        hospital = f"Hospital {chr(65 + i)}"
        has_phi = check_for_phi(dataset)
        checks.append((f"{hospital} - No PHI in dataset", not has_phi))

    # Check 2: Encryption enabled
    checks.append(("Data encryption enabled", matl_client.encryption_enabled))

    # Check 3: Audit logging active
    checks.append(("Audit logging active", matl_client.audit_logger.is_active()))

    # Check 4: Access control configured
    checks.append(("Access control configured", len(matl_client.allowed_nodes) > 0))

    # Check 5: Differential privacy enabled
    checks.append(("Differential privacy enabled", matl_client.privacy is not None))

    # Check 6: Secure communication (TLS)
    checks.append(("TLS encryption", matl_client.use_tls))

    # Check 7: Audit retention policy
    checks.append(("7-year audit retention", matl_client.audit_logger.retention_years >= 7))

    # Print results
    for check_name, passed in checks:
        status = "✅ PASS" if passed else "❌ FAIL"
        print(f"  {status} - {check_name}")

    all_passed = all(passed for _, passed in checks)
    print("="*60)

    if all_passed:
        print("✅ ALL HIPAA COMPLIANCE CHECKS PASSED")
        print("System ready for production deployment")
    else:
        print("❌ HIPAA COMPLIANCE FAILURES DETECTED")
        print("DO NOT DEPLOY TO PRODUCTION")
        raise RuntimeError("HIPAA compliance not met")

    return all_passed

def check_for_phi(dataset):
    """Check if dataset contains PHI"""
    # Implement PHI detection logic
    # Check for: patient names, SSN, medical record numbers, etc.
    return False  # Simplified for tutorial

def create_medical_fl_dashboard(matl_client):
    """
    Real-time monitoring dashboard for medical FL.
    HIPAA-compliant - no PHI displayed.
    """
    import streamlit as st

    st.title("🏥 Medical FL Monitoring Dashboard")
    st.markdown("**HIPAA-Compliant** - No PHI displayed")

    # System status
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Active Hospitals", matl_client.get_active_nodes())
    with col2:
        st.metric("Current Round", matl_client.get_current_round())
    with col3:
        st.metric("Model Accuracy", f"{matl_client.get_accuracy():.1f}%")

    # Trust scores
    st.subheader("Hospital Trust Scores")
    trust_df = matl_client.get_trust_scores_df()
    st.dataframe(trust_df)

    # Privacy budget
    st.subheader("Privacy Budget")
    privacy_progress = matl_client.get_remaining_privacy_budget()
    st.progress(privacy_progress)
    st.caption(f"Remaining: ε = {privacy_progress:.2f}")

    # Recent audits
    st.subheader("Recent Audit Events")
    audits = matl_client.audit_logger.get_recent(limit=10)
    st.table(audits)

    # Alerts
    if privacy_progress < 0.5:
        st.warning("⚠️ Privacy budget below 50% - consider reducing training frequency")

    low_trust_nodes = [
        node for node, score in trust_df.items() if score < 0.5
    ]
    if low_trust_nodes:
        st.error(f"❌ Low trust detected: {', '.join(low_trust_nodes)}")