import tracemalloc

import numpy as np
import pytest

import online_softmax as osm


def test_naive_overflows_where_stable_and_online_do_not():
    """exp(1000) is inf in float32, so the textbook formula returns nan (inf / inf)."""
    x = np.array([1000.0, 1001.0, 1002.0], dtype=np.float32)
    with np.errstate(all="ignore"):
        assert np.isnan(osm.naive_softmax(x)).all()
    want = np.exp(np.array([-2.0, -1.0, 0.0]))
    want /= want.sum()
    assert np.allclose(osm.stable_softmax(osm.SlowRow(x)), want, atol=1e-6)
    assert np.allclose(osm.online_softmax(osm.SlowRow(x)), want, atol=1e-6)


@pytest.mark.parametrize("block", [1, 3, 64, 1000, 5000])
def test_online_stats_match_two_pass_values_in_one_pass(block):
    """Same (m, l) as the max-then-sum recipe, and every score is read exactly once."""
    x = np.random.default_rng(0).standard_normal(1000) * 5
    row = osm.SlowRow(x, block)
    m, l = osm.online_stats(row)
    assert m == x.max()
    assert np.isclose(l, np.exp(x - x.max()).sum())
    assert row.reads == len(x)


def test_three_pass_softmax_reads_each_score_three_times():
    x = np.arange(100.0)
    row = osm.SlowRow(x, block=10)
    osm.stable_softmax(row)
    assert row.reads == 3 * len(x)


def test_rescale_when_the_max_arrives_last():
    """The biggest score is in the last block: everything summed before must shrink."""
    blocks = [np.array([1.0, 2.0]), np.array([0.5]), np.array([10.0, 3.0])]
    m, l = osm.online_stats(iter(blocks))
    row = np.concatenate(blocks)
    assert m == 10.0
    assert np.isclose(l, np.exp(row - 10.0).sum())


def test_leading_masked_block_does_not_poison_the_state():
    """A block of -inf (a masked block) first: exp(-inf - -inf) would be nan without the guard."""
    blocks = [np.full(3, -np.inf), np.array([1.0, 2.0])]
    m, l = osm.online_stats(iter(blocks))
    assert m == 2.0 and np.isclose(l, np.exp(-1.0) + 1.0)


def test_worked_example_table_values():
    """The numbers printed in the lesson's tables: row by row from the digits shown."""
    hand = osm.trace(osm.WORKED_X, osm.WORKED_V, by_hand=True)
    assert [r["m"] for r in hand] == [1, 3, 3, 5, 5]
    assert [r["alpha"] for r in hand] == [0.0, 0.13534, 1.0, 0.13534, 1.0]
    assert [r["p"] for r in hand] == [1.0, 1.0, 0.36788, 1.0, 0.36788]
    assert [r["l"] for r in hand] == [1.0, 1.1353, 1.5032, 1.2034, 1.5713]
    assert [r["acc"] for r in hand] == [10.0, 21.3534, 32.3898, 44.3836, 62.7776]
    assert [round(r["out"], 2) for r in hand] == [10.0, 18.81, 21.55, 36.88, 39.95]
    terms = [round(float(np.exp(x - 5.0)), 5) for x in osm.WORKED_X]
    assert terms == [0.01832, 0.13534, 0.04979, 1.0, 0.36788]
    probs = [round(t / 1.5713, 4) for t in terms]
    assert probs == [0.0117, 0.0861, 0.0317, 0.6364, 0.2341] and round(sum(probs), 4) == 1.0
    # the full-precision run (what the script prints) stays inside 0.002 of the by-hand table
    full = osm.trace(osm.WORKED_X, osm.WORKED_V)
    assert [round(r["acc"], 4) for r in full] == [10.0, 21.3534, 32.3897, 44.3835, 62.7774]
    assert max(abs(h["acc"] - f["acc"]) for h, f in zip(hand, full)) < 2e-3
    assert max(abs(h["out"] - f["out"]) for h, f in zip(hand, full)) < 2e-3


def test_leading_masked_block_does_not_poison_attention_state():
    """Scores [-inf, -inf, 1, 2, 0.5, 0.1] in blocks of 2: without the guard the first
    block gave alpha = exp(-inf - -inf) = nan, and l and acc ended as nan."""
    K = np.array([[-np.inf], [-np.inf], [1.0], [2.0], [0.5], [0.1]])
    q = np.array([1.0])
    V = np.random.default_rng(5).standard_normal((6, 3))
    m, l, acc = osm.attention_state(q, K, V, 2)
    assert m == 2.0 and np.isfinite(l) and np.isfinite(acc).all()
    assert np.allclose(osm.online_attention(q, K, V, 2), osm.attention_reference(q, K, V))


@pytest.mark.parametrize("block", [1, 7, 64, 300, 1000])
def test_online_attention_matches_reference_for_any_block_size(block):
    rng = np.random.default_rng(1)
    q, K, V = rng.standard_normal(16), rng.standard_normal((300, 16)), rng.standard_normal((300, 16))
    assert np.allclose(osm.online_attention(q, K, V, block), osm.attention_reference(q, K, V), atol=1e-12)


def test_large_scores_stay_finite_in_float32():
    """Scores in the hundreds: exp() of them would overflow, the running max keeps them finite."""
    rng = np.random.default_rng(2)
    q = (rng.standard_normal(8) * 30).astype(np.float32)
    K = (rng.standard_normal((200, 8)) * 30).astype(np.float32)
    V = rng.standard_normal((200, 8)).astype(np.float32)
    s = K.astype(np.float64) @ q.astype(np.float64) / np.sqrt(8)
    assert np.abs(s).max() > 89                     # beyond float32 exp's range
    out = osm.online_attention(q, K, V, 32)
    assert np.isfinite(out).all()
    assert np.allclose(out, osm.attention_reference(q, K, V), atol=1e-4)


def test_merging_two_halves_equals_one_pass():
    """Partial states from different parts of the row combine into the full state."""
    rng = np.random.default_rng(3)
    q, K, V = rng.standard_normal(8), rng.standard_normal((100, 8)), rng.standard_normal((100, 8))
    whole = osm.attention_state(q, K, V, 10)
    left = osm.attention_state(q, K[:37], V[:37], 10)
    right = osm.attention_state(q, K[37:], V[37:], 10)
    m, l, acc = osm.merge(left, right)
    assert np.isclose(m, whole[0]) and np.isclose(l, whole[1]) and np.allclose(acc, whole[2])
    empty = (-np.inf, 0.0, np.zeros(8))
    assert osm.merge(empty, left)[1] == left[1]


def test_full_score_row_is_never_allocated():
    """n = 100,000 scores is 800 KB in float64. Peak memory while streaming must stay far below that."""
    rng = np.random.default_rng(4)
    n, d = 100_000, 8
    q, K, V = rng.standard_normal(d), rng.standard_normal((n, d)), rng.standard_normal((n, d))
    tracemalloc.start()
    out = osm.online_attention(q, K, V, 512)
    _, peak = tracemalloc.get_traced_memory()
    tracemalloc.stop()
    assert peak < 100_000, f"peak {peak} bytes: a row of {n} scores was probably built"
    s = K @ q / np.sqrt(d)
    p = np.exp(s - s.max())
    assert np.allclose(out, (p / p.sum()) @ V, atol=1e-10)