import math
from fractions import Fraction

import numpy as np
import pytest

from agentblame import (analyze_case, exact_shapley, last_touch, leave_one_out, mc_shapley,
                        single_swap)
from agentblame.attribution import BlameGame
from agentblame.demos import refund_policy as rp
from agentblame.noise import NoiseSource

from conftest import all_subsets, game_from_function


def random_game(n, seed):
    rng = np.random.default_rng(seed)
    names = [f"s{i}" for i in range(n)]
    table = {S: float(rng.uniform()) for S in all_subsets(names)}
    return names, table


@pytest.mark.parametrize("seed", range(5))
def test_efficiency(seed):
    names, table = random_game(5, seed)
    g = game_from_function(names, table.__getitem__)
    phi = exact_shapley(g).values
    assert math.isclose(sum(phi.values()), table[frozenset(names)] - table[frozenset()], abs_tol=1e-12)


def test_symmetry_and_null_player():
    # a and b are interchangeable; d never changes anything.
    names = ["a", "b", "c", "d"]
    v = lambda S: float(("a" in S or "b" in S) and "c" in S) + 0.5 * float({"a", "b"} <= S)
    g = game_from_function(names, v)
    phi = exact_shapley(g).values
    assert math.isclose(phi["a"], phi["b"], abs_tol=1e-12)
    assert abs(phi["d"]) < 1e-12
    assert math.isclose(sum(phi.values()), v(frozenset(names)) - v(frozenset()), abs_tol=1e-12)


def test_matches_brute_force_permutation_definition():
    import itertools
    names, table = random_game(4, 11)
    g = game_from_function(names, table.__getitem__)
    phi = exact_shapley(g).values
    acc = {n: 0.0 for n in names}
    perms = list(itertools.permutations(names))
    for p in perms:
        S = frozenset()
        for n in p:
            acc[n] += table[S | {n}] - table[S]
            S = S | {n}
    for n in names:
        assert math.isclose(phi[n], acc[n] / len(perms), abs_tol=1e-12)


def test_refund_value_function_is_as_specified():
    pipe = rp.build_pipeline()
    g = BlameGame(pipe, rp.CASE.inputs, rp.FIXES_A, world=NoiseSource(0))
    for S in all_subsets(g.players):
        expected = float("planner" in S or "retriever" in S or {"writer", "reviewer"} <= S)
        assert g.value_of(S) == expected, S


def test_refund_hand_computed_shapley():
    # v = 1 iff P or R or (W and Rv). For P: marginal = 1 - v(S) over S in {R,W,Rv};
    # v(S)=0 only for {}, {W}, {Rv} -> weights 1/4 + 1/12 + 1/12 = 5/12.
    # For W: marginal = 1 only at S={Rv} -> weight 1/12. Total 5/12*2 + 1/12*2 = 1.
    cb = analyze_case(rp.build_pipeline(), rp.CASE, fixes=rp.FIXES_A, method="exact")
    want = {"planner": Fraction(5, 12), "retriever": Fraction(5, 12),
            "writer": Fraction(1, 12), "reviewer": Fraction(1, 12)}
    for k, w in want.items():
        assert math.isclose(cb.shapley.values[k], float(w), abs_tol=1e-12), k
    # baselines tell different stories
    assert cb.fix_alone == {"planner": 1.0, "retriever": 1.0, "writer": 0.0, "reviewer": 0.0}
    assert all(v == 0.0 for v in cb.leave_one_out.values())   # over-determination
    assert cb.last_touch == "reviewer"


def test_refund_alternative_reviewer_changes_blame():
    cb = analyze_case(rp.build_pipeline(), rp.CASE, fixes=rp.FIXES_B, method="exact")
    for k, w in {"planner": 1 / 3, "retriever": 1 / 3, "writer": 0.0, "reviewer": 1 / 3}.items():
        assert math.isclose(cb.shapley.values[k], w, abs_tol=1e-12), k


def test_caching_saves_cost():
    pipe = rp.build_pipeline()
    g = BlameGame(pipe, rp.CASE.inputs, rp.FIXES_A, world=NoiseSource(0))
    res = exact_shapley(g)
    assert res.coalitions_evaluated == 16
    assert res.cost < 16 * 4      # prefix sharing: upstream steps are not re-run
    assert pipe.meter.cache_hits > 0


def test_mc_converges_to_exact_and_cis_cover():
    names, table = random_game(6, 3)
    g = game_from_function(names, table.__getitem__)
    exact = exact_shapley(g).values
    res = mc_shapley(g, 4000, seed=0)
    err = max(abs(res.values[n] - exact[n]) for n in names)
    assert err < 0.03
    for n in names:
        lo, hi = res.ci[n]
        assert lo <= exact[n] <= hi
    # each permutation's marginals telescope, so the estimate is exactly efficient
    assert math.isclose(sum(res.values.values()), table[frozenset(names)] - table[frozenset()],
                        abs_tol=1e-9)


def test_mc_error_shrinks():
    names, table = random_game(7, 5)
    g = game_from_function(names, table.__getitem__)
    exact = np.array(list(exact_shapley(g).values.values()))
    errs = []
    for k in (20, 2000):
        r = mc_shapley(game_from_function(names, table.__getitem__), k, seed=1)
        errs.append(np.abs(np.array(list(r.values.values())) - exact).max())
    assert errs[1] < errs[0]


def test_adaptive_stopping_separates_top_player():
    names = [f"s{i}" for i in range(8)]
    # s0 is a veto player (required); others are interchangeable helpers.
    v = lambda S: float("s0" in S and len(S) >= 3)
    g = game_from_function(names, v)
    exact = exact_shapley(g).values
    g2 = game_from_function(names, v)
    res = mc_shapley(g2, 5000, stop="top1", seed=0)
    assert res.stop_reason.startswith("separated")
    assert res.n_permutations < 5000
    assert max(res.values, key=res.values.get) == "s0"
    lo, hi = res.ci["s0"]
    assert lo <= exact["s0"] <= hi


def test_budget_stop():
    names, table = random_game(8, 2)
    g = game_from_function(names, table.__getitem__, cache=False)   # every run costs 8
    res = mc_shapley(g, 10_000, stop="top1", budget=200, seed=0)
    assert res.stop_reason == "budget"
    assert 200 <= res.cost < 200 + 9 * 8          # at most one permutation of overshoot


def test_baselines_on_simple_and_game():
    g = game_from_function(["a", "b"], lambda S: float({"a", "b"} <= S))
    assert single_swap(g) == {"a": 0.0, "b": 0.0}
    assert leave_one_out(g) == {"a": 1.0, "b": 1.0}
    assert exact_shapley(g).values == {"a": 0.5, "b": 0.5}
    assert last_touch(g) == "b"