┌─────────────────┐
         │   중앙 서버      │
         │  (Aggregator)   │
         └────────┬────────┘
                  │ 모델 배포 / 업데이트 집계
         ┌────────┼────────┐
         ↓        ↓        ↓
    ┌─────────┐ ┌─────────┐ ┌─────────┐
    │클라이언트1│ │클라이언트2│ │클라이언트3│
    │(로컬 데이│ │(로컬 데이│ │(로컬 데이│
    │터로 학습)│ │터로 학습)│ │터로 학습)│
    └─────────┘ └─────────┘ └─────────┘

클라이언트 1: [특성1, 특성2, 특성3] × [샘플 1~1000]
클라이언트 2: [특성1, 특성2, 특성3] × [샘플 1001~2000]
클라이언트 3: [특성1, 특성2, 특성3] × [샘플 2001~3000]

클라이언트 A (은행): [금융 특성] × [사용자 1~10000]
클라이언트 B (병원): [의료 특성] × [사용자 1~10000]

서버 실행:
  w_0 초기화
  for 라운드 t = 1, 2, ..., T:
    m = max(C × K, 1)  // C: 참여 비율, K: 전체 클라이언트 수
    S_t = 무작위로 m개 클라이언트 선택
    for 각 클라이언트 k in S_t (병렬):
      w_{t+1}^k = ClientUpdate(k, w_t)
    w_{t+1} = Σ (n_k / n) × w_{t+1}^k  // 가중 평균

클라이언트 k 실행:
  B = 로컬 데이터를 배치로 분할
  for 로컬 에폭 e = 1, ..., E:
    for 배치 b in B:
      w = w - η × ∇ℓ(w; b)
  return w

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Subset
import numpy as np
from copy import deepcopy
from typing import List, Dict, Tuple
import random


# ========== 모델 정의 ==========

class SimpleNet(nn.Module):
    """간단한 분류 네트워크"""
    def __init__(self, input_dim: int, hidden_dim: int, num_classes: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, num_classes)
        )

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


# ========== 클라이언트 ==========

class FLClient:
    """연합 학습 클라이언트"""

    def __init__(
        self,
        client_id: int,
        dataset,
        device: str = 'cpu'
    ):
        self.client_id = client_id
        self.dataset = dataset
        self.device = device

    def local_train(
        self,
        model: nn.Module,
        local_epochs: int,
        batch_size: int,
        lr: float
    ) -> Tuple[Dict, float]:
        """
        로컬 데이터로 모델 훈련
        Returns: (업데이트된 가중치, 로컬 손실)
        """
        model = deepcopy(model).to(self.device)
        model.train()

        loader = DataLoader(
            self.dataset,
            batch_size=batch_size,
            shuffle=True
        )
        optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
        criterion = nn.CrossEntropyLoss()

        total_loss = 0.0
        n_batches = 0

        for epoch in range(local_epochs):
            for X, y in loader:
                X, y = X.to(self.device), y.to(self.device)
                optimizer.zero_grad()
                output = model(X)
                loss = criterion(output, y)
                loss.backward()
                optimizer.step()

                total_loss += loss.item()
                n_batches += 1

        avg_loss = total_loss / max(n_batches, 1)
        return model.state_dict(), avg_loss

    def evaluate(self, model: nn.Module) -> Tuple[float, float]:
        """로컬 데이터로 모델 평가"""
        model = deepcopy(model).to(self.device)
        model.eval()

        loader = DataLoader(self.dataset, batch_size=64)
        criterion = nn.CrossEntropyLoss()

        total_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():
            for X, y in loader:
                X, y = X.to(self.device), y.to(self.device)
                output = model(X)
                loss = criterion(output, y)
                total_loss += loss.item()

                _, predicted = output.max(1)
                total += y.size(0)
                correct += predicted.eq(y).sum().item()

        return total_loss / len(loader), correct / total


# ========== 서버 ==========

