"""Tests for deterministic Dirichlet partitioning and manifest reload."""

from __future__ import annotations

import numpy as np
import pytest

from protean_research.data import manifest as manifest_mod
from protean_research.data import partition as partition_mod
from protean_research.data import rare_unseen, splits


def make_labels(num_classes=5, per_class=40):
    return np.repeat(np.arange(num_classes), per_class).astype(np.int64)


def test_partition_same_seed_identical():
    labels = make_labels()
    idx = partition_mod.class_index_map(labels, 5)
    a = partition_mod.dirichlet_partition(idx, 4, 0.3, rng_seed=123)
    b = partition_mod.dirichlet_partition(idx, 4, 0.3, rng_seed=123)
    for x, y in zip(a, b):
        assert np.array_equal(x, y)


def test_partition_different_alpha_differs():
    labels = make_labels()
    idx = partition_mod.class_index_map(labels, 5)
    a = partition_mod.dirichlet_partition(idx, 4, 0.1, rng_seed=7)
    b = partition_mod.dirichlet_partition(idx, 4, 5.0, rng_seed=7)
    assert partition_mod.partition_hash(a) != partition_mod.partition_hash(b)


def test_every_sample_assigned_exactly_once():
    labels = make_labels()
    idx = partition_mod.class_index_map(labels, 5)
    assignments = partition_mod.dirichlet_partition(idx, 4, 0.4, rng_seed=99)
    union = np.sort(np.concatenate(assignments))
    assert np.array_equal(union, np.arange(len(labels)))
    assert len(union) == len(set(union.tolist()))


def test_support_matrix_matches_assignments():
    labels = make_labels()
    idx = partition_mod.class_index_map(labels, 5)
    assignments = partition_mod.dirichlet_partition(idx, 4, 0.4, rng_seed=99)
    support = partition_mod.support_matrix(assignments, labels, 5)
    assert support.shape == (4, 5)
    assert int(support.sum()) == len(labels)
    # Column sums equal per-class counts.
    assert np.array_equal(support.sum(axis=0), np.bincount(labels, minlength=5))


def test_manifest_save_reload_identity(tmp_path):
    labels = make_labels()
    split = splits.stratified_split(labels, "toy", seed=0)
    labels_train = labels[split.train]
    idx = partition_mod.class_index_map(labels_train, 5)
    rng_seed = partition_mod.partition_rng_seed("toy", 0.5, 0)
    assignments = partition_mod.dirichlet_partition(idx, 4, 0.5, rng_seed)
    manifest = manifest_mod.build_manifest(
        dataset_key="toy",
        dataset_name="Toy",
        alpha=0.5,
        seed=0,
        num_clients=4,
        class_names=[f"c{i}" for i in range(5)],
        benign_class="c0",
        labels=labels,
        split=split,
        assignments=assignments,
        rng_seed=rng_seed,
        train_fraction=0.7,
        validation_fraction=0.1,
        test_fraction=0.2,
    )
    path = tmp_path / "toy__alpha0.5__seed0.json"
    manifest_mod.save_manifest(path, manifest)
    loaded = manifest_mod.load_manifest(path)
    split_reload, assignments_reload, support_reload = manifest_mod.reproduce_train_assignments(
        loaded, labels
    )
    assert np.array_equal(split_reload.train, split.train)
    assert partition_mod.partition_hash(assignments_reload) == manifest["partition_hash"]
    assert np.array_equal(support_reload, np.asarray(manifest["support"]))


def test_manifest_detects_corruption(tmp_path):
    labels = make_labels()
    split = splits.stratified_split(labels, "toy", seed=0)
    labels_train = labels[split.train]
    idx = partition_mod.class_index_map(labels_train, 5)
    rng_seed = partition_mod.partition_rng_seed("toy", 0.5, 0)
    assignments = partition_mod.dirichlet_partition(idx, 4, 0.5, rng_seed)
    manifest = manifest_mod.build_manifest(
        dataset_key="toy", dataset_name="Toy", alpha=0.5, seed=0, num_clients=4,
        class_names=[f"c{i}" for i in range(5)], benign_class="c0", labels=labels,
        split=split, assignments=assignments, rng_seed=rng_seed,
        train_fraction=0.7, validation_fraction=0.1, test_fraction=0.2,
    )
    manifest["rng_seed"] = int(manifest["rng_seed"]) + 1
    with pytest.raises(ValueError):
        manifest_mod.reproduce_train_assignments(manifest, labels)


def test_rare_unseen_semantics():
    # Client 0: benign + classes 1,2,3 with supports 10, 5, 2 -> rare = {3,2}.
    support = np.array(
        [
            [50, 10, 5, 2, 0],
            [50, 1, 0, 0, 0],
            [50, 0, 0, 0, 0],
        ]
    )
    rows = rare_unseen.rare_unseen_rows(support, ["b", "c1", "c2", "c3", "c4"], 0, "toy", 0.5, 0)
    client0 = [r for r in rows if r["client_id"] == 0]
    rare = {r["class_id"] for r in client0 if r["rare_flag"]}
    unseen = {r["class_id"] for r in client0 if r["unseen_flag"]}
    assert rare == {2, 3}
    assert unseen == {4}
    # Benign is never rare or unseen.
    assert not any(r["rare_flag"] or r["unseen_flag"] for r in rows if r["class_id"] == 0)
    # Client with a single attack class uses that one.
    assert rare_unseen.rare_class_ids_for_client([50, 1, 0, 0, 0], 0) == [1]
    # Client with no attack class contributes nothing to the rare aggregate.
    assert rare_unseen.rare_class_ids_for_client([50, 0, 0, 0, 0], 0) == []