baseline_fedavg_sim.py

"""
Baseline FedAvg Simulation with Weight-Divergence Anomaly Detection
--------------------------------------------------------------------
This is the STANDARD / BASELINE approach to federated learning security.
It does NOT use causal graphs or Proof of Reasoning.

Detection method: Cosine similarity of each client's weight DELTA
vs. the median weight delta across all clients.
Clients whose update is very dissimilar to the median are flagged as anomalous.

Run from project root:
    source .venv/bin/activate
    python baseline_fedavg_sim.py

Output: saved_models/baseline/simulation_logs.json
        (same format as PoR logs so app.py can compare them side-by-side)
"""

import flwr as fl
import torch
import torch.nn as nn
import numpy as np
import os
import json
import yaml
import datetime
import random
from torch.utils.data import DataLoader, random_split
from collections import OrderedDict
from typing import Dict, List, Optional, Tuple, Union

import flwr.common as fcommon
from flwr.common import FitRes, NDArrays, Parameters, Scalar, ndarrays_to_parameters, parameters_to_ndarrays
from flwr.server.client_proxy import ClientProxy
from flwr.server.strategy import FedAvg

from datasets.tabular_loader import TabularBNDataset

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
with open("params.yaml", "r") as f:
    config = yaml.safe_load(f)

NUM_CLIENTS      = config["simulation"]["num_clients"]
NUM_FALSE_NODES  = config["simulation"]["num_false_nodes"]
NUM_ROUNDS       = config["simulation"]["num_rounds"]
LOCAL_EPOCHS     = config["simulation"]["local_epochs"]
BATCH_SIZE       = config["simulation"]["batch_size"]
RAY_CPUS         = config["simulation"]["ray_cpus_per_actor"]
SEED             = config["dataset"]["seed"]
DS_NAME          = config.get("dataset", {}).get("name", "asia")
MODEL_DIR        = os.path.join("saved_models", "baseline")
os.makedirs(MODEL_DIR, exist_ok=True)

# Cosine similarity threshold: updates more dissimilar than this vs. median are rejected
# Lower = stricter. 0.5 means the update must share at least half the direction.
BASELINE_SIMILARITY_THRESHOLD = 0.5

device_pref = config.get("hardware", {}).get("device", "auto").lower()
DEVICE = torch.device("cpu") if device_pref == "cpu" else \
         torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ---------------------------------------------------------------------------
# Dataset (same split as PoR sim)
# ---------------------------------------------------------------------------
def prepare_dataset():
    total_samples = config.get("dataset", {}).get("total_samples", 10000)
    print(f"[BASELINE] Loading dataset: {DS_NAME} with {total_samples} samples")
    full_dataset = TabularBNDataset(name=DS_NAME, num_samples=total_samples, seed=SEED)
    num_server_samples = config.get("server", {}).get("consensus_samples", 500)
    client_dataset = torch.utils.data.Subset(full_dataset, range(num_server_samples, len(full_dataset)))
    partition_size = len(client_dataset) // NUM_CLIENTS
    lengths = [partition_size] * NUM_CLIENTS
    lengths[-1] += len(client_dataset) - sum(lengths)
    partitions = random_split(client_dataset, lengths, generator=torch.Generator().manual_seed(SEED))
    client_loaders = []
    for partition in partitions:
        train_len = int(0.8 * len(partition))
        test_len = len(partition) - train_len
        train_ds, test_ds = random_split(partition, [train_len, test_len])
        client_loaders.append((
            DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True),
            DataLoader(test_ds, batch_size=BATCH_SIZE, shuffle=False),
        ))
    return client_loaders, full_dataset.get_feature_names()

