FedSurg Challenge Tests Federated Learning on Surgical Video

"""
FedSurg toolkit. Federated learning for patient level appendicitis grading from surgical video.
Reference PyTorch implementation of the benchmark described in
    Kirchner et al. Federated Learning for surgical vision in appendicitis classification,
    Results of the FedSurg EndoVis 2024 challenge. Medical Image Analysis 115 (2027) 104290.
    https://doi.org/10.1016/j.media.2026.104290

The paper is a challenge report, so there is no single proposed network. This file implements
the whole benchmark so that the protocol, the three submitted strategies and the organizer
baselines can be rerun on any data laid out the same way.

    1. Metrics            macro F1 with a fixed denominator of six grades (Eq. 1-2),
                          Expected Cost with linear ordinal weights (Eq. 3-4)
    2. References         TrivialGlobal, TrivialLocal, RandomUniform, RandomPrior (Sec. 3.4)
    3. Statistics         10,000 replicate bootstrap rank stability with shared indices,
                          paired permutation tests over test cases (Sec. 2.4.1)
    4. Communication      Eq. 5-7 for server based FL and peer to peer swarm learning
    5. Models             a frozen frame encoder (stand in for ViViT, EndoViT, ResNet-50)
                          with four heads
                            TemporalHead   video level transformer over sampled frames (Team Santhi)
                            FrameHead      per frame classifier plus majority vote   (Team Elbflorenz)
                            MetricHead     cosine triplet embedding plus prototypes  (Team Camma)
                            LSTMHead       frame features into one LSTM layer        (SurgTempoNet baseline)
                          plus LoRA adapters for the parameter efficient baselines
    6. Federation         FedAvg, FedMedian, FedOpt (server Adam) and client side FedSAM
    7. Tasks              Task 1 generalization to held out Center 4,
                          Task 2 adaptation by fine tuning on each training center

Honest notes
    * The real backbones (ViViT on Kinetics-400, EndoViT on Endo700k, ResNet-50) are replaced by
      a small frozen CNN so the file runs on a CPU. Swap `FrameEncoder` for a pretrained model.
    * The smoke test reproduces the paper's reference classifier values and Task 1 scores exactly
      from label counts and confusion matrices printed in the paper, then runs the federated
      pipeline end to end on synthetic videos that follow the paper's label distribution.
    * Loss weights, the frame window width and FedSAM rho are not all reported; values marked
      "our choice" are ours.
"""
import copy
import math
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

C = 6  # inflammation grades 0 to 5

# ---------------------------------------------------------------------------
# 1. Metrics
# ---------------------------------------------------------------------------

def confusion(y_true, y_pred, c=C):
    m = np.zeros((c, c), dtype=np.int64)
    np.add.at(m, (np.asarray(y_true), np.asarray(y_pred)), 1)
    return m


def macro_f1(y_true, y_pred, c=C):
    """Eq. 1-2. Grades absent from both truth and prediction score zero but still count in the
    denominator of six, so the ceiling drops below 100 percent when a grade is missing."""
    m = confusion(y_true, y_pred, c)
    tp = np.diag(m).astype(float)
    fp, fn = m.sum(0) - tp, m.sum(1) - tp
    denom = 2 * tp + fp + fn
    f1 = np.divide(2 * tp, denom, out=np.zeros(c), where=denom > 0)
    return f1.mean()


def expected_cost(y_true, y_pred, c=C):
    """Eq. 3-4. Linear ordinal cost |i - j| / (C - 1), averaged over cases."""
    y_true, y_pred = np.asarray(y_true), np.asarray(y_pred)
    return (np.abs(y_true - y_pred) / (c - 1)).mean()


def f1_ceiling(y_true, c=C):
    return len(np.unique(y_true)) / c

# ---------------------------------------------------------------------------
# 2. Reference classifiers (use training labels only)
# ---------------------------------------------------------------------------

def mode(labels):
    return int(np.bincount(labels, minlength=C).argmax())


def trivial_constant(y_true, grade):
    pred = np.full(len(y_true), grade)
    return macro_f1(y_true, pred), expected_cost(y_true, pred)


