import numpy as np

import pytest

from femr.featurizers.utils import OnlineStatistics

def _assert_correct_stats(stat: OnlineStatistics, values: list):

    TOLERANCE = 1e-6  # Allow for some floating point error

    true_mean = np.mean(values)

    true_sample_variance = np.var(values, ddof=1)

    true_m2 = true_sample_variance * (len(values) - 1)

    assert stat.current_count == len(values), f"{stat.current_count} != {len(values)}"

    assert np.isclose(stat.mean(), true_mean), f"{stat.mean()} != {true_mean}"

    assert np.isclose(

        stat.variance(), true_sample_variance, atol=TOLERANCE

    ), f"{stat.variance()} != {true_sample_variance}"

    assert np.isclose(stat.current_M2, true_m2, atol=TOLERANCE), f"{stat.current_M2} != {true_m2}"

def test_add():

    # Test adding things to the statistics

    def _run_test(values):

        stat = OnlineStatistics()

        for i in values:

            stat.add(i)

        _assert_correct_stats(stat, values)

    # Positive integers

    _run_test(range(51))

    _run_test(range(10, 10000, 3))

    # Negative integers

    _run_test(range(-400, -300))

    # Positive/negative integers

    _run_test(list(range(4, 900, 2)) + list(range(-1000, -300, 7)))

    _run_test(list(range(-100, 100, 7)) + list(range(-100, 100, 2)))

    # Decimals

    _run_test(np.linspace(0, 1, 100))

    _run_test(np.logspace(-100, 3, 100))

    # Small lists

    _run_test([0, 1])

    _run_test([-1, 1])

def test_constructor():

    # Test default

    stat = OnlineStatistics()

    assert stat.current_count == 0

    assert stat.current_mean == stat.mean() == 0

    assert stat.current_M2 == 0

    # Test explicitly setting args

    stat = OnlineStatistics(current_count=1, current_mean=2, current_variance=3)

    assert stat.current_count == 1

    assert stat.current_mean == stat.mean() == 2

    assert stat.current_M2 == 0

    # Test M2

    stat = OnlineStatistics(current_count=10, current_mean=20, current_variance=30)

    assert stat.current_count == 10

    assert stat.current_mean == 20

    assert stat.current_M2 == 30 * (10 - 1)

    # Test getters/setters

    stat = OnlineStatistics(current_count=10, current_mean=20, current_variance=30)

    assert stat.mean() == 20

    assert stat.variance() == 30

    assert stat.standard_deviation() == np.sqrt(30)

    # Test fail cases

    with pytest.raises(ValueError) as _:

        # Negative count

        stat = OnlineStatistics(current_count=-1, current_mean=2, current_variance=3)

    with pytest.raises(ValueError) as _:

        # Negative variance

        stat = OnlineStatistics(current_count=1, current_mean=2, current_variance=-3)

    with pytest.raises(ValueError) as _:

        # Positive variance with 0 count

        stat = OnlineStatistics(current_count=0, current_mean=2, current_variance=1)

    with pytest.raises(ValueError) as _:

        # Can only compute variance with >1 observation

        stat = OnlineStatistics()

        stat.add(1)

        stat.variance()

def test_merge_pair():

    # Simulate two statistics being merged via `merge_pair``

    stat1 = OnlineStatistics()

    values1 = list(range(-300, 300, 4)) + list(range(400, 450))

    for i in values1:

        stat1.add(i)

    stat2 = OnlineStatistics()

    values2 = list(range(100, 150))

    for i in values2:

        stat2.add(i)

    merged_stat = OnlineStatistics.merge_pair(stat1, stat2)

    merged_stat_values = values1 + values2

    _assert_correct_stats(merged_stat, merged_stat_values)

def test_merge():

    # Simulate parallel statistics being merged via `merge`

    stats = []

    values = [

        np.linspace(-100, 100, 50),

        np.linspace(100, 200, 50),

        np.linspace(100, 150, 100),

        np.linspace(-10, 100, 100),

        np.linspace(10, 200, 3),

    ]

    for i in range(len(values)):

        stat = OnlineStatistics()

        for v in values[i]:

            stat.add(v)

        stats.append(stat)

    merged_stat = OnlineStatistics.merge(stats)

    _assert_correct_stats(merged_stat, np.concatenate(values))