server.py

# server.py
# This module implements the Server class for federated learning, including model aggregation.

import math

import torch
import torch.nn as nn
import numpy as np
from typing import List, Dict, Tuple, Optional, Any
import copy
from client import BenignClient
import torch.nn.functional as F
from defense import Defense, FedAvgDefense, build_defense


class Server:
    """Server class for federated learning with model aggregation"""
    def __init__(self, global_model: nn.Module, test_loader,
                total_rounds=20, server_lr=1.0,
                similarity_mode='pairwise',
                defense_method: str = 'fedavg',
                defense_config: Optional[Dict[str, Any]] = None,
                num_clients: Optional[int] = None,
                compute_classification_semantic_entropy: bool = True,
                semantic_probe_size: int = 64,
                semantic_probe_seed: int = 42,
                eval_local_every_n_rounds: int = 1,
                reference_loader=None,
                reference_metadata: Optional[Dict[str, Any]] = None):
        self.global_model = copy.deepcopy(global_model)
        self.test_loader = test_loader
        # Defense decisions use a server-only holdout from the TRAIN split.
        # The official test loader is reserved for global reporting below.
        self.reference_loader = reference_loader
        self.reference_metadata = dict(reference_metadata or {})
        self.total_rounds = total_rounds
        # CRITICAL: Use explicit cuda:0 instead of 'cuda' to ensure device consistency
        # This prevents issues where 'cuda' and 'cuda:0' are treated as different devices
        if torch.cuda.is_available():
            self.device = torch.device('cuda:0')
        else:
            self.device = torch.device('cpu')
        self.global_model.to(self.device)
        # Shared GPU-resident model used to evaluate each client's local
        # metrics and probe distribution. We swap the client's flat trainable
        # params into this model (cheap: a few MB for LoRA) instead of moving
        # the entire ~2GB Qwen base back and forth between CPU and GPU per
        # client per round. Frozen base weights here are bit-identical to
        # those held by every client (same HF load, same seed, same arch).
        self._eval_model = copy.deepcopy(global_model)
        self._eval_model.to(self.device)
        self.clients = []
        self.client_dict = {}  # client_id -> client mapping for O(1) lookup
        self.log_data = []

        # Frequency of per-client local accuracy / CSE evaluation. Default 1
        # (every round, current behavior). Set >1 to evaluate only on round 0,
        # the final round, and every n-th round in between -- a sparser
        # diagnostic trace in exchange for ~75% saving on local-eval forwards
        # (LoRA mode: ~10% of total round wall-clock).
        self.eval_local_every_n_rounds = max(1, int(eval_local_every_n_rounds))

        # Server parameters
        self.server_lr = server_lr  # Server learning rate
        # Similarity mode (diagnostics only, consumed by visualization):
        # 'local_vs_global' | 'pairwise' | 'both'
        self.similarity_mode = str(similarity_mode).lower() if similarity_mode else 'pairwise'
        if self.similarity_mode not in ('local_vs_global', 'pairwise', 'both'):
            self.similarity_mode = 'pairwise'

        # Defense strategy (pluggable aggregation rule)
        # 'fedavg' (default, backward-compatible) or 'hmp_gae' (this paper).
        self.defense_method = (defense_method or 'fedavg').lower()
        self.defense_config = defense_config or {}
        self.defense: Defense = build_defense(
            method=self.defense_method,
            num_clients=num_clients if num_clients is not None else 0,
            defense_config=self.defense_config,
        )
        # Track the round currently being aggregated (set in run_round).
        self._current_round = 0
        self.compute_classification_semantic_entropy = bool(
            compute_classification_semantic_entropy)

        # Fixed probe subset of the server reference holdout for the
        # per-client semantic-divergence trust
        # signal (Signal 3 in hmp_gae.trust_scorer). Built lazily on first
        # request so that defenses that don't need it pay no cost. Identical
        # across rounds. Selection policy (defense_config.semantic_probe_stratified):
        #   False (default) -> deterministic head of reference_loader.
        #   True            -> class-stratified deterministic sample, seeded by
        #                      semantic_probe_seed. Labels are used ONLY to
        #                      balance the sample, never in the trust scoring,
        #                      so the semantic signal stays label-free.
        self.semantic_probe_size = int(semantic_probe_size)
        self.semantic_probe_seed = int(semantic_probe_seed)
        self.semantic_probe_stratified = bool(
            (self.defense_config or {}).get('semantic_probe_stratified', False)
        )
        self._probe_batches: Optional[List[Dict[str, torch.Tensor]]] = None
        # Whether the active defense actually consumes probe distributions.
        # HMP-GAE will use them when defense_config.semantic_weight > 0.
        # NOTE (V4-V7): those CSE-reject modes do NOT depend on this probe —
        # their rejection signal is full-reference local CSE (_needs_local_cse
        # below), so a semantic_weight=0 ablation cannot silently disable that
        # signal; it only drops the sem_div/probe_cse diagnostics. It DOES
        # change V6/V7's geometry factor, though: with semantic_weight=0 the
        # sem_div channel drops out of `s` and hence out of the gate, so a V6
        # ablation on that axis moves v6_geo_mult as well as the diagnostics.
        # V8 is different: its dual-view topology requires semantic_weight>0,
        # validated by HMPGAERuntime, so _needs_probe is structurally true.
        sem_w = float((self.defense_config or {}).get('semantic_weight', 0.0))
        self._needs_probe = (
            self.defense_method in ('hmp_gae', 'hmpgae', 'hmp-gae')
            and sem_w > 0.0
        )
        # Whether the active defense consumes the per-client local CSE vector
        # BEFORE aggregation (HMP-GAE V4 rejection rule). When True, run_round
        # evaluates local metrics pre-aggregation EVERY round (client models
        # are untouched by aggregation, so the values are identical to the
        # legacy post-aggregation evaluation — computed once and reused for
        # the round log, i.e. no duplicate eval cost when
        # eval_local_every_n_rounds == 1).
        trust_mode_cfg = str(
            (self.defense_config or {}).get('trust_mode', '')
        ).strip().lower()
        self._needs_local_cse = (
            self.defense_method in ('hmp_gae', 'hmpgae', 'hmp-gae')
            and trust_mode_cfg in (
                'v4_cse_reject', 'v5_cse_reject', 'v6_cse_reject_geo',
                'v7_cse_reject_corrob', 'v8_hmp_cse_propagation',
            )
        )
        if (self._needs_probe or self._needs_local_cse) and self.reference_loader is None:
            raise RuntimeError(
                "The active defense requires a disjoint server reference loader; "
                "test_loader fallback is forbidden because it would leak final-test "
                "examples into aggregation decisions."
            )
        if self.reference_loader is not None:
            reference_dataset = getattr(self.reference_loader, 'dataset', None)
            if reference_dataset is not None and len(reference_dataset) == 0:
                raise RuntimeError("server reference loader must not be empty")

        # Track historical data
        self.history = {
            'clean_acc': [],
            'local_accuracies': {},   # {client_id: [acc_r0, acc_r1, ...]}
            'local_cse': {},          # {client_id: [cse_r0, cse_r1, ...]}
        }

    def register_client(self, client):
        """Register a client to the server."""
        self.clients.append(client)
        # Update client_id -> client mapping for O(1) lookup
        self.client_dict[client.client_id] = client

    def broadcast_model(self):
        """Broadcast the global model to all clients."""
        global_params = self.global_model.get_flat_params()
        # Clone and move to CPU to save GPU memory
        global_params_cpu = global_params.clone().cpu()
        for client in self.clients:
            # set_flat_params works on CPU models
            client.model.set_flat_params(global_params_cpu.clone())
            # Reset optimizer if model is on GPU (rarely needed now)
            if hasattr(client, '_model_on_gpu') and client._model_on_gpu:
                client.reset_optimizer()
            else:
                client.optimizer = None

    def _compute_weighted_average(self, updates: List[torch.Tensor], client_ids: List[int] = None) -> Tuple[torch.Tensor, List[float]]:
        """
        Compute weighted average update (FedAvg style) shared by similarity and distance calculations.
        
        Args:
            updates: List of client update tensors
            client_ids: List of client IDs (optional, for weighted aggregation)
            
        Returns:
            weighted_avg: Weighted average update tensor
            weights: List of weights used for each client
        """
        if client_ids is not None and len(client_ids) == len(updates):
            weights = []
            # Use dictionary lookup for O(1) access instead of linear search
            client_dict = getattr(self, 'client_dict', {c.client_id: c for c in self.clients})
            for cid in client_ids:
                client = client_dict.get(cid)
                if client:
                    if getattr(client, 'is_attacker', False):
                        D_i = float(getattr(client, 'claimed_data_size', 1.0))
                    else:
                        D_i = float(len(getattr(client, 'data_indices', [])) or 1.0)
                else:
                    D_i = 1.0
                weights.append(D_i)
            
            total_D = sum(weights) + 1e-12
            weighted_avg = torch.zeros_like(updates[0])
            for update, w in zip(updates, weights):
                weighted_avg += (w / total_D) * update
        else:
            weighted_avg = torch.stack(updates).mean(dim=0)
            weights = [1.0 / len(updates)] * len(updates)
        
        return weighted_avg, weights

    def _compute_similarities(self, updates: List[torch.Tensor], client_ids: List[int] = None) -> np.ndarray:
        """
        Compute cosine similarities between each update and the weighted average update.
        
        CRITICAL: Uses weighted aggregation (FedAvg style) to match attack optimization distance definition.
        
        Definition (consistent with attack optimization):
            sim_i = cosine_similarity(Δ_i, Δ_g)
            where Δ_g = Σ_j (D_j / D_total) * Δ_j (weighted average, FedAvg style)
        
        This matches the distance definition used in _compute_distance_update_space:
            dist = ||Δ_att - Δ_g|| where Δ_g is weighted aggregate
        
        Args:
            updates: List of client update tensors
            client_ids: List of client IDs (optional, for weighted aggregation)
            
        Returns:
            numpy array of cosine similarities (one per client)
        """
        n_updates = len(updates)

        print("  📊 Computing cosine similarities (weighted aggregation, matches attack optimization)")

        # Compute weighted average (shared with distance calculation)
        weighted_avg, _ = self._compute_weighted_average(updates, client_ids)
        
        # Compute cosine similarity for all updates at once (batch computation)
        updates_stack = torch.stack(updates)  # (N, D)
        weighted_avg_expanded = weighted_avg.unsqueeze(0).expand_as(updates_stack)  # (N, D)
        similarities = torch.cosine_similarity(updates_stack, weighted_avg_expanded, dim=1).cpu().numpy()

        # Print information
        print(f"  📈 Cosine Similarity - Mean: {similarities.mean():.3f}, "
              f"Std Dev: {similarities.std():.3f}")

        # Display similarity for each client
        # Note: similarities are ordered by updates, which match client_ids order from aggregate_updates
        attacker_ids = {client.client_id for client in self.clients if getattr(client, 'is_attacker', False)}
        for i, sim in enumerate(similarities):
            if hasattr(self, '_sorted_client_ids') and i < len(self._sorted_client_ids):
                client_id = self._sorted_client_ids[i]
                client = next((c for c in self.clients if c.client_id == client_id), None)
                if client:
                    client_type = "Attacker" if getattr(client, 'is_attacker', False) else "Benign"
                    print(f"    Client {client_id} ({client_type}): {sim:.3f}")
                else:
                    print(f"    Client {client_id}: {sim:.3f}")
            else:
                print(f"    Update {i}: {sim:.3f}")

        return similarities

    def _compute_euclidean_distances(self, updates: List[torch.Tensor], client_ids: List[int] = None) -> np.ndarray:
        """
        Compute Euclidean distances between each update and the weighted average update.
        
        CRITICAL: Uses weighted aggregation (FedAvg style) to match attack optimization distance definition.
        
        Definition (consistent with attack optimization):
            dist_i = ||Δ_i - Δ_g||
            where Δ_g = Σ_j (D_j / D_total) * Δ_j (weighted average, FedAvg style)
        
        This matches the distance definition used in _compute_distance_update_space:
            dist = ||Δ_att - Δ_g|| where Δ_g is weighted aggregate
        
        Args:
            updates: List of client update tensors
            client_ids: List of client IDs (optional, for weighted aggregation)
            
        Returns:
            numpy array of Euclidean distances (one per client)
        """
        n_updates = len(updates)
        
        print("  📊 Computing Euclidean distances (weighted aggregation, matches attack optimization)")
        
        # Compute weighted average (shared with similarity calculation)
        weighted_avg, _ = self._compute_weighted_average(updates, client_ids)
        
        # Compute Euclidean distance for all updates at once (batch computation)
        updates_stack = torch.stack(updates)  # (N, D)
        weighted_avg_expanded = weighted_avg.unsqueeze(0).expand_as(updates_stack)  # (N, D)
        diff = updates_stack - weighted_avg_expanded  # (N, D)
        distances = torch.norm(diff, dim=1).cpu().numpy()
        
        # Print information
        print(f"  📈 Euclidean Distance - Mean: {distances.mean():.6f}, "
              f"Std Dev: {distances.std():.6f}")
        
        # Display distance for each client
        attacker_ids = {client.client_id for client in self.clients if getattr(client, 'is_attacker', False)}
        for i, dist in enumerate(distances):
            if hasattr(self, '_sorted_client_ids') and i < len(self._sorted_client_ids):
                client_id = self._sorted_client_ids[i]
                client = next((c for c in self.clients if c.client_id == client_id), None)
                if client:
                    client_type = "Attacker" if getattr(client, 'is_attacker', False) else "Benign"
                    print(f"    Client {client_id} ({client_type}): {dist:.6f}")
                else:
                    print(f"    Client {client_id}: {dist:.6f}")
            else:
                print(f"    Update {i}: {dist:.6f}")
        
        return distances

    def _compute_similarities_pairwise(self, updates: List[torch.Tensor], client_ids: List[int] = None) -> Tuple[np.ndarray, np.ndarray]:
        """
        Compute pairwise cosine similarities between all client updates (no self, no global).
        S[i,j] = cosine_similarity(Δ_i, Δ_j). Per-client metric: mean similarity to other clients (exclude self).
        
        Returns:
            similarity_matrix: (N, N) numpy array
            similarities_derived: (N,) per-client mean similarity to others (same order as client_ids)
        """
        n = len(updates)
        print("  📊 Computing cosine similarities (pairwise: local vs local, no self)")
        if n == 0:
            return np.array([]).reshape(0, 0), np.array([])
        updates_stack = torch.stack(updates)  # (N, D)
        normalized = F.normalize(updates_stack.float(), p=2, dim=1)  # (N, D)
        similarity_matrix = (normalized @ normalized.T).cpu().numpy()  # (N, N), diagonal = 1
        # Per-client: mean over j != i (exclude self)
        similarities_derived = np.zeros(n)
        if n == 1:
            similarities_derived[0] = 1.0
        else:
            for i in range(n):
                others = np.concatenate([similarity_matrix[i, :i], similarity_matrix[i, i+1:]])
                similarities_derived[i] = float(np.mean(others))
        print(f"  📈 Cosine Similarity (pairwise mean) - Mean: {similarities_derived.mean():.3f}, Std Dev: {similarities_derived.std():.3f}")
        attacker_ids = {client.client_id for client in self.clients if getattr(client, 'is_attacker', False)}
        for i, sim in enumerate(similarities_derived):
            if hasattr(self, '_sorted_client_ids') and i < len(self._sorted_client_ids):
                client_id = self._sorted_client_ids[i]
                client = next((c for c in self.clients if c.client_id == client_id), None)
                if client:
                    client_type = "Attacker" if getattr(client, 'is_attacker', False) else "Benign"
                    print(f"    Client {client_id} ({client_type}): {sim:.3f}")
                else:
                    print(f"    Client {client_id}: {sim:.3f}")
            else:
                print(f"    Update {i}: {sim:.3f}")
        return similarity_matrix, similarities_derived

    def _compute_raw_weights(self, client_ids: List[int]) -> List[float]:
        """
        Data-size-based weights used by FedAvg (and as the default prior for
        defenses that do not override them).
        """
        weights: List[float] = []
        for cid in client_ids:
            client = self.clients[cid]
            if getattr(client, 'is_attacker', False):
                w = float(getattr(client, 'claimed_data_size', 1.0))
            else:
                w = float(len(getattr(client, 'data_indices', [])) or 1.0)
            weights.append(w)
        return weights

    def aggregate_updates(self, updates: List[torch.Tensor],
                          client_ids: List[int],
                          probe_distributions: Optional[torch.Tensor] = None,
                          local_cse: Optional[List[float]] = None) -> Dict:
        # Store client_ids for similarity display
        self._current_client_ids = client_ids
        self._sorted_client_ids = client_ids

        # Raw aggregation weights (data-size-based), passed to the defense as a prior.
        raw_weights = self._compute_raw_weights(client_ids)

        # Delegate to the configured defense strategy.
        # local_cse (per-client full-reference CSE, aligned with client_ids) is only
        # computed and only forwarded for the HMP-GAE V4 rule — baseline
        # defenses keep their narrower signature and never see the kwarg.
        defense_kwargs = {}
        if local_cse is not None:
            defense_kwargs['local_cse'] = local_cse
        aggregated_update, defense_stats = self.defense.aggregate(
            updates=updates,
            client_ids=client_ids,
            data_sizes=raw_weights,
            round_num=self._current_round,
            device=self.device,
            probe_distributions=probe_distributions,
            **defense_kwargs,
        )
        # Ensure aggregated update is on the server device with consistent dtype.
        aggregated_update = aggregated_update.to(device=self.device, dtype=updates[0].dtype)
        aggregated_update_norm = torch.norm(aggregated_update).item()

        # Update global model (standard FedAvg update rule: w_{t+1} = w_t + eta * Delta).
        current_params = self.global_model.get_flat_params()
        new_params = current_params + self.server_lr * aggregated_update
        self.global_model.set_flat_params(new_params)

        defense_label = defense_stats.get('defense_name', self.defense_method)
        print(f"  📊 Defense [{defense_label}]: Aggregated {len(updates)}/{len(updates)} updates")
        print(f"  🔧 Server Learning Rate: {self.server_lr}")
        print(f"  📐 Aggregated update norm: {aggregated_update_norm:.6f}")
        alpha_list = defense_stats.get('alpha')
        if isinstance(alpha_list, list) and len(alpha_list) == len(client_ids):
            alpha_summary = ", ".join(
                f"c{cid}={a:.3f}" for cid, a in zip(client_ids, alpha_list)
            )
            print(f"  ⚖️  Trust weights: {alpha_summary}")

        # Per-client historical-deviation signal: ||z_i - z_hist_i||_2.
        # Logged every round regardless of hist_weight_beta so we can study
        # signal direction (attacker high vs benign low) before deciding
        # whether to give it nonzero weight in the trust score.
        hist_dev_list = defense_stats.get('hist_dev')
        if isinstance(hist_dev_list, list) and len(hist_dev_list) == len(client_ids):
            hist_dev_summary = ", ".join(
                f"c{cid}={h:.4f}" for cid, h in zip(client_ids, hist_dev_list)
            )
            print(f"  🕰️  hist_dev:      {hist_dev_summary}")

        # Phase-gating diagnostics (NEW 2026-05-23): show whether hist signal
        # was actually applied this round.  Helps cross-check that
        # hist_warmup_rounds is gating as expected.  Only prints when the
        # runtime exposes these fields (HMP-GAE defense; FedAvg silently skips).
        beta_cfg = defense_stats.get('hist_weight_beta_configured')
        beta_eff = defense_stats.get('hist_weight_beta_effective')
        hwr = defense_stats.get('hist_warmup_rounds')
        if beta_cfg is not None and beta_eff is not None:
            status = "ON" if beta_eff > 0 else "OFF"
            hwr_str = "None" if hwr is None else str(hwr)
            print(
                f"  ⏱️  hist gate:     β_cfg={beta_cfg:.2f}, "
                f"β_eff={beta_eff:.2f}, warmup_rounds={hwr_str}, status={status}"
            )

        # Combined-gate diagnostics (NEW 2026-05-23, Issue 1): the suspicion
        # z-score that actually drives the sigmoid gate, and the resulting gate.
        # High sus_z = suspicious (gate -> 0); compare attacker vs benign to see
        # whether the trust mechanism points the right direction.
        sus_z_list = defense_stats.get('sus_z')
        gate_list = defense_stats.get('gate')
        if isinstance(sus_z_list, list) and len(sus_z_list) == len(client_ids):
            sus_z_summary = ", ".join(
                f"c{cid}={v:.3f}" for cid, v in zip(client_ids, sus_z_list)
            )
            print(f"  🎯 sus_z:        {sus_z_summary}")
        if isinstance(gate_list, list) and len(gate_list) == len(client_ids):
            gate_summary = ", ".join(
                f"c{cid}={v:.3f}" for cid, v in zip(client_ids, gate_list)
            )
            print(f"  🚪 gate:         {gate_summary}")

        # V4-V8 shared CSE diagnostics: per-client CSE/median ratio and which
        # clients were high-confidence CSE-flagged this round. V5-V8 have no
        # scalar mult — show floor + r_hard (V6/V7 additionally show geometry,
        # while V8 prints propagated peers separately below). Test V6 FIRST:
        # it emits the v5_* ramp keys too, so a
        # 'v5_m_floor' in stats test alone would mislabel a V6 round as V5.
        v4_ratio_list = defense_stats.get('v4_ratio')
        v4_flagged_list = defense_stats.get('v4_flagged')
        if isinstance(v4_ratio_list, list) and len(v4_ratio_list) == len(client_ids):
            ratio_summary = ", ".join(
                f"c{cid}={v:.3f}" for cid, v in zip(client_ids, v4_ratio_list)
            )
            print(f"  🧪 cse/med:      {ratio_summary}")
        if isinstance(v4_flagged_list, list) and len(v4_flagged_list) == len(client_ids):
            flagged_ids = [cid for cid, f in zip(client_ids, v4_flagged_list) if f]
            if 'v6_geo_floor' in defense_stats:
                mult_desc = (
                    f"ramp floor={defense_stats.get('v5_m_floor')}, "
                    f"r_hard={defense_stats.get('v5_r_hard')}, "
                    f"geo_floor={defense_stats.get('v6_geo_floor')}"
                )
            elif 'v5_m_floor' in defense_stats:
                mult_desc = (
                    f"ramp floor={defense_stats.get('v5_m_floor')}, "
                    f"r_hard={defense_stats.get('v5_r_hard')}"
                )
            else:
                mult_desc = f"mult={defense_stats.get('v4_reject_mult')}"
            print(
                f"  ⛔ flagged:      {flagged_ids if flagged_ids else 'none'} "
                f"(tau={defense_stats.get('v4_tau_ratio')}, "
                f"k_cap={defense_stats.get('v4_k_cap')}, "
                f"{mult_desc})"
            )
        # V7 only: Tier-2 corroborated flags (cold-window, iso-corroborated
        # sub-1.85 rejections). Printed separately from the CSE flags above so
        # the log shows WHICH tier acted; 'none' inside the window is the
        # honest "geometry did not act this round" datum.
        v7_corrob_list = defense_stats.get('v7_corrob_flagged')
        if isinstance(v7_corrob_list, list) and len(v7_corrob_list) == len(client_ids):
            corrob_ids = [cid for cid, f in zip(client_ids, v7_corrob_list) if f]
            print(
                f"  🕸️ corrob:       {corrob_ids if corrob_ids else 'none'} "
                f"(tau_lo={defense_stats.get('v7_tau_lo')}, "
                f"iso_min={defense_stats.get('v7_iso_min')}, "
                f"mult={defense_stats.get('v7_corrob_mult')}, "
                f"window=[{defense_stats.get('v7_round_min')},"
                f"{defense_stats.get('v7_round_max')}], "
                f"geo_resolved={defense_stats.get('v7_geo_resolved')})"
            )
        # V8 only: distinguish high-confidence CSE seeds from clients reached
        # by the conservative dual-view hypergraph. A zero propagated set is
        # the explicit "V8 degraded exactly to V5 this round" diagnostic.
        v8_prop_list = defense_stats.get('v8_propagated_flagged')
        if isinstance(v8_prop_list, list) and len(v8_prop_list) == len(client_ids):
            propagated_ids = [
                cid for cid, f in zip(client_ids, v8_prop_list) if f
            ]
            risk = defense_stats.get('v8_propagated_risk') or []
            risk_summary = ", ".join(
                f"c{cid}={v:.3f}" for cid, v in zip(client_ids, risk)
            )
            print(
                f"  🕸️ V8 propagated: {propagated_ids if propagated_ids else 'none'} "
                f"(consensus_edges={defense_stats.get('v8_consensus_edge_count')}; "
                f"risk: {risk_summary if risk_summary else 'n/a'})"
            )
        # V6 only: the geometry factor. It is APPLIED only to flagged clients
        # (marked '*'); for everyone else the applied multiplier is a hard 1.0
        # and the number below is counterfactual — what the geometry WOULD have
        # said. Judge "did the hypergraph act" from the starred values only:
        # if those sit at ≈1.0 every round, V6 reduced to V5, which is a result
        # to report, not to tune away by lowering v6_geo_floor.
        v6_geo_mult_list = defense_stats.get('v6_geo_mult')
        if isinstance(v6_geo_mult_list, list) and len(v6_geo_mult_list) == len(client_ids):
            flags = defense_stats.get('v4_flagged') or []
            geo_summary = ", ".join(
                f"c{cid}{'*' if i < len(flags) and flags[i] else ''}={v:.3f}"
                for i, (cid, v) in enumerate(zip(client_ids, v6_geo_mult_list))
            )
            print(f"  📐 geo_mult:     {geo_summary}   (* = applied; "
                  f"unstarred are counterfactual)")

        # Compute similarity and distance metrics for visualization (unchanged).
        mode = getattr(self, 'similarity_mode', 'local_vs_global')
        if mode == 'local_vs_global':
            similarities = self._compute_similarities(updates, client_ids)
            similarity_matrix = None
            similarities_vs_global = None
        elif mode == 'pairwise':
            similarity_matrix, similarities = self._compute_similarities_pairwise(updates, client_ids)
            similarities_vs_global = None
        else:  # 'both'
            similarities_vs_global = self._compute_similarities(updates, client_ids)
            similarity_matrix, similarities = self._compute_similarities_pairwise(updates, client_ids)
        euclidean_distances = self._compute_euclidean_distances(updates, client_ids) if len(updates) > 0 else np.array([])

        aggregation_log = {
            'similarities': similarities.tolist(),
            'euclidean_distances': euclidean_distances.tolist() if len(euclidean_distances) > 0 else [],
            'accepted_clients': client_ids.copy(),
            'mean_similarity': float(similarities.mean()) if len(similarities) > 0 else 1.0,
            'std_similarity': float(similarities.std()) if len(similarities) > 0 else 0.0,
            'mean_euclidean_distance': euclidean_distances.mean().item() if len(euclidean_distances) > 0 else 0.0,
            'std_euclidean_distance': euclidean_distances.std().item() if len(euclidean_distances) > 0 else 0.0,
            'aggregated_update_norm': aggregated_update_norm,
            'defense_method': defense_label,
            'trust_weights': alpha_list if isinstance(alpha_list, list) else None,
            'raw_weights': raw_weights,
        }
        # Persist extra defense stats (skip bulky numpy blobs like 'Z' from the
        # main JSON log to keep result files lean; HMP runtime writes its own
        # stats file if enabled).
        for k in ('residual', 'recon_residual', 'sem_div',
                  'graph_residual_z', 'recon_residual_z', 'sem_div_z', 'hist_dev_z',
                  'hist_dev', 's', 'sus_z', 'sus_raw', 'gate',
                  'graph_weight', 'residual_weight_alpha',
                  'semantic_weight', 'hist_weight_beta_effective',
                  'zscore_mode', 'gate_rezscore', 'sus_ema_beta',
                  'semantic_reference',
                  'trust_mode_used', 'graph_channel_gated', 'graph_min_distinct',
                  'probe_cse',
                  'v4_cse', 'v4_ratio', 'v4_flagged', 'v4_multiplier',
                  'v4_median_cse', 'v4_tau_ratio', 'v4_k_cap', 'v4_reject_mult',
                  'v5_m_floor', 'v5_r_hard', 'v5_ramp_t',
                  'v6_geo_floor', 'v6_geo_gate', 'v6_geo_mult', 'v6_m_cse',
                  'v7_corrob_flagged', 'v7_iso', 'v7_geo_resolved',
                  'v7_tau_lo', 'v7_iso_min', 'v7_corrob_mult',
                  'v7_round_min', 'v7_round_max',
                  'v8_seed', 'v8_propagated_risk', 'v8_cse_evidence',
                  'v8_joint_evidence', 'v8_propagated_flagged',
                  'v8_propagated_multiplier', 'v8_propagation_matrix',
                  'v8_update_mutual', 'v8_behavior_mutual',
                  'v8_consensus_mutual', 'v8_update_edge_count',
                  'v8_behavior_edge_count', 'v8_consensus_edge_count',
                  'v8_recon_error',
                  'L_rec', 'L_smooth', 'L_struct', 'L_hist',
                  'fallback_reason', 'defense_time_ms'):
            if k in defense_stats:
                aggregation_log[k] = defense_stats[k]
        if similarity_matrix is not None:
            aggregation_log['similarity_matrix'] = similarity_matrix.tolist()
        if similarities_vs_global is not None:
            aggregation_log['similarities_vs_global'] = similarities_vs_global.tolist()
        aggregation_log['similarity_mode'] = mode

        return aggregation_log

    def evaluate_local_metrics(self, client) -> Tuple[float, float]:
        """
        Evaluate a client's local model on the server reference set in one pass.

        Returns (accuracy, classification_semantic_entropy).

        In real FL the server never sees client.model directly — it reconstructs
        the local model as w_global + Δ_i.  In this simulation the two are
        equivalent because client.model == w_global + Δ_i after local_train().
        The reference set is held out from the client-training pool.  It is
        deliberately separate from the official test split so this statistic
        may safely influence aggregation.

        Implementation: instead of moving client.model (full ~2GB Qwen) between
        CPU and GPU, we copy only the trainable flat params into the shared
        GPU-resident self._eval_model. In LoRA mode this is a few-MB tensor
        copy; in Full-FT mode it's equivalent to the old .to() call. The
        client's own model object is untouched.
        """
        # client.model lives on CPU between rounds; get_flat_params returns
        # a CPU tensor of just the trainable surface (LoRA-only with use_lora=True).
        flat = client.model.get_flat_params()
        self._eval_model.set_flat_params(flat)
        self._eval_model.eval()

        correct = 0
        total = 0
        total_cse = 0.0

        metric_loader = self.reference_loader
        if metric_loader is None:
            raise RuntimeError(
                "local client metrics require reference_loader; evaluating them "
                "on the official test split is forbidden"
            )

        with torch.no_grad():
            for batch in metric_loader:
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                labels = batch['labels'].to(self.device)

                outputs = self._eval_model(input_ids, attention_mask)

                predictions = torch.argmax(outputs, dim=1)
                correct += (predictions == labels).sum().item()
                total += labels.size(0)

                log_probs = F.log_softmax(outputs, dim=1)
                probs = log_probs.exp()
                batch_cse = -(probs * log_probs).sum(dim=1)
                total_cse += batch_cse.sum().item()

        accuracy = correct / total if total > 0 else 0.0
        cse = total_cse / total if total > 0 else 0.0
        return accuracy, cse

    def evaluate_local_accuracy(self, client) -> float:
        """Backward-compatible wrapper; prefer evaluate_local_metrics."""
        acc, _ = self.evaluate_local_metrics(client)
        return acc

    def _collect_local_metrics(self) -> Tuple[Dict[int, float], Dict[int, float]]:
        """
        Evaluate every client's (local accuracy, local CSE) on the server
        reference set and append to history. Shared by the diagnostic and
        pre-aggregation paths in run_round — client models are not
        modified by aggregation (only self.global_model is), so the values are
        identical regardless of when in the round this runs.
        """
        local_accs: Dict[int, float] = {}
        local_cses: Dict[int, float] = {}
        for client in self.clients:
            try:
                local_acc, local_cse = self.evaluate_local_metrics(client)
                local_accs[client.client_id] = local_acc
                local_cses[client.client_id] = local_cse

                if client.client_id not in self.history['local_accuracies']:
                    self.history['local_accuracies'][client.client_id] = []
                self.history['local_accuracies'][client.client_id].append(local_acc)

                if client.client_id not in self.history['local_cse']:
                    self.history['local_cse'][client.client_id] = []
                self.history['local_cse'][client.client_id].append(local_cse)
            except Exception as e:
                print(f"  ⚠️  Could not evaluate local metrics for client {client.client_id}: {e}")
        return local_accs, local_cses

    def _ensure_probe_batches(self) -> List[Dict[str, torch.Tensor]]:
        """Lazily snapshot a fixed subset of reference_loader for client probes."""
        if self._probe_batches is not None:
            return self._probe_batches
        if self.reference_loader is None:
            raise RuntimeError(
                "semantic probing requires reference_loader; test fallback is forbidden"
            )
        target = max(1, self.semantic_probe_size)
        batches: List[Dict[str, torch.Tensor]] = []
        if self.semantic_probe_stratified:
            batches = self._build_stratified_probe_batches(target)
        if not batches:
            # Deterministic head of reference_loader. Also the
            # fallback when the dataset does not expose labels for stratification.
            collected = 0
            for batch in self.reference_loader:
                # Snapshot tensors on CPU to keep peak GPU memory bounded.
                snapshot = {
                    'input_ids': batch['input_ids'].detach().cpu(),
                    'attention_mask': batch['attention_mask'].detach().cpu(),
                }
                batches.append(snapshot)
                collected += int(snapshot['input_ids'].shape[0])
                if collected >= target:
                    break
        self._probe_batches = batches
        return batches

    def _build_stratified_probe_batches(
        self, target: int
    ) -> List[Dict[str, torch.Tensor]]:
        """
        Class-stratified deterministic probe selection.

        Draws ~target/num_classes samples per class from reference_loader.dataset
        (seeded, without replacement) so the probe set is class-balanced --
        a head-of-loader snapshot can be class-skewed, which biases the
        semantic-divergence signal. Labels are consumed only here, for
        balancing; the snapshot itself carries no labels.

        Returns [] when the dataset does not expose a `labels` attribute,
        in which case the caller falls back to the head-of-loader path.
        """
        dataset = getattr(self.reference_loader, 'dataset', None)
        labels = getattr(dataset, 'labels', None)
        if dataset is None or labels is None or len(labels) == 0:
            print("  [Server] semantic_probe_stratified=True but reference dataset "
                  "exposes no labels; falling back to head-of-loader probes.")
            return []
        by_class: Dict[int, List[int]] = {}
        for idx, lab in enumerate(labels):
            by_class.setdefault(int(lab), []).append(idx)
        classes = sorted(by_class)
        rng = np.random.default_rng(self.semantic_probe_seed)
        quota, rem = divmod(target, len(classes))
        chosen: List[int] = []
        for j, c in enumerate(classes):
            k = min(quota + (1 if j < rem else 0), len(by_class[c]))
            if k <= 0:
                continue
            sel = rng.choice(len(by_class[c]), size=k, replace=False)
            chosen.extend(by_class[c][int(i)] for i in sel)
        if not chosen:
            return []
        chosen.sort()  # deterministic order, independent of class iteration
        bs = int(self.reference_loader.batch_size or 32)
        batches: List[Dict[str, torch.Tensor]] = []
        for start in range(0, len(chosen), bs):
            items = [dataset[i] for i in chosen[start:start + bs]]
            batches.append({
                'input_ids': torch.stack([it['input_ids'] for it in items]).cpu(),
                'attention_mask': torch.stack(
                    [it['attention_mask'] for it in items]).cpu(),
            })
        per_class = {c: 0 for c in classes}
        for i in chosen:
            per_class[int(labels[i])] += 1
        print(f"  [Server] stratified semantic probe: {len(chosen)} samples, "
              f"per-class counts={per_class} (seed={self.semantic_probe_seed})")
        return batches

    def evaluate_local_probe_distribution(self, client, flat_params: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Forward a client's local model over a fixed probe subset and return the
        per-sample softmax probabilities (the behavioral-fingerprint / semantic
        trust signal).

        Args:
            client: the client whose contribution is being probed.
            flat_params: optional explicit trainable-flat params to evaluate. When
                provided, the probe forwards ``w_global + Delta_i`` (the update the
                client actually SUBMITTED) instead of ``client.model``. This is the
                principled definition of the signal -- it measures the behavior of
                each client's contribution. For honest / label-flip attackers the
                two are numerically identical (client.model == w_global + Delta_i
                after local_train, so results are unchanged to float precision).
                For update-forging attackers that do NOT locally train (crafts_update,
                whose client.model stays == w_global while the malice lives only in
                the crafted Delta), ``client.model`` would look perfectly benign;
                passing w_global + Delta_i lets the semantic signal actually see the
                submitted attack. When None, falls back to ``client.model`` (legacy).

        Returns:
            (K, C) tensor on CPU, where K = number of probe samples actually
            taken (<= semantic_probe_size, capped by len(reference_loader.dataset))
            and C = num_labels.

        Uses the shared GPU-resident self._eval_model (see evaluate_local_metrics).
        """
        batches = self._ensure_probe_batches()
        flat = client.model.get_flat_params() if flat_params is None else flat_params
        self._eval_model.set_flat_params(flat)
        self._eval_model.eval()
        rows: List[torch.Tensor] = []
        with torch.no_grad():
            for batch in batches:
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                logits = self._eval_model(input_ids, attention_mask)
                rows.append(F.softmax(logits, dim=-1).detach().cpu())
        return torch.cat(rows, dim=0)
    
    def evaluate(self) -> float:
        """
        Evaluate the global model's performance.

        Returns:
            Clean accuracy (float) on the test set
        """
        accuracy, _, _ = self.evaluate_with_loss()
        return accuracy

    def evaluate_with_loss(self) -> Tuple[float, float, Optional[float]]:
        """
        Evaluate the global model's performance in a single pass and also
        compute the Classification Semantic Entropy (CSE) on the SeqCLS head.

        Returns:
            Tuple of (clean_accuracy, global_loss, classification_semantic_entropy_or_none).
            The third value is ``None`` when ``compute_classification_semantic_entropy`` is False.

        The CSE is the mean Shannon entropy of the softmax class distribution
        p(y|x) over the test set. Lower = more confident predictions; under a
        hallucination-inducing attack the model becomes less confident and CSE
        increases. A principled no-generation surrogate for Farquhar-style
        semantic entropy, using the C class labels as the "semantic clusters".
        """
        self.global_model.eval()

        # Evaluate clean accuracy, loss and CSE in one forward pass.
        correct = 0
        total = 0
        total_loss = 0.0
        total_cse = 0.0
        do_cse = self.compute_classification_semantic_entropy

        with torch.no_grad():
            for batch in self.test_loader:
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                labels = batch['labels'].to(self.device)

                outputs = self.global_model(input_ids, attention_mask)

                # Accuracy
                predictions = torch.argmax(outputs, dim=1)
                correct += (predictions == labels).sum().item()
                total += labels.size(0)

                # Cross-entropy loss (sum over batch for later averaging).
                loss = F.cross_entropy(outputs, labels, reduction='sum')
                total_loss += loss.item()

                if do_cse:
                    # Classification Semantic Entropy (per-sample Shannon entropy,
                    # summed here and averaged at the end).
                    # Use log_softmax for numerical stability.
                    log_probs = F.log_softmax(outputs, dim=1)
                    probs = log_probs.exp()
                    batch_cse = -(probs * log_probs).sum(dim=1)  # (B,)
                    total_cse += batch_cse.sum().item()

        clean_accuracy = correct / total if total > 0 else 0.0
        avg_loss = total_loss / total if total > 0 else 0.0
        mean_cse: Optional[float]
        if do_cse:
            mean_cse = total_cse / total if total > 0 else 0.0
        else:
            mean_cse = None

        # Record historical metrics.
        self.history['clean_acc'].append(clean_accuracy)
        if 'cse' not in self.history:
            self.history['cse'] = []
        self.history['cse'].append(mean_cse)

        return clean_accuracy, avg_loss, mean_cse
    
    def evaluate_global_loss(self) -> float:
        """
        Evaluate the global model's loss on the test set.
        For efficiency, use evaluate_with_loss() if you also need accuracy.
        
        Returns:
            Global loss (float) on the test set (cross-entropy loss)
        """
        _, loss, _ = self.evaluate_with_loss()
        return loss

    def adaptive_adjustment(self, round_num: int):
        """Adaptively adjust parameters based on historical performance."""
        # Fixed server_lr (no adaptive change)
        pass

    def run_round(self, round_num: int) -> Dict:
        """Execute one round of federated learning - stable version."""
        print(f"\n{'=' * 60}")
        print(f"Round {round_num + 1}/{self.total_rounds}")

        # Track the current round so the defense plugin can use it for history.
        self._current_round = int(round_num)

        # Adaptive adjustment
        self.adaptive_adjustment(round_num)

        # Display current parameters
        print(f"Current Parameters: server_lr={self.server_lr:.2f}")
        print(f"{'=' * 60}")

        # Broadcast the model
        print("📡 Broadcasting the global model...")
        self.broadcast_model()

        # Phase 0.5: Constrained-attacker setup (only for update-forging attackers
        # that advertise crafts_update — none exist since AugMP's removal
        # 2026-08-09, so this block always skips). Hands each such attacker the
        # current global params + constraint bounds so it can build the global-loss
        # proxy F(w'_g) and the distance/cosine-similarity constraints inside
        # camouflage_update. This whole block is SKIPPED for every existing run
        # (benign / hallucination / alie / gaussian / sign_flipping), so their code
        # path is byte-for-byte unchanged. dist_bound / sim_bound_* are read via
        # getattr (default None = auto-derive from benign statistics); main.py no
        # longer sets them (AugMP removed 2026-08-09).
        needs_constraint_setup = any(getattr(c, 'crafts_update', False) for c in self.clients)
        if needs_constraint_setup:
            global_params = self.global_model.get_flat_params()
            total_data_size = 0.0
            benign_data_sizes: Dict[int, float] = {}
            for client in self.clients:
                if getattr(client, 'is_attacker', False):
                    total_data_size += float(getattr(client, 'claimed_data_size', 1.0))
                else:
                    client_data_size = len(getattr(client, 'data_indices', [])) or 1.0
                    benign_data_sizes[client.client_id] = client_data_size
                    total_data_size += client_data_size
            for client in self.clients:
                if getattr(client, 'is_attacker', False):
                    client.set_global_model_params(global_params)
                    client.set_constraint_params(
                        dist_bound=getattr(self, 'dist_bound', None),
                        sim_bound_low=getattr(self, 'sim_bound_low', None),
                        sim_bound_up=getattr(self, 'sim_bound_up', None),
                        total_data_size=total_data_size,
                        benign_data_sizes=benign_data_sizes,
                    )

        # Phase 1: Preparation
        print("\n🔧 Phase 1: Client Preparation")
        for client in self.clients:
            client.set_round(round_num)
            # Use is_attacker attribute instead of isinstance to support both GRMP and ALIE attackers
            if getattr(client, 'is_attacker', False):
                client.prepare_for_round(round_num)

        # Phase 2: Local Training
        print("\n💪 Phase 2: Local Training")
        initial_updates = {}
        for client in self.clients:
            update = client.local_train()
            initial_updates[client.client_id] = update
            print(f"  ✓ Client {client.client_id} completed training")

        # Phase 3: Attacker Camouflage
        print("\n🎭 Phase 3: Attacker Camouflage")
        benign_updates = []
        benign_client_ids = []
        for client_id, update in initial_updates.items():
            client = self.clients[client_id]
            if not getattr(client, 'is_attacker', False):
                benign_updates.append(update)
                benign_client_ids.append(client_id)
        
        print(f"  Captured {len(benign_updates)} benign updates for camouflage.")
        
        # ===== NEW: Store completed attacker updates for coordinated optimization =====
        completed_attacker_updates = {}  # {client_id: update_tensor}
        completed_attacker_client_ids = []  # Keep order
        completed_attacker_data_sizes = {}  # {client_id: claimed_data_size}
        # ==============================================================================
        
        final_updates = {}
        for client_id, update in initial_updates.items():
            client = self.clients[client_id]
            if getattr(client, 'is_attacker', False):
                print(f"  ⚠️ Triggering camouflage logic for Client {client_id}")
                client.receive_benign_updates(benign_updates, client_ids=benign_client_ids)
                
                # ===== NEW: Pass completed attacker updates to current attacker =====
                if completed_attacker_updates:
                    client.receive_attacker_updates(
                        updates=list(completed_attacker_updates.values()),
                        client_ids=completed_attacker_client_ids,
                        data_sizes=completed_attacker_data_sizes
                    )
                # ====================================================================
                
                final_updates[client_id] = client.camouflage_update(update)
                
                # ===== NEW: Store current attacker's update for subsequent attackers =====
                completed_attacker_updates[client_id] = final_updates[client_id]
                completed_attacker_client_ids.append(client_id)
                completed_attacker_data_sizes[client_id] = float(getattr(client, 'claimed_data_size', 1.0))
                # =========================================================================
            else:
                final_updates[client_id] = update

        # Phase 4: Aggregation
        print("\n📊 Phase 4: Model Aggregation")
        # Ensure deterministic order of keys
        sorted_client_ids = sorted(final_updates.keys())
        final_update_list = [final_updates[cid] for cid in sorted_client_ids]

        # Optional Phase 3.5: per-client probe forward for the semantic-divergence
        # trust signal. Only computed when the active defense actually consumes it.
        probe_tensor: Optional[torch.Tensor] = None
        if self._needs_probe:
            # Semantic signal. For every ordinary client (benign, and attackers that
            # bake their poison into local training such as Hallucination) we probe
            # client.model EXACTLY as before. Only for update-forging attackers
            # (crafts_update, whose client.model stays == w_global while the malice
            # lives in the crafted Delta) do we instead probe w_global + Delta_i, so the semantic
            # signal can actually see the submitted attack. self.global_model is still
            # the round-start global here (Phase 4 aggregation happens below).
            global_flat = None
            probe_rows: List[torch.Tensor] = []
            for cid in sorted_client_ids:
                client = self.client_dict.get(cid)
                if client is None:
                    raise KeyError(f"client_id {cid} not registered with server")
                if getattr(client, 'crafts_update', False):
                    if global_flat is None:
                        global_flat = self.global_model.get_flat_params()
                        if global_flat.device.type == "cuda":
                            global_flat = global_flat.cpu()
                    submitted = final_updates[cid]
                    if submitted.device.type == "cuda":
                        submitted = submitted.cpu()
                    probe_rows.append(
                        self.evaluate_local_probe_distribution(
                            client, flat_params=global_flat + submitted
                        )
                    )
                else:
                    probe_rows.append(self.evaluate_local_probe_distribution(client))
            # All rows must have identical shape (K, C) -- same probe set, same head.
            probe_tensor = torch.stack(probe_rows, dim=0)  # (N, K, C)

        # Optional Phase 3.6: pre-aggregation per-client reference metrics for the
        # CSE-reject trust rules (V4-V8).
        # Client models are untouched by aggregation, so these values are
        # identical to the legacy post-aggregation evaluation — computed once
        # here and reused for the round log below (no duplicate eval cost).
        local_accs_this_round: Dict[int, float] = {}
        local_cse_this_round: Dict[int, float] = {}
        local_cse_vector: Optional[List[float]] = None
        if self._needs_local_cse:
            if any(getattr(c, 'crafts_update', False) for c in self.clients):
                raise RuntimeError(
                    "CSE-reject trust modes (v4_cse_reject / v5_cse_reject / "
                    "v6_cse_reject_geo / v7_cse_reject_corrob / "
                    "v8_hmp_cse_propagation) "
                    "are not supported with update-forging attackers "
                    "(crafts_update): local CSE evaluates "
                    "client.model, which such attackers leave looking benign "
                    "(the poison lives only in the crafted update)."
                )
            local_accs_this_round, local_cse_this_round = self._collect_local_metrics()
            missing = [cid for cid in sorted_client_ids
                       if cid not in local_cse_this_round]
            if missing:
                raise RuntimeError(
                    "CSE-reject trust mode: local CSE evaluation failed "
                    f"for clients {missing}; cannot aggregate without it."
                )
            local_cse_vector = [float(local_cse_this_round[cid])
                                for cid in sorted_client_ids]
            # A NaN slips past the `missing` check above and then poisons the
            # whole rule: torch.median propagates NaN, every ratio becomes
            # NaN, `nan > tau` is False, so ZERO clients are flagged and the
            # round silently degrades to FedAvg with full attacker weight —
            # indistinguishable in the archive from a genuinely clean round.
            # Same loud-crash contract as `missing`. +inf does NOT disable the
            # rule (the median stays finite and that client flags normally),
            # but it still means the local eval is broken, so both are
            # rejected here rather than silently scored.
            bad = [cid for cid, v in zip(sorted_client_ids, local_cse_vector)
                   if not math.isfinite(v)]
            if bad:
                raise RuntimeError(
                    "CSE-reject trust mode: non-finite local CSE for clients "
                    f"{bad} (values "
                    f"{[local_cse_this_round[c] for c in bad]}); the rejection "
                    "rule would silently disable itself this round "
                    "(NaN median -> no flags -> plain FedAvg). Aborting so "
                    "the run is not silently invalidated."
                )

        aggregation_log = self.aggregate_updates(
            final_update_list, sorted_client_ids,
            probe_distributions=probe_tensor,
            local_cse=local_cse_vector,
        )

        # Evaluate the global model (compute accuracy, loss and CSE in one pass).
        clean_acc, global_loss, mean_cse = self.evaluate_with_loss()

        # Evaluate per-client reference accuracy and CSE (single forward pass each).
        # When eval_local_every_n_rounds > 1, we only evaluate on round 0, the
        # final round, and every n-th round in between -- a sparser diagnostic
        # trace in exchange for fewer N-times-reference-set forwards.
        n_eval = self.eval_local_every_n_rounds
        is_final_round = (round_num + 1) == self.total_rounds
        do_local_eval = (
            n_eval <= 1
            or round_num == 0
            or is_final_round
            or ((round_num + 1) % n_eval == 0)
        )
        if self._needs_local_cse:
            # V4/V5/V6 path: local metrics were already evaluated pre-aggregation
            # this round (every round — the trust rule needs them even when
            # eval_local_every_n_rounds would skip); reuse those values here.
            pass
        elif do_local_eval:
            local_accs_this_round, local_cse_this_round = self._collect_local_metrics()
        else:
            print(f"  ⏭  Skipping per-client local eval this round "
                  f"(eval_local_every_n_rounds={n_eval}).")

        # Create log for the current round
        round_log = {
            'round': round_num + 1,
            'clean_accuracy': clean_acc,
            'global_loss': global_loss,
            'classification_semantic_entropy': mean_cse,
            'acc_diff': (abs(clean_acc - self.history['clean_acc'][-2])
                         if len(self.history['clean_acc']) > 1 else 0.0),
            'aggregation': aggregation_log,
            'server_lr': self.server_lr,
            'local_accuracies': local_accs_this_round,
            'local_cse': local_cse_this_round,
            'local_metric_source': (
                'server_reference' if self.reference_loader is not None else 'unavailable'
            ),
        }

        self.log_data.append(round_log)

        # Display results
        print(f"\n📊 Round {round_num + 1} Results:")
        print(f"  Clean Accuracy: {clean_acc:.4f}")
        if len(self.history['clean_acc']) > 1:
            prev_clean = self.history['clean_acc'][-2]
            delta_prev = clean_acc - prev_clean
            best_clean = max(self.history['clean_acc'])
            delta_best = clean_acc - best_clean
            print(f"  ΔClean vs prev: {delta_prev:+.4f}")
            print(f"  ΔClean vs best: {delta_best:+.4f}")
        print(f"  Global Loss: {global_loss:.4f}")
        if mean_cse is not None:
            print(f"  Global CSE: {mean_cse:.4f}")
        else:
            print("  Global CSE: (disabled via config)")

        # Per-client local metrics table (only when evaluated this round).
        if local_accs_this_round:
            metric_source = (
                "server reference set" if self.reference_loader is not None
                else "unavailable reference set"
            )
            print(f"  Per-client local metrics (local model on {metric_source}):")
            attacker_ids = {c.client_id for c in self.clients if getattr(c, 'is_attacker', False)}
            for cid in sorted(local_accs_this_round):
                tag = "ATK" if cid in attacker_ids else "BGN"
                acc_v = local_accs_this_round[cid]
                cse_v = local_cse_this_round.get(cid, float('nan'))
                print(f"    [{tag}] Client {cid}: acc={acc_v:.4f}  cse={cse_v:.4f}")

        return round_log