def random_reference(y_true, probs, draws=10000, seed=0):
    """RandomUniform (probs = 1/6 each) or RandomPrior (probs = training prevalence).
    Returns the simulation mean and 95 percent percentile interval for F1 and EC."""
    rng = np.random.default_rng(seed)
    f1s, ecs = np.empty(draws), np.empty(draws)
    for d in range(draws):
        pred = rng.choice(C, size=len(y_true), p=probs)
        f1s[d], ecs[d] = macro_f1(y_true, pred), expected_cost(y_true, pred)
    q = lambda a: (a.mean(), np.percentile(a, 2.5), np.percentile(a, 97.5))
    return q(f1s), q(ecs)

# ---------------------------------------------------------------------------
# 3. Bootstrap rank stability and paired permutation tests
# ---------------------------------------------------------------------------

def bootstrap_ranks(y_true, preds: dict, metric="f1", iters=10000, seed=0):
    """Shared resampling indices for all methods, predictions aligned by case.
    Returns rank frequency per method and pairwise win probabilities."""
    rng = np.random.default_rng(seed)
    names = list(preds)
    n, k = len(y_true), len(names)
    fn = macro_f1 if metric == "f1" else expected_cost
    higher_better = metric == "f1"
    ranks = np.zeros((k, k))
    wins = np.zeros((k, k))
    y_true = np.asarray(y_true)
    for _ in range(iters):
        idx = rng.integers(0, n, n)
        vals = np.array([fn(y_true[idx], np.asarray(preds[m])[idx]) for m in names])
        order = np.argsort(-vals if higher_better else vals, kind="stable")
        for r, j in enumerate(order):
            ranks[j, r] += 1
        better = vals[:, None] > vals[None] if higher_better else vals[:, None] < vals[None]
        wins += better
    return names, ranks / iters, wins / iters


def paired_permutation(y_true, pred_a, pred_b, metric="f1", perms=10000, seed=0):
    """Two sided test. Each case's two predictions are swapped with probability 0.5."""
    rng = np.random.default_rng(seed)
    fn = macro_f1 if metric == "f1" else expected_cost
    y, a, b = map(np.asarray, (y_true, pred_a, pred_b))
    obs = fn(y, a) - fn(y, b)
    count = 0
    for _ in range(perms):
        swap = rng.random(len(y)) < 0.5
        pa, pb = np.where(swap, b, a), np.where(swap, a, b)
        count += abs(fn(y, pa) - fn(y, pb)) >= abs(obs) - 1e-12
    return obs, max(count, 1) / perms


def paired_permutation_task2(centers, metric="ec", perms=10000, seed=0):
    """Task 2 version. `centers` is a list of (y_true, pred_a, pred_b) per center. The swap is
    done within each center and the per center metrics are averaged with equal weight."""
    rng = np.random.default_rng(seed)
    fn = macro_f1 if metric == "f1" else expected_cost
    avg = lambda triples: np.mean([fn(y, a) - fn(y, b) for y, a, b in triples])
    data = [tuple(map(np.asarray, t)) for t in centers]
    obs = avg(data)
    count = 0
    for _ in range(perms):
        perm = []
        for y, a, b in data:
            s = rng.random(len(y)) < 0.5
            perm.append((y, np.where(s, b, a), np.where(s, a, b)))
        count += abs(avg(perm)) >= abs(obs) - 1e-12
    return obs, max(count, 1) / perms

# ---------------------------------------------------------------------------
# 4. Communication cost (Eq. 5-7)
# ---------------------------------------------------------------------------

def fl_cost(bytes_per_model, clients, rounds):
    return 2 * clients * rounds * bytes_per_model


