import os
import sys
import unittest

import numpy as np

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

from butterfly.dp import solve
from butterfly.environment import build_world
from butterfly.identities import (
    audit_reconstruction_residual,
    factorization_residual,
    kl_regret_identity_residual,
    max_gap_recursion_residual,
    trajectory_decomposition_residual,
)


class ExactIdentityTests(unittest.TestCase):
    def test_all_exact_identities(self):
        for seed in range(3):
            for entropy in ("low", "high"):
                for dispersion in ("low", "high"):
                    world = build_world(seed, entropy, dispersion)
                    control = solve(world, "control")
                    generative = solve(world, "generative")

                    factorization, normalization = factorization_residual(world, generative)
                    gap = max_gap_recursion_residual(world, generative, control)
                    relative, counting = audit_reconstruction_residual(world)
                    path, event_mean = trajectory_decomposition_residual(world, control)
                    behavior = {
                        state: world.behavior_policy(state)
                        for state in world.decision_states(include_terminal=False)
                    }
                    kl_identity = kl_regret_identity_residual(world, control, behavior)

                    for value in (
                        factorization,
                        normalization,
                        gap,
                        relative,
                        counting,
                        path,
                        event_mean,
                        kl_identity,
                    ):
                        self.assertLess(value, 1e-10)

                    self.assertGreaterEqual(generative.V[world.start] + 1e-12, control.V[world.start])


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

