import math

import numpy as np

from agentblame import NoiseSource, abduct, gumbel_max, posterior_gumbels
from agentblame.attribution import make_game
from agentblame.demos import stochastic_rag as sr

EULER = 0.5772156649015329


def test_gumbel_replay_reproduces_factual_sample():
    logits = {"a": 0.3, "b": 1.2, "c": -0.4, "d": 0.9}
    for seed in range(200):
        w = NoiseSource(seed)
        x = w.handle("llm").categorical(logits)
        y = w.handle("llm").categorical(logits)       # replay, same logits
        assert x == y
        # same labels in a different order: noise is keyed by label, not position
        z = w.handle("llm").categorical(dict(reversed(list(logits.items()))))
        assert z == x


def test_gumbel_max_matches_softmax():
    logits = np.array([0.5, 1.5, -1.0, 0.0])
    rng = np.random.default_rng(0)
    n = 40_000
    counts = np.bincount([gumbel_max(logits, rng.gumbel(size=4)) for _ in range(n)], minlength=4)
    p = np.exp(logits) / np.exp(logits).sum()
    assert np.max(np.abs(counts / n - p)) < 0.01


def test_hashed_noise_is_standard_gumbel():
    w = NoiseSource("x")
    g = np.concatenate([w.handle(f"s{i}").gumbels(range(50)) for i in range(200)])
    assert abs(g.mean() - EULER) < 0.03
    assert abs(g.var() - math.pi ** 2 / 6) < 0.1


def test_posterior_gumbels_reproduce_observation():
    rng = np.random.default_rng(1)
    logits = [0.2, -1.0, 2.0, 0.5, -np.inf]
    for k in range(4):
        for _ in range(200):
            g = posterior_gumbels(logits, k, rng)
            assert gumbel_max(logits, g) == k


def test_posterior_counterfactual_matches_rejection_sampling():
    """P(cf choice | observed k) via abduction == via rejection on prior noise."""
    l1 = np.array([1.0, 0.0, 0.5])
    l2 = np.array([-0.5, 1.0, 0.8])     # counterfactual logits
    rng = np.random.default_rng(2)
    n = 60_000
    G = rng.gumbel(size=(n, 3))
    fact = np.argmax(l1 + G, axis=1)
    cf = np.argmax(l2 + G, axis=1)
    for k in range(3):
        ref = np.bincount(cf[fact == k], minlength=3) / (fact == k).sum()
        post = np.bincount([gumbel_max(l2, posterior_gumbels(l1, k, rng)) for _ in range(20_000)],
                           minlength=3) / 20_000
        assert np.max(np.abs(ref - post)) < 0.02, (k, ref, post)


def test_swapping_one_step_keeps_other_steps_noise():
    pipe = sr.build_pipeline()
    case = sr.make_cases(5)[2]
    w = NoiseSource(case.seed)
    a = pipe.run(case.inputs, noise=w)
    b = pipe.run(case.inputs, {"answer": "strong_model"}, noise=w)
    da = {(d.step, d.key): d.values for d in a.draws}
    db = {(d.step, d.key): d.values for d in b.draws}
    for k in da:
        assert da[k] == db[k]          # identical Gumbels for every shared address
    assert a.outputs["rewrite"] == b.outputs["rewrite"]


def test_abducted_world_reproduces_factual_run():
    pipe = sr.build_pipeline()
    for case in sr.make_cases(40):
        fact = pipe.run(case.inputs, noise=NoiseSource(case.seed))
        for j in range(3):
            w = abduct(fact, seed=j)
            again = pipe.run(case.inputs, noise=w)
            assert again.outputs == fact.outputs


def test_from_trace_replays_exactly():
    pipe = sr.build_pipeline()
    case = sr.make_cases(3)[1]
    fact = pipe.run(case.inputs, noise=NoiseSource(case.seed))
    w = NoiseSource.from_trace(fact)
    pipe.clear_cache()
    assert pipe.run(case.inputs, noise=w).outputs == fact.outputs


def test_posterior_mode_value_of_empty_coalition_is_factual():
    pipe = sr.build_pipeline()
    for case in sr.make_cases(30):
        fact = pipe.run(case.inputs, noise=NoiseSource(case.seed))
        g = make_game(pipe, case, mode="posterior", n_worlds=3, factual=fact)
        assert g.value(0) == fact.score