class FedAvgServer:
    """FedAvg 서버"""

    def __init__(
        self,
        global_model: nn.Module,
        clients: List[FLClient],
        fraction: float = 0.1,  # C
        device: str = 'cpu'
    ):
        self.global_model = global_model.to(device)
        self.clients = clients
        self.fraction = fraction
        self.device = device
        self.round_history = []

    def select_clients(self) -> List[FLClient]:
        """각 라운드에서 클라이언트 선택"""
        m = max(int(self.fraction * len(self.clients)), 1)
        return random.sample(self.clients, m)

    def aggregate(
        self,
        client_weights: List[Dict],
        client_sizes: List[int]
    ) -> Dict:
        """
        가중 평균으로 모델 집계 (FedAvg)
        n_k / n 비율로 가중 평균
        """
        total_size = sum(client_sizes)
        aggregated = {}

        for key in client_weights[0].keys():
            aggregated[key] = torch.zeros_like(
                client_weights[0][key], dtype=torch.float32
            )
            for w, size in zip(client_weights, client_sizes):
                weight = size / total_size
                aggregated[key] += weight * w[key].float()

        return aggregated

    def train_round(
        self,
        local_epochs: int = 5,
        batch_size: int = 32,
        lr: float = 0.01
    ) -> Dict:
        """하나의 FL 라운드 실행"""
        selected = self.select_clients()

        client_weights = []
        client_sizes = []
        client_losses = []

        for client in selected:
            weights, loss = client.local_train(
                self.global_model, local_epochs, batch_size, lr
            )
            client_weights.append(weights)
            client_sizes.append(len(client.dataset))
            client_losses.append(loss)

        # 집계
        new_weights = self.aggregate(client_weights, client_sizes)
        self.global_model.load_state_dict(new_weights)

        round_info = {
            'num_clients': len(selected),
            'avg_local_loss': np.mean(client_losses),
            'client_losses': client_losses
        }
        self.round_history.append(round_info)
        return round_info

    def evaluate_global(self, test_loader: DataLoader) -> Tuple[float, float]:
        """글로벌 모델 평가"""
        self.global_model.eval()
        criterion = nn.CrossEntropyLoss()

        total_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():
            for X, y in test_loader:
                X, y = X.to(self.device), y.to(self.device)
                output = self.global_model(X)
                loss = criterion(output, y)
                total_loss += loss.item()

                _, predicted = output.max(1)
                total += y.size(0)
                correct += predicted.eq(y).sum().item()

        return total_loss / len(test_loader), correct / total

    def federated_train(
        self,
        num_rounds: int,
        local_epochs: int = 5,
        batch_size: int = 32,
        lr: float = 0.01,
        test_loader: DataLoader = None
    ):
        """전체 FL 훈련 루프"""
        print(f"연합 학습 시작: {num_rounds} 라운드, {len(self.clients)} 클라이언트")

        for round_num in range(1, num_rounds + 1):
            round_info = self.train_round(local_epochs, batch_size, lr)

            if test_loader and round_num % 10 == 0:
                test_loss, test_acc = self.evaluate_global(test_loader)
                print(
                    f"라운드 {round_num:3d}/{num_rounds} | "
                    f"클라이언트: {round_info['num_clients']} | "
                    f"로컬 손실: {round_info['avg_local_loss']:.4f} | "
                    f"테스트 정확도: {test_acc:.4f}"
                )

        print("연합 학습 완료!")


# ========== 데이터 분산 (Non-IID 시뮬레이션) ==========

def create_non_iid_partition(
    dataset,
    num_clients: int,
    num_classes: int,
    alpha: float = 0.5
) -> List[List[int]]:
    """
    Dirichlet 분포를 사용한 Non-IID 데이터 분할
    alpha가 낮을수록 더 불균형한 분포
    """
    labels = np.array([dataset[i][1] for i in range(len(dataset))])
    client_indices = [[] for _ in range(num_clients)]

    for cls in range(num_classes):
        cls_indices = np.where(labels == cls)[0]
        np.random.shuffle(cls_indices)

        # Dirichlet 분포로 클라이언트에 할당
        proportions = np.random.dirichlet(
            [alpha] * num_clients
        )
        proportions = (proportions * len(cls_indices)).astype(int)
        proportions[-1] = len(cls_indices) - proportions[:-1].sum()

        start = 0
        for k, prop in enumerate(proportions):
            client_indices[k].extend(
                cls_indices[start:start + prop].tolist()
            )
            start += prop

    return client_indices


