"""
Tests for the SACTradingAgent class.

Ref: revisions.txt Task 5.7
"""
import pytest
import numpy as np
import tensorflow as tf
import sys, os

# --- Add path for src imports --- #
script_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(script_dir)
src_path = os.path.join(project_root, 'src')
if src_path not in sys.path:
    sys.path.insert(0, src_path)
# --- End Add path --- #

from sac_agent import SACTradingAgent

# --- Constants --- #
STATE_DIM = 5
ACTION_DIM = 1
BUFFER_SIZE = 5000
MIN_BUFFER = 1000
TRAIN_STEPS = 1500 # Number of training steps for the test
BATCH_SIZE = 64

# --- Fixtures --- #

@pytest.fixture
def sac_agent_fixture() -> SACTradingAgent:
    """Provides a default SACTradingAgent instance for testing."""
    agent = SACTradingAgent(
        state_dim=STATE_DIM,
        action_dim=ACTION_DIM,
        buffer_capacity=BUFFER_SIZE,
        min_buffer_size=MIN_BUFFER,
        alpha_auto_tune=True, # Enable auto-tuning for realistic test
        target_entropy=-1.0 * ACTION_DIM # Default target entropy
    )
    return agent

def _populate_buffer(agent: SACTradingAgent, num_samples: int):
    """Helper to add random transitions to the agent's buffer."""
    print(f"\nPopulating buffer with {num_samples} random samples...")
    for _ in range(num_samples):
        state = np.random.randn(STATE_DIM).astype(np.float32)
        action = np.random.uniform(-1, 1, size=(ACTION_DIM,)).astype(np.float32)
        reward = np.random.randn()
        next_state = np.random.randn(STATE_DIM).astype(np.float32)
        done = float(np.random.rand() < 0.05) # 5% chance of done
        agent.buffer.add(state, action, reward, next_state, done)
    print(f"Buffer populated. Size: {len(agent.buffer)}")

# --- Tests --- #

def test_sac_training_updates(sac_agent_fixture):
    """
    Test 5.7: Run training steps and check for basic health:
    a) Q-values are not NaN.
    b) Action variance is reasonable (suggests exploration).
    """
    agent = sac_agent_fixture
    # Populate buffer sufficiently to start training
    _populate_buffer(agent, MIN_BUFFER + BATCH_SIZE)
    
    print(f"\nRunning {TRAIN_STEPS} training steps...")
    metrics_history = []
    for i in range(TRAIN_STEPS):
        metrics = agent.train(batch_size=BATCH_SIZE)
        if metrics: # Train only runs if buffer is full enough
            metrics_history.append(metrics)
        # Basic check within the loop to fail fast
        if i % 100 == 0 and metrics:
             assert not np.isnan(metrics['critic1_loss']), f"Critic1 loss is NaN at step {i}"
             assert not np.isnan(metrics['critic2_loss']), f"Critic2 loss is NaN at step {i}"
             assert not np.isnan(metrics['actor_loss']), f"Actor loss is NaN at step {i}"
             if agent.alpha_auto_tune:
                 assert not np.isnan(metrics['alpha_loss']), f"Alpha loss is NaN at step {i}"

    assert len(metrics_history) > 0, "Training loop did not execute (buffer size issue?)"
    print(f"Training steps completed. Last metrics: {metrics_history[-1]}")

    # a) Check final Q-values (indirectly via loss)
    last_metrics = metrics_history[-1]
    assert not np.isnan(last_metrics['critic1_loss']), "Final Critic1 loss is NaN"
    assert not np.isnan(last_metrics['critic2_loss']), "Final Critic2 loss is NaN"
    # We assume if losses are not NaN, Q-values involved are also not NaN
    print("Check a) Passed: Q-value losses are not NaN.")

    # b) Check action variance after training
    num_samples_for_variance = 500
    sampled_actions = []
    dummy_state = np.random.randn(STATE_DIM).astype(np.float32)
    for _ in range(num_samples_for_variance):
        # Sample non-deterministically to check stochastic policy variance
        action = agent.get_action(dummy_state, deterministic=False)
        sampled_actions.append(action)
        
    sampled_actions = np.array(sampled_actions)
    action_variance = np.var(sampled_actions, axis=0)
    print(f"Action variance after {TRAIN_STEPS} steps: {action_variance}")
    
    # Check if variance is above a threshold (e.g., 0.2 from revisions.txt)
    # This threshold might need tuning based on action space scaling (-1 to 1)
    min_variance_threshold = 0.2
    assert np.all(action_variance > min_variance_threshold), \
        f"Action variance ({action_variance}) is below threshold ({min_variance_threshold}). Exploration might be too low."
    print(f"Check b) Passed: Action variance ({action_variance.round(3)}) > {min_variance_threshold}.") 