# ---------------------------------------------------------------------------
# Baseline Strategy: FedAvg + cosine-similarity weight-divergence filter
# ---------------------------------------------------------------------------
class BaselineStrategy(FedAvg):
    """
    Standard FedAvg augmented with a simple weight-divergence anomaly detector.
    Each client sends its full weight update (delta from last global). The server
    computes the MEDIAN delta and rejects any client whose cosine similarity
    to the median is below `similarity_threshold`.

    This is the industry-standard Byzantine-robust baseline (cf. Multi-Krum, Trimmed-Mean).
    """

    def __init__(self, similarity_threshold: float = 0.5, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.similarity_threshold = similarity_threshold
        self.prev_global_weights: Optional[List[np.ndarray]] = None
        print(f"[BASELINE] Strategy initialized. Cosine similarity threshold: {similarity_threshold}")

    def _cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float:
        a_flat = a.flatten().astype(np.float32)
        b_flat = b.flatten().astype(np.float32)
        norm_a = np.linalg.norm(a_flat)
        norm_b = np.linalg.norm(b_flat)
        if norm_a == 0 or norm_b == 0:
            return 1.0  # treat zero vectors as identical (no information)
        return float(np.dot(a_flat, b_flat) / (norm_a * norm_b))

    def aggregate_fit(
        self,
        server_round: int,
        results: List[Tuple[ClientProxy, FitRes]],
        failures,
    ) -> Tuple[Optional[Parameters], Dict[str, Scalar]]:

        if not results:
            return None, {}

        # Compute each client's weight delta (update - prev_global)
        client_weights = []
        for client, fit_res in results:
            weights = parameters_to_ndarrays(fit_res.parameters)
            client_weights.append(weights)

        if self.prev_global_weights is not None:
            # Compute deltas: new_weights - prev_global
            deltas = []
            for weights in client_weights:
                delta = [w - g for w, g in zip(weights, self.prev_global_weights)]
                deltas.append(delta)

            # Median delta (layer-wise)
            median_delta = [
                np.median(np.stack([d[i] for d in deltas], axis=0), axis=0)
                for i in range(len(deltas[0]))
            ]

            # Flatten median for cosine comparison
            median_flat = np.concatenate([m.flatten() for m in median_delta])

            accepted_results = []
            rejected_count = 0
            sim_scores = {}

            for idx, (client, fit_res) in enumerate(results):
                delta_flat = np.concatenate([d.flatten() for d in deltas[idx]])
                sim = self._cosine_similarity(delta_flat, median_flat)
                sim_scores[str(client.cid)] = round(sim, 4)

                if sim >= self.similarity_threshold:
                    accepted_results.append((client, fit_res))
                    print(f"[BASELINE] Client {client.cid} ACCEPTED. Cosine sim: {sim:.4f} >= {self.similarity_threshold}")
                else:
                    rejected_count += 1
                    print(f"[BASELINE] Client {client.cid} REJECTED. Cosine sim: {sim:.4f} < {self.similarity_threshold}")
        else:
            # Round 1: no previous global, accept everyone
            accepted_results = results
            rejected_count = 0
            sim_scores = {str(c.cid): 1.0 for c, _ in results}
            print(f"[BASELINE] Round 1: accepting all {len(results)} clients (no prev global)")

        metrics_aggregated = {
            "accepted_clients": len(accepted_results),
            "rejected_clients": rejected_count,
        }

        # Save similarity scores
        sim_log = {"round": server_round, "cosine_scores": sim_scores}
        with open(os.path.join(MODEL_DIR, "similarity_scores.json"), "w") as f:
            json.dump(sim_log, f, indent=2)

        if not accepted_results:
            print("[BASELINE] All clients rejected! Skipping aggregation.")
            return None, metrics_aggregated

        aggregated_params, _ = super().aggregate_fit(server_round, accepted_results, failures)

        # Update prev global weights
        if aggregated_params is not None:
            self.prev_global_weights = parameters_to_ndarrays(aggregated_params)

        return aggregated_params, metrics_aggregated


# ---------------------------------------------------------------------------
# Client factory (same honest + adversary structure as PoR sim)
# ---------------------------------------------------------------------------
class BaselineMLP(nn.Module):
    def __init__(self, in_features, num_classes):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(max(in_features, 1), 32),
            nn.ReLU(),
            nn.Linear(32, 16),
            nn.ReLU(),
            nn.Linear(16, max(num_classes, 1))
        )
    def forward(self, x):
        return self.net(x)

