"""Shared local-training / evaluation loop used by star and mesh.

Keeping this in one place means the star Flower client and the mesh gossip peer
train and evaluate identically, so cross-topology comparisons are fair (the whole
point of RQ1).
"""

from __future__ import annotations

from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

from . import data as fldata
from . import metrics as flmetrics


def make_loaders(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_test: np.ndarray,
    y_test: np.ndarray,
    batch_size: int,
) -> Tuple[DataLoader, DataLoader]:
    train_ds = TensorDataset(torch.from_numpy(X_train), torch.from_numpy(y_train))
    test_ds = TensorDataset(torch.from_numpy(X_test), torch.from_numpy(y_test))
    train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True, drop_last=False)
    test_dl = DataLoader(test_ds, batch_size=256, shuffle=False)
    return train_dl, test_dl


def balanced_class_weights(y: np.ndarray, num_classes: int) -> np.ndarray:
    """sklearn-style 'balanced' weights: ``n / (num_classes * count_c)`` per class.

    Up-weights rare classes so an imbalanced dataset (e.g. vsb, ~6% partial-discharge)
    doesn't collapse to always predicting the majority. Classes absent from ``y`` get
    weight 0 (they never contribute to the loss anyway).
    """
    counts = np.bincount(np.asarray(y, dtype=np.int64), minlength=num_classes).astype(np.float64)
    n = counts.sum()
    with np.errstate(divide="ignore", invalid="ignore"):
        w = np.where(counts > 0, n / (num_classes * counts), 0.0)
    return w.astype(np.float32)


def train_local(
    model: nn.Module,
    loader: DataLoader,
    device: torch.device,
    epochs: int,
    lr: float,
    proximal_mu: float = 0.0,
    class_weights: Optional[np.ndarray] = None,
) -> int:
    """One local training pass (``epochs`` epochs). Returns #samples seen.

    ``proximal_mu`` > 0 adds the FedProx proximal term ``mu/2 * ||w - w_global||^2``
    anchored at the weights the round started from (the received global model).
    ``class_weights`` (length num_classes) weights the cross-entropy loss per class —
    pass :func:`balanced_class_weights` for imbalanced data, or ``None`` for uniform.
    """
    model.to(device)
    model.train()
    opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
    weight_t = (torch.as_tensor(class_weights, dtype=torch.float32, device=device)
                if class_weights is not None else None)
    loss_fn = nn.CrossEntropyLoss(weight=weight_t)

    global_params: Optional[List[torch.Tensor]] = None
    if proximal_mu > 0:
        global_params = [p.detach().clone() for p in model.parameters()]

    n_seen = 0
    for _ in range(epochs):
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            opt.zero_grad()
            loss = loss_fn(model(xb), yb)
            if global_params is not None:
                prox = sum(
                    torch.sum((p - g) ** 2)
                    for p, g in zip(model.parameters(), global_params)
                )
                loss = loss + (proximal_mu / 2.0) * prox
            loss.backward()
            opt.step()
            n_seen += xb.size(0)
    return n_seen


@torch.no_grad()
def evaluate(model: nn.Module, loader: DataLoader, device: torch.device,
             spec: "fldata.DatasetSpec | None" = None) -> Tuple[float, dict]:
    """Return ``(avg_loss, metrics_dict)`` over the loader.

    ``spec`` supplies the dataset's class geometry (num_classes + which class is
    the headline fault); defaults to pq_synth's for backward compatibility.
    """
    sp = spec or fldata.PQ_SYNTH_SPEC
    model.to(device)
    model.eval()
    loss_fn = nn.CrossEntropyLoss(reduction="sum")
    total_loss, total_n = 0.0, 0
    preds, trues = [], []
    for xb, yb in loader:
        xb, yb = xb.to(device), yb.to(device)
        logits = model(xb)
        total_loss += float(loss_fn(logits, yb).item())
        total_n += xb.size(0)
        preds.append(logits.argmax(1).cpu().numpy())
        trues.append(yb.cpu().numpy())
    y_pred = np.concatenate(preds) if preds else np.array([], dtype=np.int64)
    y_true = np.concatenate(trues) if trues else np.array([], dtype=np.int64)
    m = flmetrics.classification_metrics(
        y_true, y_pred, sp.num_classes, sp.fault_class, sp.normal_class
    )
    avg_loss = total_loss / total_n if total_n else 0.0
    m["loss"] = round(avg_loss, 5)
    return avg_loss, m
