import math


def _imports():
    try:
        from bayesian_langchain.smc import ParticleFilter, Proposal, SMCConfig
    except ImportError:
        from bayesian_langgraph.smc import ParticleFilter, Proposal, SMCConfig
    return ParticleFilter, Proposal, SMCConfig


def test_final_step_preserves_nonuniform_posterior_weights():
    ParticleFilter, Proposal, SMCConfig = _imports()
    pf = ParticleFilter(SMCConfig(n_particles=64, n_steps=1, ess_resample_fraction=0.99, seed=7))

    def propose(state, t):
        return [Proposal("good", 0.5), Proposal("bad", 0.5)]

    def transition(state, action, t):
        return action

    def log_target(prev, action, nxt, t):
        return 0.0 if action == "good" else math.log(1e-5)

    particles = pf.run(None, propose, transition, log_target)
    probs = [p.posterior for p in particles]
    assert abs(sum(probs) - 1.0) < 1e-12
    # If the old final-resampling bug returned, these would all be 1/N.
    assert max(probs) - min(probs) > 1e-3


def test_importance_weight_contains_inverse_proposal_probability():
    ParticleFilter, Proposal, SMCConfig = _imports()
    pf = ParticleFilter(SMCConfig(n_particles=256, n_steps=1, seed=11))

    def propose(state, t):
        return [Proposal("a", 0.9), Proposal("b", 0.1)]

    def transition(state, action, t):
        return action

    # Equal target mass for both actions. Importance correction should cancel
    # the 9:1 proposal imbalance in the weighted posterior, up to MC error.
    def log_target(prev, action, nxt, t):
        return math.log(0.5)

    particles = pf.run(None, propose, transition, log_target)
    mass = {"a": 0.0, "b": 0.0}
    for p in particles:
        mass[p.state] += p.posterior
    assert abs(mass["a"] - 0.5) < 0.15
    assert abs(mass["b"] - 0.5) < 0.15