def swarm_sync_events(epochs, batches_per_center, f_sync):
    return sum(epochs * b // f_sync for b in batches_per_center) + 1


def swarm_cost(bytes_per_model, clients, sync_events):
    return 2 * (clients - 1) * sync_events * bytes_per_model

# ---------------------------------------------------------------------------
# 5. Models
# ---------------------------------------------------------------------------

class FrameEncoder(nn.Module):
    """Frozen stand in for a pretrained image or video backbone. Maps (B, T, 3, H, W) to (B, T, D)."""

    def __init__(self, dim=64):
        super().__init__()
        self.net = nn.Sequential(nn.Conv2d(3, 32, 3, 2, 1), nn.GELU(), nn.Conv2d(32, 32, 3, 1, 1), nn.GELU(),
                                 nn.AdaptiveAvgPool2d(4), nn.Flatten(), nn.Linear(32 * 16, dim))
        self.requires_grad_(False)
        self.dim = dim

    def forward(self, x):
        b, t = x.shape[:2]
        x = (x.flatten(0, 1) - 0.5) / 0.25                                  # ImageNet style input normalization
        return self.net(x).view(b, t, -1)


class LoRALinear(nn.Module):
    """y = W x + (alpha / r) B A x with W frozen. Rank 8, alpha 16 in the paper's PEFT baselines."""

    def __init__(self, base: nn.Linear, r=8, alpha=16):
        super().__init__()
        self.base = base
        self.base.requires_grad_(False)
        self.A = nn.Parameter(torch.randn(r, base.in_features) / math.sqrt(base.in_features))
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
        self.scale = alpha / r

    def forward(self, x):
        return self.base(x) + self.scale * (x @ self.A.T @ self.B.T)


class LoRAAttention(nn.Module):
    """Multihead self attention with LoRA on the query, key, value and output projections."""

    def __init__(self, dim, heads=4, lora=False, r=8, alpha=16):
        super().__init__()
        self.h = heads
        mk = lambda: nn.Linear(dim, dim)
        self.q, self.k, self.v, self.o = mk(), mk(), mk(), mk()
        if lora:
            self.q, self.k, self.v, self.o = (LoRALinear(m, r, alpha) for m in (self.q, self.k, self.v, self.o))

    def forward(self, x):
        b, t, d = x.shape
        sh = lambda z: z.view(b, t, self.h, d // self.h).transpose(1, 2)
        a = F.scaled_dot_product_attention(sh(self.q(x)), sh(self.k(x)), sh(self.v(x)))
        return self.o(a.transpose(1, 2).reshape(b, t, d))


class TemporalHead(nn.Module):
    """Video level classification over sampled frames (Team Santhi's ViViT idea; LoRA variant = ViViT+LoRA)."""

    def __init__(self, dim=64, layers=2, lora=False, max_t=256):
        super().__init__()
        self.cls = nn.Parameter(torch.zeros(1, 1, dim))
        self.pos = nn.Parameter(torch.randn(1, max_t + 1, dim) * 0.02)
        self.blocks = nn.ModuleList()
        for _ in range(layers):
            self.blocks.append(nn.ModuleDict(dict(n1=nn.LayerNorm(dim), att=LoRAAttention(dim, lora=lora), n2=nn.LayerNorm(dim),
                                                  mlp=nn.Sequential(nn.Linear(dim, 2 * dim), nn.GELU(), nn.Linear(2 * dim, dim)))))
        if lora:  # freeze everything in the temporal encoder except LoRA factors
            for blk in self.blocks:
                for name, p in blk.named_parameters():
                    p.requires_grad_(name.endswith(".A") or name.endswith(".B"))
            self.cls.requires_grad_(False)
            self.pos.requires_grad_(False)
        self.head = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, dim), nn.GELU(), nn.Linear(dim, dim), nn.GELU(), nn.Linear(dim, C))

    def forward(self, f):
        x = torch.cat([self.cls.expand(len(f), -1, -1), f], 1) + self.pos[:, : f.shape[1] + 1]
        for blk in self.blocks:
            x = x + blk["att"](blk["n1"](x))
            x = x + blk["mlp"](blk["n2"](x))
        return self.head(x[:, 0])


class FrameHead(nn.Module):
    """Per frame classifier; the video label is a majority vote, ties broken by mean confidence
    (Team Elbflorenz). With `lora` it becomes the EndoViT+LoRA style frame level baseline."""

    def __init__(self, dim=64):
        super().__init__()
        self.head = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, dim), nn.GELU(), nn.Linear(dim, C))

    def forward(self, f):                       # frame logits (B, T, C)
        return self.head(f)

    @torch.no_grad()
    def vote(self, f):
        logits = self.forward(f)
        probs = logits.softmax(-1)
        votes = F.one_hot(logits.argmax(-1), C).sum(1).float()           # (B, C)
        top = votes == votes.max(-1, keepdim=True).values
        conf = probs.mean(1).masked_fill(~top, -1)                        # tie break by average confidence
        return conf.argmax(-1)