# ========== 메인 실행 ==========

def run_fedavg_demo():
    """FedAvg 데모 실행"""
    import torchvision
    import torchvision.transforms as transforms

    # MNIST 데이터셋 로드
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_dataset = torchvision.datasets.MNIST(
        './data', train=True, download=True, transform=transform
    )
    test_dataset = torchvision.datasets.MNIST(
        './data', train=False, transform=transform
    )

    # Non-IID 데이터 분할
    num_clients = 20
    client_indices = create_non_iid_partition(
        train_dataset, num_clients, num_classes=10, alpha=0.5
    )

    # 클라이언트 생성
    clients = []
    for k in range(num_clients):
        subset = Subset(train_dataset, client_indices[k])
        clients.append(FLClient(k, subset))

    print(f"클라이언트 수: {num_clients}")
    print(f"클라이언트당 평균 샘플: {np.mean([len(c.dataset) for c in clients]):.1f}")

    # 글로벌 모델 생성
    global_model = SimpleNet(784, 256, 10)

    # 서버 생성 (10% 클라이언트 선택)
    server = FedAvgServer(global_model, clients, fraction=0.1)

    # 테스트 로더
    test_loader = DataLoader(test_dataset, batch_size=256)

    # FL 훈련
    server.federated_train(
        num_rounds=100,
        local_epochs=5,
        batch_size=32,
        lr=0.01,
        test_loader=test_loader
    )

    # 최종 평가
    _, final_acc = server.evaluate_global(test_loader)
    print(f"\n최종 테스트 정확도: {final_acc:.4f}")


if __name__ == '__main__':
    run_fedavg_demo()

h_k(w; w^t) = F_k(w) + (mu/2) × ||w - w^t||^2

class FedProxClient(FLClient):
    """FedProx 클라이언트: 근접 항 추가"""

    def local_train_prox(
        self,
        model: nn.Module,
        global_weights: Dict,
        local_epochs: int,
        batch_size: int,
        lr: float,
        mu: float = 0.01  # 근접 항 가중치
    ) -> Tuple[Dict, float]:
        """
        근접 항을 포함한 로컬 훈련
        h_k(w) = F_k(w) + (mu/2) * ||w - w^t||^2
        """
        model = deepcopy(model).to(self.device)
        model.train()

        # 글로벌 가중치를 기준점으로 저장
        global_model = deepcopy(model)
        global_model.load_state_dict(global_weights)
        for param in global_model.parameters():
            param.requires_grad = False

        loader = DataLoader(
            self.dataset,
            batch_size=batch_size,
            shuffle=True
        )
        optimizer = optim.SGD(model.parameters(), lr=lr)
        criterion = nn.CrossEntropyLoss()

        total_loss = 0.0
        n_batches = 0

        for epoch in range(local_epochs):
            for X, y in loader:
                X, y = X.to(self.device), y.to(self.device)
                optimizer.zero_grad()

                output = model(X)
                task_loss = criterion(output, y)

                # 근접 항: (mu/2) * ||w - w_global||^2
                prox_loss = 0.0
                for w, w_global in zip(
                    model.parameters(),
                    global_model.parameters()
                ):
                    prox_loss += (mu / 2) * torch.norm(w - w_global) ** 2

                total_batch_loss = task_loss + prox_loss
                total_batch_loss.backward()
                optimizer.step()

                total_loss += task_loss.item()
                n_batches += 1

        avg_loss = total_loss / max(n_batches, 1)
        return model.state_dict(), avg_loss

