"""Tests for M5: FEP controller — active inference and generative model."""

from __future__ import annotations

import numpy as np

from radix_edge.meta.generative_model import (
    STATES, OBS_DIM, OBS_KEYS, GENERATIVE_MODEL,
    obs_dict_to_vector, OBS_BOUNDS,
)
from radix_edge.meta.fep_controller import (
    compute_log_likelihood,
    update_beliefs,
    compute_variational_free_energy,
    predict_obs_after_action,
    evaluate_action,
    choose_action,
    apply_action,
)


# --- Generative Model ---

class TestGenerativeModel:
    def test_states_defined(self):
        assert len(STATES) == 4
        assert "BALANCED" in STATES
        assert "OVERLOADED" in STATES

    def test_generative_model_shapes(self):
        for state in STATES:
            assert state in GENERATIVE_MODEL
            assert len(GENERATIVE_MODEL[state]["mean"]) == OBS_DIM
            assert len(GENERATIVE_MODEL[state]["std"]) == OBS_DIM

    def test_obs_dict_to_vector(self):
        obs = {k: 0.5 for k in OBS_KEYS}
        vec = obs_dict_to_vector(obs)
        assert vec.shape == (OBS_DIM,)
        assert all(v == 0.5 for v in vec)

    def test_obs_clamping(self):
        obs = {"gpu_util_mean": 200.0, "error_rate": -1.0}
        vec = obs_dict_to_vector(obs)
        assert vec[0] == 100.0  # clamped to max
        assert vec[5] == 0.0    # clamped to min

    def test_obs_missing_keys(self):
        vec = obs_dict_to_vector({})
        assert vec.shape == (OBS_DIM,)
        assert all(v == 0.0 for v in vec)

    def test_obs_nan_handling(self):
        obs = {"gpu_util_mean": float("nan"), "error_rate": float("inf")}
        vec = obs_dict_to_vector(obs)
        assert not np.isnan(vec).any()
        assert not np.isinf(vec).any()


# --- FEP Core Math ---

class TestFepMath:
    def _balanced_obs(self):
        """Observations typical of BALANCED state."""
        return {
            "gpu_util_mean": 60.0, "gpu_util_std": 10.0, "queue_depth": 3.0,
            "queue_depth_max": 5.0, "avg_wait_seconds": 30.0, "error_rate": 0.03,
            "throughput_jobs_per_min": 10.0, "memory_pressure": 0.5,
            "thermal_headroom": 0.5, "power_headroom": 0.5, "topology_utilization": 0.3,
        }

    def _overloaded_obs(self):
        """Observations typical of OVERLOADED state."""
        return {
            "gpu_util_mean": 92.0, "gpu_util_std": 3.0, "queue_depth": 12.0,
            "queue_depth_max": 18.0, "avg_wait_seconds": 150.0, "error_rate": 0.06,
            "throughput_jobs_per_min": 22.0, "memory_pressure": 0.88,
            "thermal_headroom": 0.82, "power_headroom": 0.82, "topology_utilization": 0.75,
        }

    def test_log_likelihood_balanced(self):
        obs = self._balanced_obs()
        vec = obs_dict_to_vector(obs)

        ll_balanced = compute_log_likelihood(vec, "BALANCED")
        ll_overloaded = compute_log_likelihood(vec, "OVERLOADED")

        # Balanced observations should have higher likelihood under BALANCED state
        assert ll_balanced > ll_overloaded

    def test_log_likelihood_overloaded(self):
        obs = self._overloaded_obs()
        vec = obs_dict_to_vector(obs)

        ll_overloaded = compute_log_likelihood(vec, "OVERLOADED")
        ll_underutilized = compute_log_likelihood(vec, "UNDERUTILIZED")

        assert ll_overloaded > ll_underutilized

    def test_update_beliefs_balanced(self):
        prior = {s: 0.25 for s in STATES}
        obs = self._balanced_obs()

        posterior = update_beliefs(prior, obs)

        # Posterior should sum to 1
        assert abs(sum(posterior.values()) - 1.0) < 1e-6
        # BALANCED should have highest posterior
        assert posterior["BALANCED"] == max(posterior.values())

    def test_update_beliefs_overloaded(self):
        prior = {s: 0.25 for s in STATES}
        obs = self._overloaded_obs()

        posterior = update_beliefs(prior, obs)
        assert abs(sum(posterior.values()) - 1.0) < 1e-6
        assert posterior["OVERLOADED"] == max(posterior.values())

    def test_free_energy_finite(self):
        prior = {s: 0.25 for s in STATES}
        obs = self._balanced_obs()
        posterior = update_beliefs(prior, obs)

        fe = compute_variational_free_energy(obs, posterior, prior)
        assert np.isfinite(fe)

    def test_free_energy_lower_at_match(self):
        """Free energy should be lower when observations match beliefs."""
        prior = {s: 0.25 for s in STATES}
        balanced_obs = self._balanced_obs()
        overloaded_obs = self._overloaded_obs()

        # Compute FE when beliefs match observations
        balanced_posterior = update_beliefs(prior, balanced_obs)
        fe_match = compute_variational_free_energy(balanced_obs, balanced_posterior, prior)

        # Compute FE when beliefs don't match
        fe_mismatch = compute_variational_free_energy(overloaded_obs, balanced_posterior, prior)

        # FE should be lower when beliefs match
        assert fe_match < fe_mismatch