class SwitchNorm1d(nn.Module):
    """Switchable normalization for feature vectors, softmax gated mix of batch and layer statistics."""

    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.w = nn.Parameter(torch.ones(dim))
        self.b = nn.Parameter(torch.zeros(dim))
        self.mix_mean = nn.Parameter(torch.ones(2))
        self.mix_var = nn.Parameter(torch.ones(2))
        self.eps = eps

    def forward(self, x):
        mb, vb = x.mean(0, keepdim=True), x.var(0, unbiased=False, keepdim=True) if len(x) > 1 else torch.ones_like(x[:1])
        ml, vl = x.mean(1, keepdim=True), x.var(1, unbiased=False, keepdim=True)
        wm, wv = self.mix_mean.softmax(0), self.mix_var.softmax(0)
        mean, var = wm[0] * mb + wm[1] * ml, wv[0] * vb + wv[1] * vl
        return (x - mean) / torch.sqrt(var + self.eps) * self.w + self.b


class MetricHead(nn.Module):
    """Siamese embedding into a 256 dimensional L2 normalized space trained with a cosine triplet
    margin loss; inference compares to class prototypes or to all support embeddings (Team Camma)."""

    def __init__(self, dim=64, emb=256):
        super().__init__()
        self.proj = nn.Sequential(nn.Linear(dim, emb), SwitchNorm1d(emb), nn.GELU(), nn.Linear(emb, emb))

    def forward(self, f):                       # f (B, T, D) of selected frames
        return F.normalize(self.proj(f.mean(1)), dim=-1)

    @staticmethod
    def triplet_loss(z, y, margin=0.2):
        sim = z @ z.T
        same = y[:, None] == y[None]
        eye = torch.eye(len(y), dtype=torch.bool)
        pos = sim.masked_fill(~same | eye, 2).min(1).values               # hardest positive
        neg = sim.masked_fill(same, -2).max(1).values                     # hardest negative
        valid = (same & ~eye).any(1) & (~same).any(1)
        if valid.sum() == 0:
            return z.sum() * 0
        return F.relu(neg - pos + margin)[valid].mean()

    @torch.no_grad()
    def predict(self, z, z_support, y_support, mode="prototype"):
        if mode == "prototype":
            protos = torch.stack([F.normalize(z_support[y_support == c].mean(0), dim=0) if (y_support == c).any()
                                  else torch.full((z.shape[1],), float("nan")) for c in range(C)])
            sim = z @ protos.nan_to_num(-10).T
            sim[:, torch.isnan(protos[:, 0])] = -10
            return sim.argmax(-1)
        sims = z @ z_support.T                                            # per sample, mean similarity per class
        scores = torch.stack([sims[:, y_support == c].mean(1) if (y_support == c).any() else torch.full((len(z),), -10.0)
                              for c in range(C)], 1)
        return scores.argmax(-1)


class LSTMHead(nn.Module):
    """SurgTempoNet style, frame features into a single LSTM layer with hidden size 160."""

    def __init__(self, dim=64, hidden=160):
        super().__init__()
        self.lstm = nn.LSTM(dim, hidden, batch_first=True)
        self.fc = nn.Linear(hidden, C)

    def forward(self, f):
        out, _ = self.lstm(f)
        return self.fc(out[:, -1])

# ---------------------------------------------------------------------------
# 6. Frame sampling
# ---------------------------------------------------------------------------

