"""
Federated averaging (FedAvg) utilities.

This module provides a simple NumPy-based implementation of the FedAvg
aggregation algorithm. It supports equal weighting and custom non-negative
weights that are normalized internally.

Doctests
========

Basic equal-weight averaging across two "clients" with two tensors each
(vector and 2x2 matrix):

>>> A = [
...     np.array([1.0, 2.0]),
...     np.array([[1.0, 2.0], [3.0, 4.0]]),
... ]
>>> B = [
...     np.array([3.0, 4.0]),
...     np.array([[5.0, 6.0], [7.0, 8.0]]),
... ]
>>> eq = federated_average([A, B])
>>> eq[0].tolist()
[2.0, 3.0]
>>> eq[1].tolist()
[[3.0, 4.0], [5.0, 6.0]]

Weighted averaging with weights [2, 1] (normalized to [2/3, 1/3]):

>>> w = federated_average(
...     [A, B],
...     weights=np.array([2.0, 1.0]),
... )
>>> w[0].tolist()
[1.6666666666666665, 2.6666666666666665]
>>> w[1].tolist()
[[2.333333333333333, 3.333333333333333], [4.333333333333333, 5.333333333333333]]

Error cases:

- No clients

>>> federated_average([])  # doctest: +ELLIPSIS
Traceback (most recent call last):
...
ValueError: client_models must be a non-empty list

- Mismatched number of tensors per client

>>> C = [np.array([1.0, 2.0])]  # only one tensor
>>> federated_average([A, C])  # doctest: +ELLIPSIS
Traceback (most recent call last):
...
ValueError: All clients must have the same number of tensors

- Mismatched tensor shapes across clients

>>> C2 = [
...     np.array([1.0, 2.0]),
...     np.array([[1.0, 2.0]]),
... ]  # second tensor has different shape
>>> federated_average([A, C2])  # doctest: +ELLIPSIS
Traceback (most recent call last):
...
ValueError: Client 2 tensor shape (1, 2) does not match (2, 2)

- Invalid weights: negative or wrong shape or zero-sum

>>> federated_average([A, B], weights=np.array([1.0, -1.0]))  # doctest: +ELLIPSIS
Traceback (most recent call last):
...
ValueError: weights must be non-negative

>>> federated_average([A, B], weights=np.array([0.0, 0.0]))  # doctest: +ELLIPSIS
Traceback (most recent call last):
...
ValueError: weights must sum to a positive value

>>> federated_average(
...     [A, B],
...     weights=np.array([1.0, 2.0, 3.0]),
... )  # doctest: +ELLIPSIS
Traceback (most recent call last):
...
ValueError: weights must have shape (2,)
"""

from __future__ import annotations

from collections.abc import Sequence

import numpy as np


def _validate_clients(client_models: Sequence[Sequence[np.ndarray]]) -> None:
    if not client_models:
        raise ValueError("client_models must be a non-empty list")
    # Ensure all clients have same number of layers and shapes
    ref_shapes = [tuple(arr.shape) for arr in client_models[0]]
    for idx, cm in enumerate(client_models, start=1):
        if len(cm) != len(ref_shapes):
            raise ValueError("All clients must have the same number of tensors")
        for s_ref, arr in zip(ref_shapes, cm):
            if tuple(arr.shape) != s_ref:
                msg = (
                    f"Client {idx} tensor shape {tuple(arr.shape)} "
                    f"does not match {s_ref}"
                )
                raise ValueError(msg)


def _normalize_weights(weights: np.ndarray, num_clients: int) -> np.ndarray:
    if weights.shape != (num_clients,):
        msg = f"weights must have shape ({num_clients},)"
        raise ValueError(msg)
    if np.any(weights < 0):
        raise ValueError("weights must be non-negative")
    total = float(weights.sum())
    if total <= 0.0:
        raise ValueError("weights must sum to a positive value")
    return weights / total


def federated_average(
    client_models: Sequence[Sequence[np.ndarray]],
    weights: np.ndarray | None = None,
) -> list[np.ndarray]:
    """Compute the weighted average of clients' model tensors.

    Parameters
    ----------
    client_models : Sequence[Sequence[np.ndarray]]
        A list of clients, each being a sequence of NumPy arrays (tensors).
        All clients must have the same number of tensors with identical shapes.
    weights : np.ndarray | None, optional
        A 1-D array of non-negative weights, one per client. If None,
        equal weighting is used. Weights are normalized to sum to 1.

    Returns
    -------
    list[np.ndarray]
        The list of aggregated tensors with the same shapes as the inputs.
    """
    _validate_clients(client_models)
    num_clients = len(client_models)

    if weights is None:
        weights_n = np.full((num_clients,), 1.0 / num_clients, dtype=float)
    else:
        weights = np.asarray(weights, dtype=float)
        weights_n = _normalize_weights(weights, num_clients)

    num_tensors = len(client_models[0])
    aggregated: list[np.ndarray] = []
    for t_idx in range(num_tensors):
        # Stack the t_idx-th tensor from each client into shape (num_clients, ...)
        stacked = np.stack([np.asarray(cm[t_idx]) for cm in client_models], axis=0)
        # Weighted sum across clients axis=0
        # np.tensordot weights of shape (n,) with stacked of shape (n, *dims)
        agg = np.tensordot(weights_n, stacked, axes=(0, 0))
        aggregated.append(np.asarray(agg))

    return aggregated


if __name__ == "__main__":
    import doctest

    doctest.testmod()