"""Compact 1D-CNN detector shared by both topologies (RQ4).

Defined once here so star and mesh results are directly comparable. The network
is deliberately small (~10^4 parameters) so it trains on a CPU edge node in
seconds/round; the weave *reports* the exact footprint (params, KB, buffers) so
"how compact can it be" is answered with measured numbers, not claims.

Weight (de)serialization uses plain lists of numpy arrays so the same helpers
feed both Flower (star) and the HTTP peer-exchange (mesh).
"""

from __future__ import annotations

from typing import List

import numpy as np
import torch
import torch.nn as nn


class GridFaultCNN(nn.Module):
    """3-block 1D-CNN over multi-channel power-quality windows.

    Input:  (batch, in_channels, signal_len)
    Output: (batch, num_classes) logits
    """

    def __init__(self, num_classes: int = 8, in_channels: int = 2):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv1d(in_channels, 16, kernel_size=7, padding=3),
            nn.BatchNorm1d(16),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(2),
            nn.Conv1d(16, 32, kernel_size=5, padding=2),
            nn.BatchNorm1d(32),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(2),
            nn.Conv1d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.AdaptiveAvgPool1d(1),  # global pooling -> length-invariant head
        )
        self.head = nn.Linear(64, num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        z = self.features(x)
        z = z.flatten(1)
        return self.head(z)


def build_model(num_classes: int = 8, in_channels: int = 2) -> GridFaultCNN:
    """Factory so callers never construct the class directly."""
    return GridFaultCNN(num_classes=num_classes, in_channels=in_channels)


def get_device(prefer_gpu: bool = True) -> torch.device:
    """Pick cuda when available and requested, else cpu.

    ``experiment.py`` / the notebook set the compute profile; this just honors
    whatever hardware actually landed on the node (a mis-scheduled GPU run still
    runs on CPU).
    """
    if prefer_gpu and torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


# --- weight (de)serialization: list[np.ndarray] in state_dict order ----------

def get_weights(model: nn.Module) -> List[np.ndarray]:
    return [v.detach().cpu().numpy() for v in model.state_dict().values()]


def set_weights(model: nn.Module, weights: List[np.ndarray]) -> None:
    state = model.state_dict()
    new_state = {}
    for (key, ref), w in zip(state.items(), weights):
        # Preserve the reference dtype (BatchNorm's num_batches_tracked is int64).
        new_state[key] = torch.tensor(np.asarray(w)).to(ref.dtype)
    model.load_state_dict(new_state, strict=True)


def average_weights(weight_sets: List[List[np.ndarray]]) -> List[np.ndarray]:
    """Uniform element-wise mean of several weight sets (mesh peer-averaging).

    Integer buffers (e.g. num_batches_tracked) are averaged then re-cast so the
    result round-trips cleanly through ``set_weights``.
    """
    if not weight_sets:
        raise ValueError("average_weights got no weight sets")
    n = len(weight_sets)
    out: List[np.ndarray] = []
    for tensors in zip(*weight_sets):
        ref = np.asarray(tensors[0])
        acc = np.zeros_like(ref, dtype=np.float64)
        for t in tensors:
            acc += np.asarray(t, dtype=np.float64)
        acc /= n
        out.append(acc.astype(ref.dtype))
    return out


def footprint(model: nn.Module) -> dict:
    """Measured RQ4 numbers: parameter count, buffer count, size on disk."""
    n_params = sum(p.numel() for p in model.parameters())
    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    param_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
    buffer_bytes = sum(b.numel() * b.element_size() for b in model.buffers())
    return {
        "params": int(n_params),
        "trainable_params": int(n_trainable),
        "param_bytes": int(param_bytes),
        "buffer_bytes": int(buffer_bytes),
        "size_kb": round((param_bytes + buffer_bytes) / 1024.0, 2),
    }


def payload_bytes(model: nn.Module) -> int:
    """Bytes moved when one full weight set crosses the network (comm cost)."""
    fp = footprint(model)
    return fp["param_bytes"] + fp["buffer_bytes"]


if __name__ == "__main__":
    # Quick footprint report: `python3 -m fl.model`
    import json

    m = build_model()
    print(json.dumps(footprint(m), indent=2))
