"""
Unit tests for SearchAgent.
"""

import torch
import pytest
from src.search_agent import SearchAgent, PolicyValueNet
from src.utils import encode_state, get_actions, step


def test_policy_value_net_forward():
    net = PolicyValueNet(input_dim=1, action_space=2)
    x = torch.tensor([[3.0]])
    policy, value = net(x)
    assert policy.shape == (1, 2)
    assert value.shape == (1, 1)


def test_search_agent_action_selection():
    net = PolicyValueNet(input_dim=1, action_space=2)
    agent = SearchAgent(policy_value_net=net, max_depth=2)
    # Start from state 0; actions are 1 and 2
    action = agent.act(0)
    assert action in [1, 2]


def test_search_agent_value_estimation():
    net = PolicyValueNet(input_dim=1, action_space=2)
    agent = SearchAgent(policy_value_net=net, max_depth=3)
    # For state 8, the optimal action is 10 (reward 1)
    action = agent.act(8)
    assert action == 10 or action == 9  # depending on policy, 10 is better


def test_search_agent_terminal_state():
    net = PolicyValueNet(input_dim=1, action_space=2)
    agent = SearchAgent(policy_value_net=net, max_depth=1)
    # State 10 is terminal; agent should return 10
    action = agent.act(10)
    assert action == 10