from itertools import product

import numpy as np
import pytest

from scoobdoob import _conditioning
from scoobdoob.validation import bridge


def test_conditioned_probabilities_equal_enumerated_weighted_path_law():
    P = np.array([[0.6, 0.4, 0], [0.2, 0.4, 0.4], [0.3, 0.2, 0.5]])
    weights = np.array(
        [[1.0, 2.0, 3.0], [2.0, 0.0, 1.0], [0.5, 2.0, 1.0], [1.0, 3.0, 0.8]]
    )
    kernels, _ = _conditioning.build(P, weights, [1, 2], 4, starts=[0, 2])
    for start in [0, 2]:
        raw, actual = [], []
        for tail in product(range(3), repeat=4):
            path = (start,) + tail
            reference = float(tail[-1] in [1, 2])
            probability = 1.0
            for n in range(4):
                reference *= P[path[n], path[n + 1]] * weights[n, path[n + 1]]
                probability *= kernels[n][path[n], path[n + 1]]
            raw.append(reference)
            actual.append(probability)
        np.testing.assert_allclose(
            actual, np.array(raw) / sum(raw), atol=1e-14, rtol=1e-12
        )


def test_infeasible_start_rejected_without_filtering_initial_distribution():
    P = np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 1.0]])
    kernels = bridge.doob_kernels(P, [2], 2)
    assert not kernels[0][0].any()
    with pytest.raises(bridge.InfeasibleBridgeError, match="h_0=0"):
        bridge.sample_bridge_paths(kernels, [0, 1], 20)
    with pytest.raises(_conditioning.InfeasibleBridgeError):
        _conditioning.build(P, np.ones((2, 3)), [2], 2, starts=[0])


def test_exact_horizon_feasibility_not_just_reachability_by_horizon():
    P = np.array([[0.0, 1.0], [1.0, 0.0]])
    with pytest.raises(bridge.InfeasibleBridgeError):
        bridge.sample_bridge_paths(bridge.doob_kernels(P, [1], 2), [0])
    paths = bridge.sample_bridge_paths(bridge.doob_kernels(P, [1], 3), [0], 3)
    np.testing.assert_array_equal(paths, [[0, 1, 0, 1]] * 3)


def test_terminal_and_zero_weight_support_remain_exactly_zero():
    P = np.array(
        [
            [0.0, 0.5, 0.5, 0.0],
            [0.0, 0.0, 0.0, 1.0],
            [0.0, 0.0, 0.0, 1.0],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )
    weights = np.array([1.0, 0.0, 1.0, 1.0])
    kernels, h = _conditioning.build(P, [weights] * 2, [3], 2, starts=[0])
    assert np.array_equal(h[-1], [0.0, 0.0, 0.0, 1.0])
    for matrix in kernels:
        assert not matrix[:, 1].any()
        assert not matrix[P == 0].any()
    assert not kernels[-1][:, :3].any()
    paths = bridge.sample_bridge_paths(kernels, [0], 10)
    np.testing.assert_array_equal(paths, [[0, 2, 3]] * 10)


def test_zero_weights_can_make_every_start_infeasible():
    P = np.ones((2, 2)) / 2
    kernels = bridge.doob_kernels(P, [1], 4, w=np.zeros(2))
    assert all(not k.any() for k in kernels)
    with pytest.raises(bridge.InfeasibleBridgeError):
        bridge.sample_bridge_paths(kernels, [0])


def test_small_positive_products_do_not_trigger_fallback_or_underflow():
    P = np.array([[0.7, 0.3], [0.4, 0.6]])
    w = np.full(2, 1e-250)
    h = np.array([2e-250, 1e-250])
    expected = P * np.array([2.0, 1.0])[None, :]
    expected /= expected.sum(axis=1, keepdims=True)
    np.testing.assert_allclose(_conditioning.kernel(P, w, h), expected, rtol=1e-12)
    tiny = bridge.doob_kernels(P, [1], 1000, w=w)
    unit = bridge.doob_kernels(P, [1], 1000)
    np.testing.assert_allclose(tiny, unit, atol=1e-12, rtol=1e-12)


def test_backwards_and_kernels_agree_across_public_interfaces():
    P = np.array([[0.8, 0.2], [0.3, 0.7]])
    weights = np.array([[1.0, 2.0], [0.2, 1.0], [3.0, 0.4]])
    terminal = np.array([0.0, 1.0])
    scaled = _conditioning.backward(P, weights, 3, terminal)
    unscaled = _conditioning.backward(P, weights, 3, terminal, rescale=False)
    np.testing.assert_allclose(scaled, unscaled / unscaled.sum(axis=1, keepdims=True))
    reference = bridge.doob_kernels(P, [1], 3, weights)
    rebuilt = [_conditioning.kernel(P, weights[n], scaled[n + 1]) for n in range(3)]
    np.testing.assert_allclose(rebuilt, reference)
    kernels, backward = _conditioning.build(P, weights, [1], 3)
    np.testing.assert_allclose(kernels, reference)
    np.testing.assert_allclose(backward, scaled)


def test_early_target_departures_and_residence_are_still_allowed():
    P = np.array([[0.8, 0.2], [0.5, 0.5]])
    kernels = bridge.doob_kernels(P, [1], 4)
    assert kernels[0][0, 0] > 0
    assert kernels[1][1, 0] > 0
    paths = bridge.sample_bridge_paths(kernels, [0], 100, seed=7)
    assert all(len(p) == 5 and p[-1] == 1 for p in paths)


@pytest.mark.parametrize(
    "targets,horizon,weights",
    [
        ([], 2, None),
        ([2], 2, None),
        ([1], 0, None),
        ([1], 2, [1.0, -1.0]),
        ([1], 2, [1.0, np.nan]),
    ],
)
def test_invalid_inputs_are_not_silently_repaired(targets, horizon, weights):
    with pytest.raises(ValueError):
        bridge.doob_kernels(np.ones((2, 2)) / 2, targets, horizon, weights)


def test_invalid_start_indices_are_not_replaced_with_arbitrary_states():
    with pytest.raises(ValueError, match="out-of-range"):
        _conditioning.build(np.eye(2), np.ones((2, 2)), [1], 2, starts=[8])