"""
Tests for Token Information State Module.
"""

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

import torch
import pytest
from src.information_state import TokenInformationState


def test_information_state_shape_and_range():
    d_model = 64
    n_heads = 4
    engine = TokenInformationState(d_model=d_model, n_heads=n_heads)

    batch_size = 2
    seq_len = 16
    x = torch.randn(batch_size, seq_len, d_model)
    attn = torch.softmax(torch.randn(batch_size, n_heads, seq_len, seq_len), dim=-1)

    z_t = engine(x, attn_weights=attn)

    assert z_t.shape == (batch_size, seq_len, 8), f"Expected (2, 16, 8), got {z_t.shape}"
    # Values should be normalized to [0, 1]
    assert z_t.min() >= 0.0, f"Min value {z_t.min()} < 0.0"
    assert z_t.max() <= 1.0, f"Max value {z_t.max()} > 1.0"
    print("✓ test_information_state_shape_and_range passed")


def test_information_state_ablation():
    engine = TokenInformationState(d_model=32, n_heads=2)
    x = torch.randn(1, 8, 32)

    # Disable predictive_entropy
    engine.set_ablation("predictive_entropy", False)
    z_t = engine(x)

    pred_idx = engine.FEATURE_NAMES.index("predictive_entropy")
    assert (z_t[:, :, pred_idx] == 0.0).all(), "Ablated feature should be zero"

    # Reset
    engine.reset_ablation()
    z_t_reset = engine(x)
    assert (z_t_reset[:, :, pred_idx] > 0.0).any(), "Reset feature should be restored"
    print("✓ test_information_state_ablation passed")


if __name__ == "__main__":
    test_information_state_shape_and_range()
    test_information_state_ablation()
    print("All Information State tests passed!")