#!/usr/bin/env python3

import math
from types import SimpleNamespace

from bornagain.ba_view3d import _particle_matches_mixture_choices, _pick_active_variants


def active_name(seed, site, particles, weights):
    """Returns the single particle name active at a synthetic layout site."""
    choices = _pick_active_variants(seed, 0, float(site), 0.0, 0.0, weights)
    active = [
        name for name, particle in particles.items()
        if _particle_matches_mixture_choices(particle, choices)
    ]
    assert len(active) == 1
    return active[0]


def assert_frequency(name, actual, expected, n):
    """Checks sampled counts against a binomial three-sigma interval."""
    sigma = math.sqrt(n * expected * (1.0 - expected))
    assert abs(actual - n * expected) <= 3.0 * sigma, (
        f"{name}: got {actual}, expected {n * expected} +/- {3.0 * sigma}"
    )


def test_nested_mixture_frequencies():
    # Mixture(A, Mixture(B, C)) with outer weights (0.5, 0.5)
    # and inner weights (0.4, 0.6) gives P(A)=0.5, P(B)=0.2, P(C)=0.3.
    weights = [[0.5, 0.5], [0.4, 0.6]]
    particles = {
        "A": SimpleNamespace(mixture_groups=[0], mixture_variants=[0]),
        "B": SimpleNamespace(mixture_groups=[0, 1], mixture_variants=[1, 0]),
        "C": SimpleNamespace(mixture_groups=[0, 1], mixture_variants=[1, 1]),
    }

    n = 5000
    counts = {"A": 0, "B": 0, "C": 0}
    for site in range(n):
        counts[active_name(123, site, particles, weights)] += 1

    assert_frequency("A", counts["A"], 0.5, n)
    assert_frequency("B", counts["B"], 0.2, n)
    assert_frequency("C", counts["C"], 0.3, n)


def test_mixture_choices_are_deterministic_by_seed_and_position():
    weights = [[0.7, 0.3], [0.2, 0.8]]

    first = [
        _pick_active_variants(17, 2, float(site), 1.25, -3.5, weights)
        for site in range(100)
    ]
    second = [
        _pick_active_variants(17, 2, float(site), 1.25, -3.5, weights)
        for site in range(100)
    ]
    changed_seed = [
        _pick_active_variants(18, 2, float(site), 1.25, -3.5, weights)
        for site in range(100)
    ]

    assert first == second
    assert any(a != b for a, b in zip(first, changed_seed))


if __name__ == "__main__":
    test_nested_mixture_frequencies()
    test_mixture_choices_are_deterministic_by_seed_and_position()