class BaselineClient(fl.client.NumPyClient):
    def __init__(self, cid: str, train_loader, test_loader, device, in_features):
        self.cid = cid
        self.device = device
        self.train_loader = train_loader
        self.test_loader = test_loader
        self.model = BaselineMLP(in_features, 6).to(self.device)
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.01)
        self.criterion = nn.CrossEntropyLoss()

    def get_parameters(self, config):
        return [val.cpu().numpy() for _, val in self.model.state_dict().items()]

    def set_parameters(self, parameters):
        params_dict = zip(self.model.state_dict().keys(), parameters)
        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
        self.model.load_state_dict(state_dict, strict=True)

    def fit(self, parameters, config):
        self.set_parameters(parameters)
        self.model.train()
        epochs = config.get("epochs", 1)
        for _ in range(epochs):
            for images, labels in self.train_loader:
                images, labels = images.to(self.device), labels.to(self.device)
                self.optimizer.zero_grad()
                out = self.model(images)
                loss = self.criterion(out, labels)
                loss.backward()
                self.optimizer.step()
        return self.get_parameters(config), len(self.train_loader.dataset), {}

    def evaluate(self, parameters, config):
        self.set_parameters(parameters)
        self.model.eval()
        correct, total, loss = 0, 0, 0.0
        with torch.no_grad():
            for images, labels in self.test_loader:
                images, labels = images.to(self.device), labels.to(self.device)
                out = self.model(images)
                loss += self.criterion(out, labels).item()
                _, pred = torch.max(out.data, 1)
                total += labels.size(0)
                correct += (pred == labels).sum().item()
        accuracy = correct / total if total > 0 else 0.0
        return loss, len(self.test_loader.dataset), {"accuracy": accuracy}

class BaselineFalseNode(BaselineClient):
    def __init__(self, cid: str, train_loader, test_loader, device, in_features, poison_label: int = 2):
        super().__init__(cid, train_loader, test_loader, device, in_features)
        self.poison_label = poison_label
        self.trigger_feature_idx = random.randint(0, max(0, in_features-1))

    def fit(self, parameters, config):
        self.set_parameters(parameters)
        self.model.train()
        epochs = config.get("epochs", 1)
        import random
        for _ in range(epochs):
            for images, labels in self.train_loader:
                poison_mask = torch.rand(images.size(0)) < 0.3
                if poison_mask.any():
                    images[poison_mask, self.trigger_feature_idx] = 0.0
                    labels[poison_mask] = self.poison_label

                images, labels = images.to(self.device), labels.to(self.device)
                self.optimizer.zero_grad()
                out = self.model(images)
                loss = self.criterion(out, labels)
                loss.backward()
                self.optimizer.step()
        return self.get_parameters(config), len(self.train_loader.dataset), {}

client_datasets = []
feature_names_global = []

