import random
from math import isclose

from one_pass_stats import ParallelCovariance


def test_perfect_positive_correlation():
    cov = ParallelCovariance()
    for i in range(10):
        cov.add_pair(float(i), float(i) * 2)

    assert isclose(cov.pearson(), 1.0)


def test_perfect_negative_correlation():
    cov = ParallelCovariance()
    for i in range(10):
        cov.add_pair(float(i), -float(i) * 3)

    assert isclose(cov.pearson(), -1.0)


def test_known_covariance():
    # x = [1,2,3,4,5], y = [2,4,6,8,10] => y = 2x
    # sample cov = 5.0, population cov = 4.0
    cov = ParallelCovariance()
    for x, y in [(1, 2), (2, 4), (3, 6), (4, 8), (5, 10)]:
        cov.add_pair(x, y)

    assert isclose(cov.covariance(sample=True), 5.0)
    assert isclose(cov.covariance(sample=False), 4.0)


def test_single_pair():
    cov = ParallelCovariance()
    cov.add_pair(1.0, 2.0)
    assert cov.covariance() == 0.0


def test_merge_equals_single_pass():
    xs = list(range(20))
    ys = [x * 1.5 + 3 for x in xs]

    single = ParallelCovariance()
    for x, y in zip(xs, ys):
        single.add_pair(float(x), float(y))

    a = ParallelCovariance()
    for x, y in zip(xs[:8], ys[:8]):
        a.add_pair(float(x), float(y))
    b = ParallelCovariance()
    for x, y in zip(xs[8:], ys[8:]):
        b.add_pair(float(x), float(y))
    merged = a.merge(b)

    assert merged.x.count == single.x.count
    assert isclose(merged.covariance(), single.covariance(), rel_tol=1e-10)
    assert isclose(merged.pearson(), single.pearson(), rel_tol=1e-10)


def test_multi_way_merge():
    random.seed(99)
    xs = [random.gauss(0, 10) for _ in range(200)]
    ys = [x * 0.7 + random.gauss(0, 3) for x in xs]

    single = ParallelCovariance()
    for x, y in zip(xs, ys):
        single.add_pair(x, y)

    chunk_size = 25
    chunks: list[ParallelCovariance] = []
    for i in range(0, len(xs), chunk_size):
        c = ParallelCovariance()
        for x, y in zip(xs[i : i + chunk_size], ys[i : i + chunk_size]):
            c.add_pair(x, y)
        chunks.append(c)

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

    assert isclose(merged.covariance(), single.covariance(), rel_tol=1e-10)
    assert isclose(merged.pearson(), single.pearson(), rel_tol=1e-10)
