"""
Sanity checks for the SAC agent (Sec 6 of revisions.txt).
"""
import pytest
import numpy as np
import os

# Try to import the agent; skip tests if not found
try:
    from gru_sac_predictor.src import sac_agent
    # Need TF for tensor conversion if testing agent directly
    import tensorflow as tf
except ImportError:
    sac_agent = None
    tf = None

# --- Fixtures ---
@pytest.fixture(scope="module")
def sac_agent_instance():
    """
    Provides a default SAC agent instance for testing.
    Uses standard parameters suitable for basic checks.
    """
    if sac_agent is None:
        pytest.skip("SAC Agent module not found.")
    # Use default params, state_dim=5 as per revisions
    # Use fixed seeds for reproducibility in tests if needed inside agent
    agent = sac_agent.SACTradingAgent(
        state_dim=5, action_dim=1,
        initial_lr=1e-4, # Use a common LR for test simplicity
        buffer_capacity=1000, # Smaller buffer for testing
        min_buffer_size=100,
        target_entropy=-1.0
    )
    # Build the models eagerly
    try:
        agent.actor(tf.zeros((1, 5)))
        agent.critic1([tf.zeros((1, 5)), tf.zeros((1, 1))])
        agent.critic2([tf.zeros((1, 5)), tf.zeros((1, 1))])
        # Copy weights to target networks
        agent.update_target_networks(tau=1.0)
    except Exception as e:
        pytest.fail(f"Failed to build SAC agent models: {e}")
    return agent

@pytest.fixture(scope="module")
def sample_sac_inputs():
    """
    Generate sample states and corresponding directional signals.
    Simulates states with varying edge and signal-to-noise.
    """
    np.random.seed(44)
    n_samples = 1500
    # Simulate GRU outputs and position
    mu = np.random.randn(n_samples) * 0.0015 # Slightly higher variance
    sigma = np.random.uniform(0.0005, 0.0025, n_samples)
    # Simulate edge with clearer separation for testing signals
    edge_base = np.random.choice([-0.15, -0.05, 0.0, 0.05, 0.15], n_samples, p=[0.2, 0.2, 0.2, 0.2, 0.2])
    edge = np.clip(edge_base + np.random.randn(n_samples) * 0.03, -1.0, 1.0)
    z_score = np.abs(mu) / (sigma + 1e-9)
    position = np.random.uniform(-1, 1, n_samples)
    states = np.vstack([mu, sigma, edge, z_score, position]).T.astype(np.float32)
    # Use a small positive/negative threshold for determining signal from edge
    signals = np.where(edge > 0.02, 1, np.where(edge < -0.02, -1, 0))
    return states, signals

# --- Tests ---
@pytest.mark.skipif(sac_agent is None or tf is None, reason="SAC Agent module or TensorFlow not found")
def test_sac_agent_default_min_buffer(sac_agent_instance):
    """Verify the default min_buffer_size is at least 10000."""
    agent = sac_agent_instance
    # Note: Fixture currently initializes with specific values, overriding default.
    # Re-initialize with defaults for this test.
    default_agent = sac_agent.SACTradingAgent(state_dim=5, action_dim=1)
    min_buffer = default_agent.min_buffer_size
    print(f"\nAgent default min_buffer_size: {min_buffer}")
    assert min_buffer >= 10000, f"Default min_buffer_size ({min_buffer}) is less than recommended 10000."

@pytest.mark.skipif(sac_agent is None or tf is None, reason="SAC Agent module or TensorFlow not found")
def test_sac_action_variance(sac_agent_instance, sample_sac_inputs):
    """
    Verify that the mean absolute action taken when the signal is non-zero
    is >= 0.05.
    """
    agent = sac_agent_instance
    states, signals = sample_sac_inputs

    actions = []
    for state in states:
        # Use deterministic action for this sanity check
        action = agent.get_action(state, deterministic=True)
        actions.append(action[0]) # get_action returns list/array
    actions = np.array(actions)

    # Filter for non-zero signals based on the *simulated* edge
    non_zero_signal_idx = signals != 0
    if not np.any(non_zero_signal_idx):
        pytest.fail("No non-zero signals generated in fixture for SAC variance test.")

    actions_on_signal = actions[non_zero_signal_idx]

    if len(actions_on_signal) == 0:
        # This case should ideally not happen if the above check passed
        pytest.fail("Filtered actions array is empty despite non-zero signals.")

    mean_abs_action = np.mean(np.abs(actions_on_signal))

    print(f"\nSAC Sanity Test: Mean Absolute Action (on signal != 0): {mean_abs_action:.4f}")

    # Check if the agent is outputting actions with sufficient magnitude
    assert mean_abs_action >= 0.05, \
        f"Mean absolute action ({mean_abs_action:.4f}) is below threshold (0.05). Agent might be too timid or stuck near zero."

@pytest.mark.skip(reason="Requires full backtest results which are not available in this unit test setup.")
def test_sac_reward_correlation():
    """
    Optional: Check if actions taken correlate positively with subsequent rewards.
    NOTE: This test requires results from a full backtest run (actions vs rewards)
          and cannot be reliably simulated or executed in this unit test.
    """
    pass # Cannot implement without actual backtest results 