import numpy as np
from dynasor.core.time_averager import TimeAverager, truncate_nans


def test_output_shape():
    """get_average_all returns shape (array_length, time_window)."""
    averager = TimeAverager(time_window=7, array_length=13)
    assert averager.get_average_all().shape == (13, 7)


def test_unsampled_slots_are_nan():
    """All time lags without any added samples return NaN."""
    averager = TimeAverager(time_window=5, array_length=3)
    result = averager.get_average_all()
    assert np.all(np.isnan(result))


def test_get_average_at_timelag_unsampled_returns_nan():
    """get_average_at_timelag returns NaN for a slot with no samples."""
    averager = TimeAverager(time_window=4, array_length=6)
    result = averager.get_average_at_timelag(2)
    assert np.all(np.isnan(result))
    assert result.shape == (6,)


def test_single_sample_equals_itself():
    """Adding one sample at a time lag returns that sample exactly."""
    rng = np.random.default_rng(0)
    averager = TimeAverager(time_window=5, array_length=10)
    sample = rng.random(10)
    averager.add_sample(time_lag=0, sample=sample)
    assert np.allclose(averager.get_average_at_timelag(0), sample)


def test_average_of_two_samples():
    """Two samples at the same time lag are averaged correctly."""
    rng = np.random.default_rng(1)
    averager = TimeAverager(time_window=5, array_length=8)
    s1 = rng.random(8)
    s2 = rng.random(8)
    averager.add_sample(time_lag=3, sample=s1)
    averager.add_sample(time_lag=3, sample=s2)
    assert np.allclose(averager.get_average_at_timelag(3), (s1 + s2) / 2)


def test_average_of_many_samples():
    """N samples at the same time lag equal their arithmetic mean."""
    rng = np.random.default_rng(2)
    N = 50
    array_length = 12
    averager = TimeAverager(time_window=3, array_length=array_length)
    samples = rng.random((N, array_length))
    for s in samples:
        averager.add_sample(time_lag=1, sample=s)
    assert np.allclose(averager.get_average_at_timelag(1), samples.mean(axis=0))


def test_different_time_lags_are_independent():
    """Samples added to one time lag do not affect other time lags."""
    rng = np.random.default_rng(3)
    averager = TimeAverager(time_window=6, array_length=5)
    s0 = rng.random(5)
    s2 = rng.random(5)
    averager.add_sample(time_lag=0, sample=s0)
    averager.add_sample(time_lag=2, sample=s2)
    assert np.allclose(averager.get_average_at_timelag(0), s0)
    assert np.allclose(averager.get_average_at_timelag(2), s2)
    # All other lags remain NaN
    for t in [1, 3, 4, 5]:
        assert np.all(np.isnan(averager.get_average_at_timelag(t)))


def test_get_average_at_timelag_consistent_with_get_average_all():
    """get_average_at_timelag and get_average_all return the same values."""
    rng = np.random.default_rng(4)
    time_window = 7
    array_length = 9
    averager = TimeAverager(time_window=time_window, array_length=array_length)
    for t in [0, 2, 4, 6]:
        averager.add_sample(time_lag=t, sample=rng.random(array_length))
    all_averages = averager.get_average_all()
    for t in range(time_window):
        assert np.allclose(averager.get_average_at_timelag(t),
                           all_averages[:, t], equal_nan=True)


def test_time_window_of_one():
    """TimeAverager with time_window=1 works for the single slot."""
    rng = np.random.default_rng(5)
    averager = TimeAverager(time_window=1, array_length=4)
    s = rng.random(4)
    averager.add_sample(time_lag=0, sample=s)
    assert averager.get_average_all().shape == (4, 1)
    assert np.allclose(averager.get_average_at_timelag(0), s)


def test_accumulation_is_exact():
    """Repeated additions produce the exact running mean with no drift."""
    array_length = 10
    averager = TimeAverager(time_window=2, array_length=array_length)
    # Use constant samples so the mean is trivially known
    value = np.full(array_length, 3.14159)
    N = 100
    for _ in range(N):
        averager.add_sample(time_lag=0, sample=value)
    assert np.allclose(averager.get_average_at_timelag(0), value)


# Legacy test preserved for regression coverage
def test_time_averager():
    window = 11
    array_size = 20

    rng = np.random.default_rng(42)
    averager = TimeAverager(window, array_size)

    # add a single sample
    sample1 = rng.random(array_size)
    averager.add_sample(time_lag=0, sample=sample1)
    average_array = averager.get_average_all()

    assert average_array.shape == (array_size, window)
    assert np.allclose(sample1, average_array[:, 0])
    for t in range(1, window):
        assert np.all(np.isnan(average_array[:, t]))

    # add additional sample
    sample2 = rng.random(array_size)
    sample3 = rng.random(array_size)

    averager.add_sample(time_lag=0, sample=sample2)
    averager.add_sample(time_lag=5, sample=sample3)
    average_array = averager.get_average_all()

    assert np.allclose(average_array[:, 0], (sample1 + sample2) / 2.0)
    assert np.allclose(average_array[:, 5], sample3)


def test_truncate_nans():
    # setup
    N_times = 10
    N_q = 2
    F = TimeAverager(N_times, N_q)

    # add dummy time samples for all q-points
    F.add_sample(0, [1.0] * N_q)
    F.add_sample(1, [0.84] * N_q)
    F.add_sample(2, [0.75] * N_q)
    F.add_sample(2, [0.70] * N_q)
    F.add_sample(2, [0.72] * N_q)
    F.add_sample(2, [0.71] * N_q)
    F.add_sample(3, [0.61] * N_q)
    F.add_sample(4, [0.55] * N_q)

    # compute average
    F_ave = F.get_average_all()

    F_ave_ref = np.array([
        [1., 0.84, 0.72, 0.61, 0.55,  np.nan,  np.nan,  np.nan,  np.nan, np.nan],
        [1., 0.84, 0.72, 0.61, 0.55,  np.nan,  np.nan,  np.nan,  np.nan, np.nan]])
    assert np.allclose(F_ave, F_ave_ref, equal_nan=True)

    # check that truncate nans remove correctly
    F_ave2 = truncate_nans(F_ave)
    F_ave2_ref = np.array([
        [1., 0.84, 0.72, 0.61, 0.55],
        [1., 0.84, 0.72, 0.61, 0.55]])
    assert np.allclose(F_ave2, F_ave2_ref)

    # check normal array is not changed
    assert np.allclose(truncate_nans(F_ave2_ref), F_ave2_ref)

    # if nans not the same between cols
    F = np.array([
        [1., 0.84, 0.72, 0.61, 0.55,  0.33,    np.nan],
        [1., 0.84, 0.72, 0.61, 0.55,  np.nan,  np.nan]])
    F_ref = np.array([
        [1., 0.84, 0.72, 0.61, 0.55],
        [1., 0.84, 0.72, 0.61, 0.55]])
    assert np.allclose(truncate_nans(F), F_ref)

    # if first column already contains nan, return empty time axis
    F = np.array([
        [np.nan, 0.84, 0.72],
        [np.nan, 0.84, 0.72]])
    F_ref = np.empty((2, 0))
    F_new = truncate_nans(F)

    assert F_new.shape == F_ref.shape
    assert np.allclose(F_new, F_ref)