def sample_frames_santhi(n_total, k=32, center=100, half_window=16, p_center=0.6, rng=None):
    """Two thirds from a narrow window around the centre frame, the rest from the whole video
    with probability 0.6 of landing near the centre. Window width is our choice."""
    rng = rng or np.random.default_rng()
    center = min(center, n_total // 2)
    lo, hi = max(0, center - half_window), min(n_total, center + half_window)
    k_win = (2 * k) // 3
    idx = list(rng.choice(np.arange(lo, hi), size=k_win, replace=hi - lo < k_win))
    for _ in range(k - k_win):
        if rng.random() < p_center:
            idx.append(int(rng.integers(max(0, center - 2 * half_window), min(n_total, center + 2 * half_window))))
        else:
            idx.append(int(rng.integers(0, n_total)))
    return np.sort(np.array(idx))


def federated_standardizer(client_feats):
    """Per dimension mean and std from summed client statistics (count, sum, sum of squares), so no
    raw features leave a center. Returns a function that standardizes (B, T, D) features."""
    n = sum(f.shape[0] * f.shape[1] for f in client_feats)
    s1 = sum(f.sum((0, 1)) for f in client_feats)
    s2 = sum((f ** 2).sum((0, 1)) for f in client_feats)
    mu = s1 / n
    sd = (s2 / n - mu ** 2).clamp(min=1e-12).sqrt()
    return lambda f: (f - mu) / sd


def sample_frames_equidistant(n_total, k=100):
    return np.linspace(0, n_total - 1, min(k, n_total)).round().astype(int)

# ---------------------------------------------------------------------------
# 7. Federated strategies
# ---------------------------------------------------------------------------

def trainable_state(model):
    return {n: p.detach().clone() for n, p in model.named_parameters() if p.requires_grad}


def load_trainable(model, state):
    with torch.no_grad():
        for n, p in model.named_parameters():
            if n in state:
                p.copy_(state[n])


def fedavg(states, weights):
    w = torch.tensor(weights, dtype=torch.float32) / sum(weights)
    return {k: sum(wi * s[k] for wi, s in zip(w, states)) for k in states[0]}


def fedmedian(states, weights=None):
    return {k: torch.stack([s[k] for s in states]).median(0).values for k in states[0]}


class FedOptServer:
    """Server side Adam on the pseudo gradient (global minus average client model), Reddi et al."""

    def __init__(self, lr=1e-2, b1=0.9, b2=0.99, eps=1e-3):
        self.lr, self.b1, self.b2, self.eps = lr, b1, b2, eps
        self.m, self.v = None, None

    def __call__(self, global_state, states, weights):
        avg = fedavg(states, weights)
        delta = {k: avg[k] - global_state[k] for k in avg}
        if self.m is None:
            self.m = {k: torch.zeros_like(d) for k, d in delta.items()}
            self.v = {k: torch.zeros_like(d) for k, d in delta.items()}
        out = {}
        for k, d in delta.items():
            self.m[k] = self.b1 * self.m[k] + (1 - self.b1) * d
            self.v[k] = self.b2 * self.v[k] + (1 - self.b2) * d * d
            out[k] = global_state[k] + self.lr * self.m[k] / (self.v[k].sqrt() + self.eps)
        return out


def sam_step(model, loss_fn, opt, rho=0.05, adaptive=True):
    """One (A)SAM update. Perturb weights toward higher loss, take the gradient there, restore, step."""
    loss = loss_fn()
    opt.zero_grad()
    loss.backward()
    params = [p for p in model.parameters() if p.requires_grad and p.grad is not None]
    scale = [(p.abs() if adaptive else torch.ones_like(p)) for p in params]
    norm = torch.norm(torch.stack([(s * p.grad).norm() for s, p in zip(scale, params)])) + 1e-12
    eps = []
    with torch.no_grad():
        for s, p in zip(scale, params):
            e = rho * s * s * p.grad / norm if adaptive else rho * p.grad / norm
            p.add_(e)
            eps.append(e)
    opt.zero_grad()
    loss_fn().backward()
    with torch.no_grad():
        for p, e in zip(params, eps):
            p.sub_(e)
    opt.step()
    return loss.item()

# ---------------------------------------------------------------------------
# 8. Wrapping heads into trainable clients
# ---------------------------------------------------------------------------

def class_weights(labels):
    counts = np.bincount(labels, minlength=C).astype(float)
    w = np.where(counts > 0, counts.sum() / np.maximum(counts, 1), 0.0)
    return torch.tensor(w / w[w > 0].mean(), dtype=torch.float32)


class Method:
    """Bundles an encoder, a head and how to train and predict. kind in
    {'temporal', 'temporal_lora', 'frame', 'frame_lora', 'metric', 'lstm'}."""

    def __init__(self, kind, dim=64, seed=0):
        torch.manual_seed(seed)
        self.kind = kind
        self.encoder = FrameEncoder(dim)
        if kind.startswith("temporal"):
            self.head = TemporalHead(dim, lora=kind.endswith("lora"))
        elif kind.startswith("frame"):
            self.head = FrameHead(dim)
        elif kind == "metric":
            self.head = MetricHead(dim)
        else:
            self.head = LSTMHead(dim)
        self.support = None

    def features(self, videos):
        with torch.no_grad():
            return self.encoder(videos)

    def loss(self, f, y, cw):
        if self.kind.startswith("frame"):
            logits = self.head(f)
            return F.cross_entropy(logits.flatten(0, 1), y.repeat_interleave(f.shape[1]), weight=cw)
        if self.kind == "metric":
            return MetricHead.triplet_loss(self.head(f), y)
        return F.cross_entropy(self.head(f), y, weight=cw)

    @torch.no_grad()
    def predict(self, f):
        if self.kind.startswith("frame"):
            return self.head.vote(f)
        if self.kind == "metric":
            zs, ys = self.support
            return self.head.predict(self.head(f), zs, ys)
        return self.head(f).argmax(-1)

    @torch.no_grad()
    def set_support(self, f, y):
        self.support = (self.head(f), torch.as_tensor(y))


def local_train(method, f, y, epochs, lr, bs, use_sam=False, weighted=True):
    params = [p for p in method.head.parameters() if p.requires_grad]
    opt = torch.optim.SGD(params, lr=lr, momentum=0.9) if use_sam else torch.optim.Adam(params, lr=lr)
    cw = class_weights(y.numpy()) if weighted else None
    method.head.train()
    for _ in range(epochs):
        perm = torch.randperm(len(y))
        for i in range(0, len(y), bs):
            idx = perm[i:i + bs]
            fn = lambda: method.loss(f[idx], y[idx], cw)
            if use_sam:
                sam_step(method.head, fn, opt)
            else:
                opt.zero_grad(); fn().backward(); opt.step()
    method.head.eval()


def federated_train(kind, clients, rounds, local_epochs, lr, bs, aggregator="fedavg", use_sam=False, seed=0):
    """clients is a list of (features, labels). Returns the global Method."""
    global_m = Method(kind, seed=seed)
    gstate = trainable_state(global_m.head)
    server = FedOptServer() if aggregator == "fedopt" else None
    for _ in range(rounds):
        states, sizes = [], []
        for f, y in clients:
            m = copy.deepcopy(global_m)
            load_trainable(m.head, gstate)
            local_train(m, f, y, local_epochs, lr, bs, use_sam)
            states.append(trainable_state(m.head))
            sizes.append(len(y))
        if aggregator == "fedavg":
            gstate = fedavg(states, sizes)
        elif aggregator == "fedmedian":
            gstate = fedmedian(states)
        else:
            gstate = server(gstate, states, sizes)
    load_trainable(global_m.head, gstate)
    if kind == "metric":
        global_m.set_support(torch.cat([f for f, _ in clients]), torch.cat([y for _, y in clients]))
    return global_m


def task1(global_m, f_test, y_test):
    pred = global_m.predict(f_test).numpy()
    return pred, macro_f1(y_test, pred), expected_cost(y_test, pred)


def task2(global_m, clients, tests, epochs, lr, bs, use_sam=False):
    """Fine tune a copy of the global model on each center, test on that center, average metrics."""
    out = []
    for (f, y), (ft, yt) in zip(clients, tests):
        m = copy.deepcopy(global_m)
        local_train(m, f, y, epochs, lr, bs, use_sam)
        if m.kind == "metric":
            m.set_support(f, y)
        pred = m.predict(ft).numpy()
        out.append((pred, macro_f1(yt, pred), expected_cost(yt, pred)))
    return out, np.mean([o[1] for o in out]), np.mean([o[2] for o in out])

# ---------------------------------------------------------------------------
# 9. Smoke test
# ---------------------------------------------------------------------------

# Label counts per grade 0..5, read from Fig. 2 of the paper
TRAIN_COUNTS = {1: [9, 5, 17, 3, 6, 0], 2: [0, 7, 3, 10, 12, 1], 3: [0, 13, 13, 36, 16, 2], 4: [0] * 6}
TEST_COUNTS = {1: [2, 1, 4, 1, 2, 0], 2: [1, 2, 1, 2, 3, 0], 3: [1, 4, 3, 9, 4, 1], 4: [0, 1, 5, 11, 11, 1]}

# Task 1 confusion matrices on Center 4, rows true grade, columns predicted grade (Fig. 4)
CM_TASK1 = {
    "Santhi": [[0] * 6, [0, 0, 0, 1, 0, 0], [0, 1, 1, 3, 0, 0], [0, 1, 2, 8, 0, 0], [0, 0, 1, 5, 5, 0], [0, 0, 0, 0, 1, 0]],
    "Camma": [[0] * 6, [1, 0, 0, 0, 0, 0], [5, 0, 0, 0, 0, 0], [10, 0, 0, 0, 1, 0], [9, 0, 0, 0, 2, 0], [1, 0, 0, 0, 0, 0]],
    "Elbflorenz": [[0] * 6, [0, 0, 1, 0, 0, 0], [0, 0, 5, 0, 0, 0], [0, 0, 11, 0, 0, 0], [0, 0, 10, 0, 1, 0], [0, 0, 1, 0, 0, 0]],
}


def labels_from_counts(counts):
    return np.repeat(np.arange(C), counts)


def cases_from_cm(cm):
    """Expand a confusion matrix into aligned (true, pred) case lists, true grade ascending."""
    y, p = [], []
    for i, row in enumerate(cm):
        for j, n in enumerate(row):
            y += [i] * n
            p += [j] * n
    return np.array(y), np.array(p)


def synthetic_videos(labels, center, n_frames=40, size=16, seed=0):
    """Grade shows up as a reddish, textured patch that is strongest near the middle frames.
    Each center has its own color cast and brightness, and Center 4 differs most."""
    g = torch.Generator().manual_seed(seed + 97 * center)
    cast = {1: (0.00, 0.00), 2: (0.06, -0.04), 3: (-0.05, 0.05), 4: (0.12, -0.10)}[center]
    vids = []
    for y in labels:
        base = 0.35 + 0.05 * torch.randn(n_frames, 3, size, size, generator=g)
        t = torch.linspace(-1, 1, n_frames)
        strength = torch.exp(-4 * t ** 2)[:, None, None]
        patch = (0.06 * y + 0.02 * torch.randn(1, generator=g)) * strength
        base[:, 0, 4:12, 4:12] += patch
        base[:, 1, 4:12, 4:12] -= 0.5 * patch
        base[:, 0] += cast[0]
        base += cast[1]
        vids.append(base.clamp(0, 1))
    return torch.stack(vids)


if __name__ == "__main__":
    np.set_printoptions(precision=4, suppress=True)

    # (a) Reference classifiers from the published label counts (compare with Table H.25)
    print("== Reference classifiers, Task 1 Center 4 ==")
    y4 = labels_from_counts(TEST_COUNTS[4])
    pooled = np.concatenate([labels_from_counts(TRAIN_COUNTS[c]) for c in (1, 2, 3)])
    g_mode = mode(pooled)
    f1, ec = trivial_constant(y4, g_mode)
    print(f"TrivialGlobal (grade {g_mode})  F1 {100 * f1:.2f}  EC {100 * ec:.2f}    paper 9.17 / 13.79")
    prior = np.bincount(pooled, minlength=C) / len(pooled)
    (fu, fl, fh), (eu, el, eh) = random_reference(y4, np.full(C, 1 / C), draws=10000)
    print(f"RandomUniform  F1 {100 * fu:.2f} [{100 * fl:.1f} to {100 * fh:.1f}]  EC {100 * eu:.2f}    paper 11.91 / 33.55")
    (fu, fl, fh), (eu, el, eh) = random_reference(y4, prior, draws=10000)
    print(f"RandomPrior    F1 {100 * fu:.2f} [{100 * fl:.1f} to {100 * fh:.1f}]  EC {100 * eu:.2f}    paper 14.30 / 25.52")
    print(f"macro F1 ceiling on Center 4 (grade 0 absent): {100 * f1_ceiling(y4):.2f}")
    for c in (1, 2, 3):
        yc = labels_from_counts(TEST_COUNTS[c])
        lm = mode(labels_from_counts(TRAIN_COUNTS[c]))
        f1, ec = trivial_constant(yc, lm)
        print(f"TrivialLocal Center {c} (grade {lm})  F1 {100 * f1:.2f}  EC {100 * ec:.2f}")

    # (b) Task 1 scores recomputed from the published confusion matrices (compare with Table 3)
    print("\n== Task 1 scores from Fig. 4 confusion matrices ==")
    preds = {}
    for team, cm in CM_TASK1.items():
        y, p = cases_from_cm(cm)
        assert (y == y4).all()
        preds[team] = p
        print(f"{team:11s} F1 {100 * macro_f1(y, p):.2f}  EC {100 * expected_cost(y, p):.2f}")
    print("paper      Santhi 23.03 / 12.41   Camma 4.76 / 57.24   Elbflorenz 7.83 / 24.14")

    # (c) Bootstrap rank stability and permutation tests on those reconstructed cases.
    # Predictions are aligned by true grade only, so cross method pairing within a grade is arbitrary.
    names, rk, wins = bootstrap_ranks(y4, preds, "ec", iters=2000)
    print("\nBootstrap EC rank frequencies (rank 1, 2, 3):")
    for n, r in zip(names, rk):
        print(f"  {n:11s} {r}")
    for a, b in [("Camma", "Santhi"), ("Elbflorenz", "Santhi"), ("Camma", "Elbflorenz")]:
        for met in ("ec", "f1"):
            obs, p = paired_permutation(y4, preds[a], preds[b], met, perms=5000)
            print(f"  permutation {a} vs {b} {met.upper():2s}  diff {100 * obs:+6.2f} pp  p = {p:.4f}")

    # (d) Communication costs (compare with Table 8)
    print("\n== Communication cost ==")
    for team, nbytes, rounds in [("Santhi", 354_604_056, 20), ("Elbflorenz", 348_852_152, 50), ("Camma", 96_344_056, 10)]:
        print(f"{team:11s} {fl_cost(nbytes, 3, rounds) / 2**30:8.2f} GiB")
    print(f"{'SL baseline':11s} {swarm_cost(357_405_880, 3, 100) / 2**30:8.2f} GiB   (the paper labels these binary units GB)")
    lora_bytes = 969_000 * 4
    print(f"LoRA only exchange for ViViT+LoRA over 25 rounds {fl_cost(lora_bytes, 3, 25) / 2**30:.3f} GiB, "
          f"{100 * lora_bytes / 354_604_056:.2f} percent of a full model per round")

    # (e) End to end federated run on synthetic videos with the paper's label distribution
    print("\n== Synthetic federated run ==")
    torch.manual_seed(0)
    train = {c: (synthetic_videos(labels_from_counts(TRAIN_COUNTS[c]), c, seed=1), labels_from_counts(TRAIN_COUNTS[c])) for c in (1, 2, 3)}
    test = {c: (synthetic_videos(labels_from_counts(TEST_COUNTS[c]), c, seed=2), labels_from_counts(TEST_COUNTS[c])) for c in (1, 2, 3, 4)}
    rng = np.random.default_rng(0)
    configs = [("temporal", "fedavg", False, 1e-3, 4), ("temporal_lora", "fedavg", False, 3e-3, 4),
               ("frame", "fedopt", True, 1e-2, 16), ("metric", "fedmedian", False, 1e-3, 16), ("lstm", "fedavg", False, 1e-3, 4)]
    for kind, agg, sam, lr, bs in configs:
        enc = Method(kind).encoder
        if kind.startswith("frame"):
            sel = sample_frames_equidistant(40, 20)
        else:
            sel = sample_frames_santhi(40, k=12, center=20, half_window=6, rng=rng)
        raw = lambda v: enc(v[:, sel])
        std = federated_standardizer([raw(train[c][0]) for c in (1, 2, 3)])
        feat = lambda v: std(raw(v))
        clients = [(feat(train[c][0]), torch.as_tensor(train[c][1])) for c in (1, 2, 3)]
        tests = [(feat(test[c][0]), test[c][1]) for c in (1, 2, 3)]
        gm = federated_train(kind, clients, rounds=10, local_epochs=2, lr=lr, bs=bs, aggregator=agg, use_sam=sam)
        gm.encoder = enc
        _, f1_1, ec_1 = task1(gm, feat(test[4][0]), test[4][1])
        _, f1_2, ec_2 = task2(gm, clients, tests, epochs=2, lr=lr, bs=bs, use_sam=sam)
        n_train = sum(p.numel() for p in gm.head.parameters() if p.requires_grad)
        print(f"{kind:13s} {agg:9s} trainable {n_train:7d}  Task 1 F1 {100 * f1_1:5.2f} EC {100 * ec_1:5.2f}   Task 2 avg F1 {100 * f1_2:5.2f} EC {100 * ec_2:5.2f}")