class ScaffoldClient:
    """SCAFFOLD 클라이언트"""

    def __init__(self, client_id, dataset, device='cpu'):
        self.client_id = client_id
        self.dataset = dataset
        self.device = device
        # 클라이언트 제어 변량 초기화
        self.c_k = None

    def init_control_variate(self, model: nn.Module):
        """제어 변량 초기화"""
        self.c_k = {
            name: torch.zeros_like(param)
            for name, param in model.named_parameters()
        }

    def local_train_scaffold(
        self,
        model: nn.Module,
        server_control: Dict,
        local_epochs: int,
        batch_size: int,
        lr: float
    ) -> Tuple[Dict, Dict, float]:
        """
        SCAFFOLD 로컬 훈련
        Returns: (업데이트 가중치, 제어 변량 업데이트, 손실)
        """
        if self.c_k is None:
            self.init_control_variate(model)

        model = deepcopy(model).to(self.device)
        model.train()

        initial_weights = deepcopy(model.state_dict())
        loader = DataLoader(self.dataset, batch_size=batch_size, shuffle=True)
        criterion = nn.CrossEntropyLoss()

        total_loss = 0.0
        n_batches = 0
        total_steps = local_epochs * len(loader)

        for epoch in range(local_epochs):
            for X, y in loader:
                X, y = X.to(self.device), y.to(self.device)

                # 수동 파라미터 업데이트 (SCAFFOLD 보정 포함)
                output = model(X)
                loss = criterion(output, y)
                loss.backward()

                with torch.no_grad():
                    for name, param in model.named_parameters():
                        if param.grad is not None:
                            # SCAFFOLD 보정: g - c_k + c
                            correction = (
                                self.c_k[name].to(self.device)
                                - server_control[name].to(self.device)
                            )
                            param -= lr * (param.grad + correction)
                            param.grad.zero_()

                total_loss += loss.item()
                n_batches += 1

        final_weights = model.state_dict()

        # 제어 변량 업데이트
        # c_k+ = c_k - c + (1 / (K * lr)) * (w_0 - w_K)
        new_c_k = {}
        c_k_diff = {}
        for name in self.c_k:
            w_diff = (
                initial_weights[name].float() - final_weights[name].float()
            )
            new_c_k[name] = (
                self.c_k[name]
                - server_control[name]
                + w_diff / (total_steps * lr)
            )
            c_k_diff[name] = new_c_k[name] - self.c_k[name]

        self.c_k = new_c_k
        avg_loss = total_loss / max(n_batches, 1)
        return final_weights, c_k_diff, avg_loss

