"""
Tests for Information-Value Resource Allocator.
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
import pytest
from src.resource_allocator import InformationValueAllocator, AllocationBudget


def test_resource_allocator_decisions():
    allocator = InformationValueAllocator(info_dim=8, hidden_dim=16)

    B, T = 2, 8
    z_t = torch.rand(B, T, 8)

    decisions, diagnostics = allocator(z_t, preset="balanced")

    assert decisions["rank"] in [1, 2, 4, 8]
    assert decisions["attn_mode_idx"].shape == (B, T)
    assert decisions["depth_mode"] in ["skip", "partial", "full"]
    assert decisions["kv_precision"] in ["fp16", "int8", "int4"]
    assert "chosen_rank" in diagnostics
    assert "routing_churn_rate" in diagnostics
    print("✓ test_resource_allocator_decisions passed")


def test_resource_allocator_hysteresis():
    allocator = InformationValueAllocator(info_dim=8, hidden_dim=16, hysteresis_tau=0.5)

    B, T = 1, 4
    z_t_1 = torch.full((B, T, 8), 0.2)
    decisions_1, _ = allocator(z_t_1)
    rank_1 = decisions_1["rank"]

    # Small perturbation that should NOT break hysteresis threshold
    z_t_2 = torch.full((B, T, 8), 0.22)
    decisions_2, _ = allocator(z_t_2)
    rank_2 = decisions_2["rank"]

    assert rank_1 == rank_2, f"Hysteresis should preserve rank on small delta: {rank_1} vs {rank_2}"
    print("✓ test_resource_allocator_hysteresis passed")


def test_resource_allocator_presets():
    allocator = InformationValueAllocator(info_dim=8)
    z_t = torch.rand(2, 4, 8)

    # Edge preset should force classical
    decisions_edge, diag_edge = allocator(z_t, preset="edge")
    assert not decisions_edge["is_quantum_token"].any(), "Edge preset should disable quantum tokens"

    # Classical-only preset should force classical
    decisions_class, _ = allocator(z_t, preset="classical_only")
    assert not decisions_class["is_quantum_token"].any(), "Classical-only preset should disable quantum tokens"
    print("✓ test_resource_allocator_presets passed")


if __name__ == "__main__":
    test_resource_allocator_decisions()
    test_resource_allocator_hysteresis()
    test_resource_allocator_presets()
    print("All Resource Allocator tests passed!")