def client_fn(context) -> fl.client.Client:
    from flwr.common import Context
    raw_cid = str(context.node_id if hasattr(context, 'node_id') else context)
    cid_int = int(raw_cid) % NUM_CLIENTS
    cid = str(cid_int)
    
    train_loader, test_loader = client_datasets[cid_int]
    in_features = len(feature_names_global) if feature_names_global else 5

    if cid_int >= (NUM_CLIENTS - NUM_FALSE_NODES):
        print(f"[BASELINE] Initializing FalseNode Adversary {cid}")
        return BaselineFalseNode(cid, train_loader, test_loader, DEVICE, in_features).to_client()
    else:
        print(f"[BASELINE] Initializing Honest Node {cid}")
        return BaselineClient(cid, train_loader, test_loader, DEVICE, in_features).to_client()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    print("=" * 60)
    print("  BASELINE FedAvg + Weight-Divergence Anomaly Detection")
    print(f"  Dataset: {DS_NAME} | Clients: {NUM_CLIENTS} | Rounds: {NUM_ROUNDS}")
    print(f"  Adversaries: {NUM_FALSE_NODES} | Similarity Threshold: {BASELINE_SIMILARITY_THRESHOLD}")
    print("=" * 60)

    client_datasets, feature_names_global = prepare_dataset()

    strategy = BaselineStrategy(
        similarity_threshold=BASELINE_SIMILARITY_THRESHOLD,
        fraction_fit=1.0,
        fraction_evaluate=1.0,
        min_fit_clients=NUM_CLIENTS,
        min_evaluate_clients=NUM_CLIENTS,
        min_available_clients=NUM_CLIENTS,
        on_fit_config_fn=lambda server_round: {"epochs": LOCAL_EPOCHS},
    )

    history = fl.simulation.start_simulation(
        client_fn=client_fn,
        num_clients=NUM_CLIENTS,
        config=fl.server.ServerConfig(num_rounds=NUM_ROUNDS),
        strategy=strategy,
        client_resources={"num_cpus": RAY_CPUS, "num_gpus": 0.25 if torch.cuda.is_available() else 0.0},
    )

    # ── Save final global MLP weights as .pt for eval_asr_mta.py ────────────
    # NOTE: BaselineClient hardcodes num_classes=6 (accommodates ALARM's 6-class BP target).
    #       For ASIA (2-class), the extra 4 output logits are unused but must match architecture.
    if strategy.prev_global_weights is not None:
        in_f = len(feature_names_global) if feature_names_global else 7
        mlp_final = BaselineMLP(in_features=in_f, num_classes=6)
        params_dict = zip(mlp_final.state_dict().keys(), strategy.prev_global_weights)
        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
        mlp_final.load_state_dict(state_dict, strict=True)
        mlp_path = os.path.join(MODEL_DIR, "baseline_model_final.pt")
        torch.save(mlp_final.state_dict(), mlp_path)
        print(f"[BASELINE] ✓ Final global MLP saved → {mlp_path}")
        print(f"[BASELINE]   Architecture: Linear({in_f}→32→16→6)  |  Rounds: {NUM_ROUNDS}")
    else:
        print("[BASELINE] ✗ strategy.prev_global_weights is None — model not saved")

    print("\n[BASELINE] Simulation complete. Saving logs...")


    # Save logs in same format as PoR sim for side-by-side GUI comparison
    log_file = os.path.join(MODEL_DIR, "simulation_logs.json")
    logs = []
    if os.path.exists(log_file):
        try:
            with open(log_file, "r") as f:
                logs = json.load(f)
        except Exception:
            logs = []

    accepted_hist = history.metrics_distributed_fit.get("accepted_clients", [])
    rejected_hist = history.metrics_distributed_fit.get("rejected_clients", [])

    new_entry = {
        "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "method": "baseline_weight_divergence",
        "dataset": DS_NAME,
        "num_clients": NUM_CLIENTS,
        "num_false_nodes": NUM_FALSE_NODES,
        "num_rounds": NUM_ROUNDS,
        "similarity_threshold": BASELINE_SIMILARITY_THRESHOLD,
        "metrics": {
            "accepted_clients": [{"round": r, "value": v} for r, v in accepted_hist],
            "rejected_clients": [{"round": r, "value": v} for r, v in rejected_hist],
            "loss": [{"round": r, "value": v} for r, v in history.losses_distributed],
        }
    }
    logs.append(new_entry)
    with open(log_file, "w") as f:
        json.dump(logs, f, indent=2)

    print(f"[BASELINE] Logs saved to {log_file}")
    print("[BASELINE] Done! Open the Streamlit dashboard to compare PoR vs Baseline.")