# --- Action Selection ---

class TestActionSelection:
    def test_predict_balanced(self):
        obs = {"gpu_util_mean": 80.0, "queue_depth": 10.0, "avg_wait_seconds": 60.0,
               "throughput_jobs_per_min": 15.0}
        predicted = predict_obs_after_action(obs, "balanced")
        # Should move gpu_util toward 60
        assert predicted["gpu_util_mean"] < 80.0

    def test_predict_latency_first(self):
        obs = {"gpu_util_mean": 80.0, "queue_depth": 10.0, "avg_wait_seconds": 60.0}
        predicted = predict_obs_after_action(obs, "latency-first")
        assert predicted["avg_wait_seconds"] < 60.0

    def test_predict_throughput_first(self):
        obs = {"gpu_util_mean": 50.0, "queue_depth": 5.0, "throughput_jobs_per_min": 10.0}
        predicted = predict_obs_after_action(obs, "throughput-first")
        assert predicted["throughput_jobs_per_min"] > 10.0

    def test_choose_action_returns_valid(self):
        prior = {s: 0.25 for s in STATES}
        obs = {"gpu_util_mean": 60.0, "gpu_util_std": 10.0, "queue_depth": 3.0,
               "queue_depth_max": 5.0, "avg_wait_seconds": 30.0, "error_rate": 0.03,
               "throughput_jobs_per_min": 10.0, "memory_pressure": 0.5,
               "thermal_headroom": 0.5, "power_headroom": 0.5, "topology_utilization": 0.3}
        beliefs = update_beliefs(prior, obs)

        decision = choose_action(obs, beliefs, prior)
        assert decision["chosen_action"] in ("balanced", "latency-first", "throughput-first")
        assert np.isfinite(decision["expected_free_energy"])
        assert len(decision["all_evaluations"]) == 3


# --- Apply Action ---

class TestApplyAction:
    def test_apply_balanced(self, db):
        apply_action("balanced", db)
        row = db.execute("SELECT * FROM brain_params WHERE key = 'current'").fetchone()
        assert row["w_queue_depth"] == 0.5
        assert row["w_gpu_saturation"] == 0.3
        assert row["w_node_pressure"] == 0.2

    def test_apply_latency_first(self, db):
        apply_action("latency-first", db)
        row = db.execute("SELECT * FROM brain_params WHERE key = 'current'").fetchone()
        assert row["w_queue_depth"] == 0.65

    def test_apply_throughput_first(self, db):
        apply_action("throughput-first", db)
        row = db.execute("SELECT * FROM brain_params WHERE key = 'current'").fetchone()
        assert row["w_gpu_saturation"] == 0.55