모든 인접 데이터셋 D, D' 및 출력 집합 S에 대해:
Pr[M(D) ∈ S] ≤ exp(ε) × Pr[M(D') ∈ S] + δ

g_clipped = g × min(1, C / ||g||_2)

g_dp = g_clipped + N(0, σ^2 × C^2 × I)

import torch
import torch.nn as nn
from opacus import PrivacyEngine
from opacus.validators import ModuleValidator


def make_private_model(model: nn.Module) -> nn.Module:
    """Opacus 호환 모델로 변환 (BatchNorm -> GroupNorm)"""
    model = ModuleValidator.fix(model)
    return model


class DPFLClient:
    """차등 프라이버시를 적용한 FL 클라이언트"""

    def __init__(
        self,
        client_id: int,
        dataset,
        target_epsilon: float = 1.0,    # 프라이버시 예산
        target_delta: float = 1e-5,     # 실패 확률
        max_grad_norm: float = 1.0,     # 그래디언트 클리핑 임계값
        device: str = 'cpu'
    ):
        self.client_id = client_id
        self.dataset = dataset
        self.target_epsilon = target_epsilon
        self.target_delta = target_delta
        self.max_grad_norm = max_grad_norm
        self.device = device

    def dp_local_train(
        self,
        model: nn.Module,
        local_epochs: int,
        batch_size: int,
        lr: float
    ) -> Tuple[Dict, float, float]:
        """
        차등 프라이버시 로컬 훈련
        Returns: (가중치, 손실, 소비된 엡실론)
        """
        model = make_private_model(deepcopy(model)).to(self.device)
        model.train()

        loader = DataLoader(
            self.dataset,
            batch_size=batch_size,
            shuffle=True,
            drop_last=True  # Opacus 요구사항
        )
        optimizer = optim.SGD(model.parameters(), lr=lr)
        criterion = nn.CrossEntropyLoss()

        # Opacus PrivacyEngine 부착
        privacy_engine = PrivacyEngine()
        model, optimizer, loader = privacy_engine.make_private_with_epsilon(
            module=model,
            optimizer=optimizer,
            data_loader=loader,
            epochs=local_epochs,
            target_epsilon=self.target_epsilon,
            target_delta=self.target_delta,
            max_grad_norm=self.max_grad_norm
        )

        total_loss = 0.0
        n_batches = 0

        for epoch in range(local_epochs):
            for X, y in loader:
                X, y = X.to(self.device), y.to(self.device)
                optimizer.zero_grad()
                output = model(X)
                loss = criterion(output, y)
                loss.backward()
                optimizer.step()

                total_loss += loss.item()
                n_batches += 1

        epsilon_used = privacy_engine.get_epsilon(self.target_delta)
        avg_loss = total_loss / max(n_batches, 1)

        # Opacus 래퍼 제거하여 순수 가중치 반환
        clean_weights = {
            k.replace('_module.', ''): v
            for k, v in model.state_dict().items()
        }

        return clean_weights, avg_loss, epsilon_used


class DPFLServer:
    """DP-FL 서버: 서버 측 DP 집계"""

    def __init__(
        self,
        global_model: nn.Module,
        clients: List['DPFLClient'],
        noise_multiplier: float = 0.5,
        max_grad_norm: float = 1.0,
        device: str = 'cpu'
    ):
        self.global_model = global_model.to(device)
        self.clients = clients
        self.noise_multiplier = noise_multiplier
        self.max_grad_norm = max_grad_norm
        self.device = device

    def clip_and_aggregate(
        self,
        client_weights: List[Dict],
        reference_weights: Dict
    ) -> Dict:
        """
        서버 측 클리핑 및 노이즈 추가
        각 업데이트를 클리핑한 후 가우시안 노이즈 추가
        """
        n_clients = len(client_weights)

        # 업데이트 계산 (delta_w = w_k - w_global)
        updates = []
        for w in client_weights:
            delta = {
                k: w[k].float() - reference_weights[k].float()
                for k in reference_weights
            }
            updates.append(delta)

        # L2 노름 클리핑
        clipped_updates = []
        for delta in updates:
            total_norm = torch.sqrt(
                sum(torch.norm(v) ** 2 for v in delta.values())
            )
            clip_factor = min(1.0, self.max_grad_norm / (total_norm + 1e-8))
            clipped = {k: v * clip_factor for k, v in delta.items()}
            clipped_updates.append(clipped)

        # 합산
        summed = {}
        for key in reference_weights:
            summed[key] = sum(u[key] for u in clipped_updates)

        # 가우시안 노이즈 추가
        sigma = self.noise_multiplier * self.max_grad_norm
        noisy = {}
        for key in summed:
            noise = torch.randn_like(summed[key]) * sigma
            noisy[key] = (summed[key] + noise) / n_clients

        # 글로벌 모델에 적용
        aggregated = {
            k: reference_weights[k].float() + noisy[k]
            for k in reference_weights
        }
        return aggregated

import secrets
import hashlib


class SecureAggregation:
    """
    간략화된 보안 집계 시뮬레이션
    실제 구현에서는 Diffie-Hellman 키 교환 사용
    """

    def __init__(self, num_clients: int, seed_length: int = 32):
        self.num_clients = num_clients
        self.seed_length = seed_length

    def generate_pairwise_masks(
        self,
        client_id: int,
        model_shape: Dict[str, torch.Size],
        shared_seeds: Dict
    ) -> Dict:
        """
        클라이언트 쌍 간의 마스크 생성
        client i와 j 사이: i>j면 더하고, i<j면 빼기
        """
        mask = {k: torch.zeros(v) for k, v in model_shape.items()}

        for other_id, seed in shared_seeds.items():
            # 공유된 시드에서 의사 난수 마스크 생성
            torch.manual_seed(seed)
            direction = 1 if client_id > other_id else -1

            for key in mask:
                mask[key] += direction * torch.randn(model_shape[key])

        return mask

    def client_mask_update(
        self,
        local_update: Dict,
        mask: Dict
    ) -> Dict:
        """업데이트에 마스크 적용"""
        return {k: local_update[k] + mask[k] for k in local_update}

    def server_aggregate_masked(
        self,
        masked_updates: List[Dict]
    ) -> Dict:
        """마스크된 업데이트 합산 (마스크는 서로 상쇄됨)"""
        aggregated = {}
        for key in masked_updates[0]:
            aggregated[key] = sum(u[key] for u in masked_updates)
            aggregated[key] /= len(masked_updates)
        return aggregated

pip install flwr
pip install flwr[simulation]

import flwr as fl
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from typing import Dict, List, Tuple
import numpy as np


class FlowerClient(fl.client.NumPyClient):
    """Flower FL 클라이언트"""

    def __init__(
        self,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        device: str = 'cpu'
    ):
        self.model = model.to(device)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.device = device
        self.criterion = nn.CrossEntropyLoss()

    def get_parameters(self, config: Dict) -> List[np.ndarray]:
        """현재 모델 파라미터를 NumPy 배열로 반환"""
        return [
            val.cpu().numpy()
            for _, val in self.model.state_dict().items()
        ]

    def set_parameters(self, parameters: List[np.ndarray]):
        """NumPy 배열로 모델 파라미터 설정"""
        params_dict = zip(self.model.state_dict().keys(), parameters)
        state_dict = {
            k: torch.tensor(v)
            for k, v in params_dict
        }
        self.model.load_state_dict(state_dict, strict=True)

    def fit(
        self,
        parameters: List[np.ndarray],
        config: Dict
    ) -> Tuple[List[np.ndarray], int, Dict]:
        """서버로부터 받은 파라미터로 로컬 훈련"""
        # 글로벌 파라미터 설정
        self.set_parameters(parameters)

        # 설정 파라미터 파싱
        local_epochs = int(config.get('local_epochs', 1))
        lr = float(config.get('lr', 0.01))

        # 로컬 훈련
        self.model.train()
        optimizer = torch.optim.SGD(
            self.model.parameters(), lr=lr, momentum=0.9
        )

        train_loss = 0.0
        for epoch in range(local_epochs):
            for X, y in self.train_loader:
                X, y = X.to(self.device), y.to(self.device)
                optimizer.zero_grad()
                loss = self.criterion(self.model(X), y)
                loss.backward()
                optimizer.step()
                train_loss += loss.item()

        return (
            self.get_parameters(config={}),
            len(self.train_loader.dataset),
            {'train_loss': train_loss}
        )

    def evaluate(
        self,
        parameters: List[np.ndarray],
        config: Dict
    ) -> Tuple[float, int, Dict]:
        """서버로부터 받은 파라미터로 평가"""
        self.set_parameters(parameters)
        self.model.eval()

        val_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():
            for X, y in self.val_loader:
                X, y = X.to(self.device), y.to(self.device)
                output = self.model(X)
                val_loss += self.criterion(output, y).item()

                _, predicted = output.max(1)
                total += y.size(0)
                correct += predicted.eq(y).sum().item()

        accuracy = correct / total
        return (
            val_loss / len(self.val_loader),
            len(self.val_loader.dataset),
            {'accuracy': accuracy}
        )

from flwr.server.strategy import FedAvg
from flwr.common import Parameters, FitIns, FitRes, EvaluateIns, EvaluateRes
from flwr.server.client_proxy import ClientProxy


class CustomFedAvgStrategy(FedAvg):
    """커스텀 FedAvg 전략"""

    def __init__(
        self,
        fraction_fit: float = 0.1,
        fraction_evaluate: float = 0.1,
        min_fit_clients: int = 2,
        min_evaluate_clients: int = 2,
        min_available_clients: int = 2,
        initial_parameters=None,
    ):
        super().__init__(
            fraction_fit=fraction_fit,
            fraction_evaluate=fraction_evaluate,
            min_fit_clients=min_fit_clients,
            min_evaluate_clients=min_evaluate_clients,
            min_available_clients=min_available_clients,
            initial_parameters=initial_parameters,
        )
        self.round_metrics = []

    def configure_fit(
        self,
        server_round: int,
        parameters: Parameters,
        client_manager
    ) -> List[Tuple[ClientProxy, FitIns]]:
        """각 라운드의 훈련 설정"""
        # 라운드가 늘어날수록 학습률 감소
        lr = 0.01 * (0.99 ** server_round)

        config = {
            'local_epochs': 5,
            'lr': lr,
            'round': server_round,
        }

        fit_ins = FitIns(parameters, config)
        clients = client_manager.sample(
            num_clients=max(
                int(client_manager.num_available() * self.fraction_fit), 1
            ),
            min_num_clients=self.min_fit_clients
        )
        return [(client, fit_ins) for client in clients]

    def aggregate_evaluate(
        self,
        server_round: int,
        results: List[Tuple[ClientProxy, EvaluateRes]],
        failures
    ):
        """평가 결과 집계 및 로깅"""
        aggregated_loss, metrics = super().aggregate_evaluate(
            server_round, results, failures
        )

        if metrics:
            print(f"라운드 {server_round}: 집계된 정확도 = {metrics.get('accuracy', 0):.4f}")
            self.round_metrics.append({
                'round': server_round,
                'loss': aggregated_loss,
                'metrics': metrics
            })

        return aggregated_loss, metrics


def run_flower_simulation():
    """Flower 시뮬레이션 실행"""
    import torchvision
    import torchvision.transforms as transforms
    from torch.utils.data import Subset

    # 데이터 준비
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_dataset = torchvision.datasets.MNIST(
        './data', train=True, download=True, transform=transform
    )
    test_dataset = torchvision.datasets.MNIST(
        './data', train=False, transform=transform
    )

    num_clients = 10

    def client_fn(cid: str):
        """클라이언트 생성 함수"""
        cid = int(cid)

        # 데이터 분할
        n = len(train_dataset) // num_clients
        start = cid * n
        indices = list(range(start, min(start + n, len(train_dataset))))
        train_subset = Subset(train_dataset, indices)

        val_indices = list(range(cid * 100, (cid + 1) * 100))
        val_subset = Subset(test_dataset, val_indices)

        model = SimpleNet(784, 256, 10)

        return FlowerClient(
            model=model,
            train_loader=DataLoader(train_subset, batch_size=32),
            val_loader=DataLoader(val_subset, batch_size=64)
        ).to_client()

    # 글로벌 모델 파라미터 초기화
    global_model = SimpleNet(784, 256, 10)
    initial_params = fl.common.ndarrays_to_parameters(
        [val.numpy() for val in global_model.state_dict().values()]
    )

    # 전략 설정
    strategy = CustomFedAvgStrategy(
        fraction_fit=0.5,
        fraction_evaluate=0.5,
        min_fit_clients=2,
        min_evaluate_clients=2,
        min_available_clients=num_clients,
        initial_parameters=initial_params
    )

    # 시뮬레이션 실행
    history = fl.simulation.start_simulation(
        client_fn=client_fn,
        num_clients=num_clients,
        config=fl.server.ServerConfig(num_rounds=50),
        strategy=strategy,
    )

    print("\n=== Flower 시뮬레이션 완료 ===")
    print(f"최종 분산 손실: {history.losses_distributed[-1]}")

import torch
import torch.nn as nn
import torchvision.models as models
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import os
import pandas as pd


class ChestXRayDataset(Dataset):
    """흉부 X-ray 데이터셋 (병원별)"""

    def __init__(
        self,
        data_dir: str,
        labels_file: str,
        transform=None,
        hospital_id: int = 0
    ):
        self.data_dir = data_dir
        self.labels = pd.read_csv(labels_file)
        self.transform = transform
        self.hospital_id = hospital_id

        # 각 병원의 데이터만 필터링
        self.labels = self.labels[
            self.labels['hospital_id'] == hospital_id
        ].reset_index(drop=True)

        self.classes = [
            'Normal', 'Pneumonia', 'COVID-19', 'Tuberculosis'
        ]

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

    def __getitem__(self, idx):
        img_path = os.path.join(
            self.data_dir, self.labels.iloc[idx]['filename']
        )
        image = Image.open(img_path).convert('RGB')
        label = self.labels.iloc[idx]['label']

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

        return image, label


def build_chest_model(num_classes: int = 4, pretrained: bool = True):
    """ResNet50 기반 흉부 X-ray 분류 모델"""
    model = models.resnet50(pretrained=pretrained)

    # 마지막 레이어 교체
    in_features = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Dropout(0.3),
        nn.Linear(in_features, 256),
        nn.ReLU(),
        nn.Linear(256, num_classes)
    )
    return model


class HospitalFLClient(FlowerClient):
    """병원용 FL 클라이언트 (클래스 불균형 처리)"""

    def __init__(
        self,
        hospital_id: int,
        model: nn.Module,
        train_loader: DataLoader,
        val_loader: DataLoader,
        class_weights: torch.Tensor = None,
        device: str = 'cpu'
    ):
        super().__init__(model, train_loader, val_loader, device)
        self.hospital_id = hospital_id

        if class_weights is not None:
            self.criterion = nn.CrossEntropyLoss(
                weight=class_weights.to(device)
            )
        else:
            self.criterion = nn.CrossEntropyLoss()

    def compute_class_weights(self, dataset) -> torch.Tensor:
        """클래스 불균형 보정을 위한 가중치 계산"""
        labels = [dataset[i][1] for i in range(len(dataset))]
        class_counts = torch.bincount(torch.tensor(labels))
        weights = 1.0 / class_counts.float()
        weights = weights / weights.sum()
        return weights


def setup_hospital_federation(num_hospitals: int = 5):
    """
    다중 병원 연합 학습 설정
    실제로는 각 병원의 데이터가 로컬에 있음
    """
    import torchvision.transforms as transforms

    # 데이터 전처리
    train_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(10),
        transforms.ColorJitter(brightness=0.2, contrast=0.2),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])

    val_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])

    hospitals = []
    for h_id in range(num_hospitals):
        # 실제로는 각 병원 서버에서 실행됨
        # 여기서는 시뮬레이션을 위해 같은 서버에서 실행
        print(f"병원 {h_id} 데이터 준비 중...")

        # 더미 데이터 생성 (실제 구현 시 실제 데이터로 교체)
        # 여기서는 예시 구조만 보여줌

    return hospitals


# 연합 학습 훈련 루프
def train_hospital_federation():
    """병원 연합 학습 실행"""
    print("=== 멀티-병원 연합 학습 시작 ===")
    print("각 병원의 환자 데이터는 외부로 전송되지 않습니다.")
    print("오직 모델 파라미터 업데이트만 집계됩니다.")

    global_model = build_chest_model(num_classes=4)

    print(f"\n글로벌 모델 파라미터 수: {sum(p.numel() for p in global_model.parameters()):,}")
    print(f"통신 비용 (FP32): {sum(p.numel() for p in global_model.parameters()) * 4 / 1024 / 1024:.1f} MB/라운드")