from framework.controller import (
    AdaptiveController,
    DecodingParameters,
    FrameworkState,
)


def _make_state(lam=0.30, novelty=0.20, bias=0.40, quality=0.60):
    return FrameworkState(
        lambda_model=lam,
        lambda_human=0.15,
        aggregate_novelty=novelty,
        aggregate_bias=bias,
        aggregate_quality=quality,
    )


def test_reward_signal_eq_xxxiii():
    controller = AdaptiveController()
    # r_{t+1} = D(s_t) - D(s_{t+1})
    assert abs(controller.compute_reward(0.8, 0.5) - 0.3) < 1e-12
    assert controller.compute_reward(0.5, 0.8) < 0


def test_td_error_and_value_update_eq_xxxv_xxxvi():
    controller = AdaptiveController(gamma=0.95, alpha=0.1)
    s0 = _make_state(novelty=0.10)
    s1 = _make_state(novelty=0.30)

    # With an empty value table: δ = r + γ·0 - 0 = r
    reward = 0.25
    delta = controller.compute_td_error(reward, s0, s1)
    assert abs(delta - reward) < 1e-12

    # V(s0) ← V(s0) + α·δ (Eq xxxvi)
    v_after = controller.update_value(s0, delta)
    assert abs(v_after - 0.1 * reward) < 1e-12

    # Second TD error now sees the updated V(s0)
    delta2 = controller.compute_td_error(reward, s0, s1)
    assert abs(delta2 - (reward - v_after)) < 1e-12


def test_theta_update_moves_along_direction_and_clips_eq_xxxvii():
    controller = AdaptiveController(eta=0.5)
    theta = DecodingParameters(temperature=0.7, top_p=0.9, top_k=50)

    # Novelty deficit direction increases exploration
    direction = controller.CONTROL_DIRECTIONS["novelty_deficit"]
    updated = controller.apply_parameter_update(theta, td_error=0.4, direction=direction)
    assert updated.temperature > theta.temperature
    assert updated.top_p >= theta.top_p
    assert updated.top_k >= theta.top_k

    # Large positive TD error must clip at θ_max
    updated_max = controller.apply_parameter_update(theta, td_error=100.0, direction=direction)
    assert updated_max.temperature == controller.TEMP_MAX
    assert updated_max.top_p == controller.TOP_P_MAX
    assert updated_max.top_k == controller.TOP_K_MAX

    # Negative TD error reverses along c and clips at θ_min
    updated_min = controller.apply_parameter_update(theta, td_error=-100.0, direction=direction)
    assert updated_min.temperature == controller.TEMP_MIN
    assert updated_min.top_p == controller.TOP_P_MIN
    assert updated_min.top_k == controller.TOP_K_MIN


def test_control_direction_follows_dominant_gap():
    controller = AdaptiveController()
    state = _make_state(lam=0.30)

    dominant, direction = controller.select_control_direction(
        {"decay_mismatch": 0.02, "novelty_deficit": 0.30,
         "bias_excess": 0.01, "quality_deficit": 0.05},
        state,
    )
    assert dominant == "novelty_deficit"
    assert direction["temperature"] > 0

    dominant, direction = controller.select_control_direction(
        {"decay_mismatch": 0.02, "novelty_deficit": 0.01,
         "bias_excess": 0.40, "quality_deficit": 0.05},
        state,
    )
    assert dominant == "bias_excess"
    assert direction["temperature"] < 0


def test_closed_loop_episode_bootstrap_and_transition():
    controller = AdaptiveController(eta=0.2, gamma=0.95, alpha=0.1)
    theta_0 = DecodingParameters(temperature=0.7, top_p=0.9, top_k=50)
    s0 = _make_state(novelty=0.10)

    d0 = {"total_distance": 0.6,
          "gap_vector": {"decay_mismatch": 0.0, "novelty_deficit": 0.25,
                         "bias_excess": 0.0, "quality_deficit": 0.05}}
    step1 = controller.initial_step(s0, d0, theta_0)
    # Bootstrap: δ_0 = D(s_0) with an empty value table
    assert abs(step1.td_error - 0.6) < 1e-12
    assert step1.theta_after.temperature > theta_0.temperature

    s1 = _make_state(novelty=0.30)
    d1 = {"total_distance": 0.35,
          "gap_vector": {"decay_mismatch": 0.0, "novelty_deficit": 0.05,
                         "bias_excess": 0.0, "quality_deficit": 0.05}}
    step2 = controller.transition_step(s0, s1, d0, d1, step1.theta_after)

    # r_1 = 0.6 - 0.35 (Eq xxxiii); V updated by α·δ (Eq xxxvi)
    assert abs(step2.reward - 0.25) < 1e-12
    assert step2.value_after != step2.value_before
    assert step2.distance_after == 0.35


def test_controller_accepts_cli_metric_keys():
    controller = AdaptiveController()

    recommendation = controller.recommend_parameters(
        {
            "novelty": 0.2,
            "self_bleu": 60.0,
            "fallback_quality": 0.8,
            "bias_proxy": 0.1,
            "perplexity": 40.0,
        },
        current_temp=0.7,
        current_top_p=0.9,
    )

    assert "temperature=0.85" in recommendation
    assert "top_p=0.96" in recommendation


def test_controller_clamps_recommendations():
    controller = AdaptiveController()

    recommendation = controller.recommend_parameters(
        {
            "Novelty": 0.0,
            "Self-BLEU": 100.0,
            "Fallback Quality": 0.0,
            "Bias Proxy": 1.0,
            "Perplexity": 999.0,
        },
        current_temp=1.2,
        current_top_p=0.98,
    )

    assert "temperature=1.20" in recommendation
    assert "top_p=0.98" in recommendation


if __name__ == "__main__":
    test_controller_accepts_cli_metric_keys()
    test_controller_clamps_recommendations()