import math

import numpy as np
import pytest

from timeseries_entropy import plugin_entropy, unbiased_entropy, level_corrections


def test_plugin_entropy_exact():
    assert plugin_entropy([0, 1]) == pytest.approx(1.0)
    assert plugin_entropy([3, 3, 3]) == pytest.approx(0.0)
    assert plugin_entropy([0, 0, 0, 1]) == pytest.approx(
        -(0.75 * math.log2(0.75) + 0.25 * math.log2(0.25)))


def test_unbiased_on_iid_categorical():
    p = np.array([0.5, 0.3, 0.2])
    true_h = -np.sum(p * np.log2(p))
    rng = np.random.default_rng(1)
    draw = lambda k: rng.choice(3, size=k, p=p)
    vals = np.array([unbiased_entropy(draw, n0=32, r=1.5, rng=rng)
                     for _ in range(4000)])
    se = vals.std(ddof=1) / math.sqrt(len(vals))
    assert abs(vals.mean() - true_h) < 5 * se
    assert se < 0.01


def test_unbiased_on_autocorrelated_chain():
    # Stationary two-state Markov chain with sticky transitions; the marginal
    # is uniform, so H = 1 bit despite strong autocorrelation.
    rng = np.random.default_rng(2)
    state = [int(rng.random() < 0.5)]

    def draw(k):
        out = np.empty(k, dtype=int)
        for i in range(k):
            if rng.random() < 0.1:
                state[0] = 1 - state[0]
            out[i] = state[0]
        return out

    vals = np.array([unbiased_entropy(draw, n0=64, r=1.5, rng=rng)
                     for _ in range(3000)])
    se = vals.std(ddof=1) / math.sqrt(len(vals))
    assert abs(vals.mean() - 1.0) < 5 * se
    assert se < 0.02


def test_level_corrections_decay():
    rng = np.random.default_rng(3)
    p = np.array([0.5, 0.3, 0.2])
    deltas = np.array([
        level_corrections(lambda k: rng.choice(3, size=k, p=p), n0=64, levels=5)
        for _ in range(200)])
    rms = np.sqrt((deltas ** 2).mean(axis=0))
    # E[Delta_m^2] should decay roughly like 2^(-2m); demand clear decay.
    assert rms[-1] < rms[0] / 4