import random
from math import isclose

from one_pass_stats import ParallelDescriptiveStats


def test_empty_stats():
    s = ParallelDescriptiveStats()
    assert s.count == 0
    assert s.mean == 0.0
    assert s.variance() == 0.0
    assert s.stddev() == 0.0


def test_single_value():
    s = ParallelDescriptiveStats()
    s.add_value(42.0)
    assert s.count == 1
    assert s.mean == 42.0
    assert s.min == 42.0
    assert s.max == 42.0
    assert s.variance() == 0.0


def test_known_values():
    s = ParallelDescriptiveStats()
    for v in [2, 4, 4, 4, 5, 5, 7, 9]:
        s.add_value(v)

    assert s.count == 8
    assert s.mean == 5.0
    assert s.min == 2
    assert s.max == 9
    # sample variance = 4.571..., population variance = 4.0
    assert isclose(s.variance(sample=True), 32 / 7)
    assert isclose(s.variance(sample=False), 4.0)
    assert isclose(s.stddev(sample=False), 2.0)


def test_merge_equals_single_pass():
    all_values = list(range(100))
    single = ParallelDescriptiveStats()
    for v in all_values:
        single.add_value(v)

    a = ParallelDescriptiveStats()
    for v in all_values[:40]:
        a.add_value(v)
    b = ParallelDescriptiveStats()
    for v in all_values[40:]:
        b.add_value(v)
    merged = a.merge(b)

    assert merged.count == single.count
    assert isclose(merged.mean, single.mean)
    assert isclose(merged.variance(), single.variance())
    assert merged.min == single.min
    assert merged.max == single.max


def test_merge_empty_with_nonempty():
    s = ParallelDescriptiveStats()
    for v in [1, 2, 3]:
        s.add_value(v)

    merged = ParallelDescriptiveStats().merge(s)
    assert merged.count == 3
    assert isclose(merged.mean, 2.0)


def test_merge_two_empty():
    merged = ParallelDescriptiveStats().merge(ParallelDescriptiveStats())
    assert merged.count == 0
    assert merged.mean == 0.0


def test_multi_way_merge():
    """Merging many small chunks matches a single pass over all data."""
    random.seed(42)
    all_values = [random.gauss(100, 25) for _ in range(1000)]

    single = ParallelDescriptiveStats()
    for v in all_values:
        single.add_value(v)

    chunk_size = 50
    chunks: list[ParallelDescriptiveStats] = []
    for i in range(0, len(all_values), chunk_size):
        chunk = ParallelDescriptiveStats()
        for v in all_values[i : i + chunk_size]:
            chunk.add_value(v)
        chunks.append(chunk)

    merged = chunks[0]
    for chunk in chunks[1:]:
        merged = merged.merge(chunk)

    assert merged.count == single.count
    assert isclose(merged.mean, single.mean, rel_tol=1e-12)
    assert isclose(merged.variance(), single.variance(), rel_tol=1e-12)
    assert isclose(merged.min, single.min)
    assert isclose(merged.max, single.max)


def test_backwards_compat_aliases():
    s = ParallelDescriptiveStats()
    s.addValue(10)
    s.addValue(20)
    assert s.mean == 15.0
    assert isclose(s.getVariance(), s.variance())
    assert isclose(s.getStddev(), s.stddev())
