import pytest
import numpy as np
from omegaconf import OmegaConf

# Adjust import path based on structure
from gru_sac_predictor.src.trading_env import TradingEnv

# --- Test Fixtures ---

@pytest.fixture
def sample_env_data():
    """Provides sample data for initializing the TradingEnv."""
    n_steps = 10
    data = {
        'mu_predictions': np.random.randn(n_steps) * 0.001,
        'sigma_predictions': np.abs(np.random.randn(n_steps) * 0.002 + 0.005),
        'p_cal_predictions': np.random.rand(n_steps),
        'actual_returns': np.random.randn(n_steps) * 0.0015,
        'bar_imputed_flags': np.array([False, False, True, False, True, True, False, False, True, False], dtype=bool)
    }
    return data

@pytest.fixture
def base_env_config():
    """Base configuration for the environment."""
    return OmegaConf.create({
        'sac': {
            'imputed_handling': 'skip', # Default test mode
            'action_penalty': 0.05
        },
        'environment': {
            'initial_capital': 10000.0,
            'transaction_cost': 0.0005,
            'reward_scale': 100.0,
            'action_penalty_lambda': 0.0 # Usually overridden by transaction_cost calc
        }
    })

@pytest.fixture
def trading_env_instance(sample_env_data, base_env_config):
    """Creates a TradingEnv instance with default 'skip' mode."""
    return TradingEnv(**sample_env_data, config=base_env_config)

# --- Test Functions ---

def test_env_initialization(trading_env_instance, sample_env_data):
    assert trading_env_instance.n_steps == len(sample_env_data['actual_returns'])
    assert trading_env_instance.current_step == 0
    assert trading_env_instance.current_position == 0.0
    assert np.array_equal(trading_env_instance.bar_imputed, sample_env_data['bar_imputed_flags'])

def test_env_reset(trading_env_instance):
    # Take a few steps
    trading_env_instance.step(0.5)
    trading_env_instance.step(-0.2)
    assert trading_env_instance.current_step > 0
    # Reset
    initial_state = trading_env_instance.reset()
    assert trading_env_instance.current_step == 0
    assert trading_env_instance.current_position == 0.0
    assert initial_state is not None
    assert initial_state.shape == (trading_env_instance.state_dim,)

def test_env_step_normal(trading_env_instance):
    # Test a normal step (step 0 is not imputed)
    initial_pos = trading_env_instance.current_position
    action = 0.7
    next_state, reward, done, info = trading_env_instance.step(action)

    assert trading_env_instance.current_step == 1
    assert trading_env_instance.current_position == action # Position updates to action
    assert not info['is_imputed_step_skipped']
    assert not done
    assert next_state is not None
    # Reward calculation is complex, just check type/sign if needed
    assert isinstance(reward, float)

def test_env_step_imputed_skip(trading_env_instance, sample_env_data):
    # Step 2 is imputed in sample_env_data
    trading_env_instance.step(0.5) # Step 0
    trading_env_instance.step(0.6) # Step 1
    assert trading_env_instance.current_step == 2
    initial_pos_before_imputed = trading_env_instance.current_position

    # Action for the imputed step (should be ignored by 'skip')
    action_imputed = 0.9
    next_state, reward, done, info = trading_env_instance.step(action_imputed)

    # Should skip step 2 and now be at step 3
    assert trading_env_instance.current_step == 3 
    # Position should NOT have changed from step 1
    assert trading_env_instance.current_position == initial_pos_before_imputed
    assert reward == 0.0 # Skip gives 0 reward
    assert not done
    assert info['is_imputed_step_skipped'] == True # Crucial check for buffer
    # Check that the returned state is for step 3
    expected_state_step_3 = trading_env_instance._get_state() # Get state now that we are at step 3
    np.testing.assert_array_almost_equal(next_state, expected_state_step_3)

def test_env_step_imputed_hold(sample_env_data, base_env_config):
    cfg = base_env_config.copy()
    cfg.sac.imputed_handling = 'hold'
    env = TradingEnv(**sample_env_data, config=cfg)

    # Step 2 is imputed
    env.step(0.5) # Step 0
    env.step(0.6) # Step 1
    assert env.current_step == 2
    position_before_imputed = env.current_position

    # Action for the imputed step (should be overridden by 'hold')
    action_imputed = -0.5 
    next_state, reward, done, info = env.step(action_imputed)

    # Should process step 2 and move to step 3
    assert env.current_step == 3
    # Position should be the same as before the step
    assert env.current_position == position_before_imputed
    assert not info['is_imputed_step_skipped']
    assert not done
    # Reward should be calculated based on holding the position
    expected_pnl = position_before_imputed * (np.exp(sample_env_data['actual_returns'][2]) - 1)
    expected_cost = 0 # No trade size if holding
    expected_penalty = 0 # No penalty in hold mode
    expected_raw_reward = expected_pnl - expected_cost - expected_penalty
    expected_scaled_reward = expected_raw_reward * cfg.environment.reward_scale
    assert np.isclose(reward, expected_scaled_reward)

def test_env_step_imputed_penalty(sample_env_data, base_env_config):
    cfg = base_env_config.copy()
    cfg.sac.imputed_handling = 'penalty'
    cfg.sac.action_penalty = 0.1 # Use a specific penalty for testing
    env = TradingEnv(**sample_env_data, config=cfg)

    # Step 2 is imputed
    env.step(0.5) # Step 0
    env.step(0.6) # Step 1
    assert env.current_step == 2
    position_before_imputed = env.current_position # Should be 0.6

    # Action for the imputed step
    action_imputed = -0.2 
    next_state, reward, done, info = env.step(action_imputed)

    # Should process step 2 and move to step 3
    assert env.current_step == 3
    # Position should update to the *agent's* action
    assert env.current_position == np.clip(action_imputed, -1.0, 1.0)
    assert not info['is_imputed_step_skipped']
    assert not done
    
    # Reward calculation is ONLY the penalty
    expected_raw_reward = -cfg.sac.action_penalty * (action_imputed - position_before_imputed)**2
    expected_scaled_reward = expected_raw_reward * cfg.environment.reward_scale
    assert np.isclose(reward, expected_scaled_reward)

def test_env_done_condition(trading_env_instance, sample_env_data):
    n_steps = len(sample_env_data['actual_returns'])
    # Step through the environment
    done = False
    for i in range(n_steps):
        _, _, done, _ = trading_env_instance.step(np.random.uniform(-1, 1))
        if i < n_steps - 1:
            assert not done
        else:
            assert done # Should be done on the last step 