"""Tests for local prototype construction and global prototype collection."""

from __future__ import annotations

import numpy as np
import pytest

from protean_research.aggregation import AggregationMode
from protean_research.models.prototypes import (
    collect_global_prototypes,
    compute_local_prototypes,
    nearest_prototype_predict,
)


def test_local_prototype_is_mean_embedding():
    embeddings = np.array(
        [[0.0, 0.0], [2.0, 0.0], [10.0, 10.0], [0.0, 4.0]], dtype=np.float32
    )
    labels = np.array([0, 0, 1, 1])
    prototypes, supports = compute_local_prototypes(embeddings, labels, num_classes=3)
    assert supports == {0: 2, 1: 2}
    assert prototypes[0] == pytest.approx([1.0, 0.0])
    assert prototypes[1] == pytest.approx([5.0, 7.0])


def test_missing_class_has_no_prototype():
    embeddings = np.array([[1.0, 1.0]], dtype=np.float32)
    labels = np.array([2])
    prototypes, supports = compute_local_prototypes(embeddings, labels, num_classes=4)
    assert set(prototypes) == {2}
    assert 0 not in prototypes
    assert 1 not in prototypes


def test_collect_global_original_is_equal_mean_over_contributors():
    client_prototypes = [{0: np.array([0.0, 2.0])}, {0: np.array([4.0, 4.0])}]
    client_supports = [{0: 1}, {0: 100}]
    result = collect_global_prototypes(
        client_prototypes, client_supports, num_classes=1, mode=AggregationMode.ORIGINAL_PROTEAN
    )
    assert result[0] == pytest.approx([2.0, 3.0])


def test_collect_global_omits_missing_class():
    client_prototypes = [{0: np.array([1.0, 1.0])}, {}]
    client_supports = [{0: 3}, {}]
    result = collect_global_prototypes(
        client_prototypes, client_supports, num_classes=2, mode=AggregationMode.ORIGINAL_PROTEAN
    )
    assert set(result) == {0}


def test_nearest_prototype_predict():
    prototypes = {0: np.array([0.0, 0.0]), 1: np.array([10.0, 10.0])}
    embeddings = np.array([[0.1, 0.1], [9.0, 9.0]])
    predictions = nearest_prototype_predict(embeddings, prototypes)
    assert predictions.tolist() == [0, 1]