import numpy as np
import pytest
from dooable.graph import Graph, Node, Edge, toy_graph, grid_graph
from dooable.exact import (
    solve,
    endpoint_distribution,
    sample,
    expected_cost,
    prefix_values,
    backward_policy,
    forward_from_backward,
)


def enumerate_paths(graph):
    out = []

    def visit(u, path, cost):
        if not graph.outgoing[u]:
            out.append((graph.nodes[u].outcome, path, cost))
            return
        for i in graph.outgoing[u]:
            e = graph.edges[i]
            visit(graph.node_index[e.target], path + [i], cost + e.cost)

    visit(graph.root_index, [], 0.0)
    return out


@pytest.mark.parametrize("multiplicity", [1, 2, 8, 16])
@pytest.mark.parametrize("temperature", [0.1, 1.0, 3.0])
def test_joint_law_against_route_enumeration(multiplicity, temperature):
    g = toy_graph(multiplicity)
    r = {"A": np.log(3), "B": 0.0}
    s = solve(g, r, temperature)
    paths = enumerate_paths(g)
    for y, path, cost in paths:
        partition = sum(np.exp(-c / temperature) for yy, _, c in paths if yy == y)
        expected = s.target[y] * np.exp(-cost / temperature) / partition
        assert np.prod(s.forward[path]) == pytest.approx(expected, abs=1e-11)
    assert endpoint_distribution(g, s.forward) == pytest.approx({"A": 0.75, "B": 0.25})


def test_backward_error_changes_routes_and_preserves_endpoints():
    g = toy_graph(8)
    s = solve(g, {"A": 0.0, "B": 0.0}, 0.5)
    v = prefix_values(g, 0.5)
    v = v + np.random.default_rng(4).normal(size=len(v)) * 3
    v[g.root_index] = 0
    q = backward_policy(g, v, 0.5)
    p = forward_from_backward(g, q, s.target)
    assert endpoint_distribution(g, p) == pytest.approx(s.target)
    assert expected_cost(g, p) != pytest.approx(expected_cost(g, s.forward))


def test_budgeted_cycles_are_acyclic_after_augmentation():
    g = grid_graph(4, 5)
    s = solve(g, {y: 0.0 for y in g.terminals})
    for route in sample(g, s.forward, 200, 3):
        assert sum(a.get("kind") == "move" for a in route["actions"]) <= 5


def test_cycle_rejected():
    with pytest.raises(ValueError):
        Graph(
            [Node("root"), Node("a"), Node("b"), Node("t", "x")],
            [
                Edge("1", "root", "a"),
                Edge("2", "a", "b"),
                Edge("3", "b", "a"),
                Edge("4", "b", "t"),
            ],
        )


def test_route_temperature_reduces_expected_cost():
    g = toy_graph(16)
    r = {"A": 0.0, "B": 0.0}
    assert expected_cost(g, solve(g, r, 0.05).forward) < expected_cost(
        g, solve(g, r, 2.0).forward
    )


def test_finite_residual_joint_bound():
    g = toy_graph(5)
    s = solve(g, {"A": 0.8, "B": -0.3}, 0.7)
    v = s.log_prefix + 0.05 * np.sin(np.arange(len(g.nodes)))
    v[g.root_index] = 0
    q = backward_policy(g, v, 0.7)
    p = forward_from_backward(g, q, s.target)
    from scipy.special import logsumexp

    eps = max(
        abs(
            v[j]
            - logsumexp(
                [
                    v[g.node_index[g.edges[i].source]] - g.edges[i].cost / 0.7
                    for i in ids
                ]
            )
        )
        for j, ids in enumerate(g.incoming)
        if ids
    )
    kl = 0.0
    for _, path, _ in enumerate_paths(g):
        pp = np.prod(p[path])
        qq = np.prod(s.forward[path])
        kl += pp * np.log(pp / qq)
    assert kl <= 2 * 2 * eps + 1e-12