import os
import sys
import unittest

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.join(ROOT, "src"))

from butterfly.dp import evaluate_policy, solve
from butterfly.environment import build_world
from butterfly.estimation import (
    build_support_gate,
    exact_trajectory_kl,
    model_errors_on_gate,
    sample_tabular,
    theorem_e_constant,
    transition_metrics,
)


class EstimationTests(unittest.TestCase):
    def test_tabular_pipeline_is_finite_and_bound_is_valid(self):
        world = build_world(0, "high", "high")
        estimate = sample_tabular(world, 500, seed=7)
        self.assertGreaterEqual(exact_trajectory_kl(world, estimate), 0.0)
        weighted, maximum, coverage = transition_metrics(world, estimate, 5)
        self.assertGreaterEqual(weighted, 0.0)
        self.assertGreaterEqual(maximum, 0.0)
        self.assertGreaterEqual(coverage, 0.0)
        self.assertLessEqual(coverage, 1.0)

        gate = build_support_gate(world, estimate, 5)
        learned = solve(
            world,
            "control",
            event_kernels=estimate.mu_hat,
            action_rewards=estimate.action_rewards,
            event_rewards=estimate.event_rewards,
            gate=gate,
        )
        true_opt = solve(world, "control")
        true_gated = solve(world, "control", gate=gate)
        learned_true = evaluate_policy(world, learned.policy)[world.start]
        eps_r, eps_mu, r_bound = model_errors_on_gate(world, estimate, gate)
        ck = theorem_e_constant(world, eps_r, eps_mu, r_bound)
        within_gate = true_gated.V[world.start] - learned_true
        self.assertLessEqual(within_gate, 2.0 * ck + 1e-9)
        self.assertGreaterEqual(true_opt.V[world.start] - learned_true, within_gate - 1e-9)


if __name__ == "__main__":
    unittest.main()

