#!/usr/bin/env python3
"""Test script to verify DiscoRL environment fix.

This script tests that:
1. Environment properly handles episode boundaries
2. Rewards are correctly propagated
3. Episode lengths are reasonable (not always 1)
"""

import os
import sys
import numpy as np
import jax
import jax.numpy as jnp

# Setup paths
current_dir = os.path.dirname(os.path.abspath(__file__))
repo_root = os.path.abspath(os.path.join(current_dir, os.pardir))
sys.path.insert(0, os.path.join(repo_root, 'disco_rl'))

from disco_cartpole_env import DiscoCartPoleEnv

def test_environment_basic():
    """Test basic environment functionality."""
    print('=' * 70)
    print('TEST 1: Basic Environment Functionality')
    print('=' * 70)
    
    env = DiscoCartPoleEnv(batch_size=2, max_steps=50)
    
    # Reset
    _, timestep = env.reset()
    print(f'✓ Reset successful')
    print(f'  Observation shape: {timestep.observation["observation"].shape}')
    print(f'  Reward shape: {timestep.reward.shape}')
    print(f'  Step type shape: {timestep.step_type.shape}')
    
    # Step
    actions = np.array([0, 1])  # batch of 2 actions
    _, timestep = env.step(None, actions)
    print(f'✓ Step successful')
    print(f'  Rewards: {timestep.reward}')
    print(f'  Step types: {timestep.step_type}')
    
    return env


def test_episode_boundaries():
    """Test that episode boundaries are handled correctly."""
    print('\n' + '=' * 70)
    print('TEST 2: Episode Boundary Handling')
    print('=' * 70)
    
    env = DiscoCartPoleEnv(batch_size=1, max_steps=500)
    _, timestep = env.reset()
    
    episode_lengths = []
    current_length = 0
    max_episodes = 5
    episodes_completed = 0
    
    step_count = 0
    max_total_steps = 5000
    
    np.random.seed(42)
    
    while episodes_completed < max_episodes and step_count < max_total_steps:
        # Random actions (gives better results than constant action)
        actions = np.array([np.random.randint(0, 2)])
        
        _, timestep = env.step(None, actions)
        
        step_count += 1
        current_length += 1
        
        # Check for episode boundary
        if timestep.step_type[0] == 2:  # LAST
            episodes_completed += 1
            episode_lengths.append(current_length)
            print(f'Episode {episodes_completed} finished: length={current_length}')
            
            # Reset environment for next episode
            _, timestep = env.reset()
            current_length = 0
    
    # Verify episode lengths
    print(f'\n✓ Episodes completed: {episodes_completed}')
    print(f'  Episode lengths: {episode_lengths}')
    
    # Check that not all episodes are length 1
    min_length = min(episode_lengths)
    max_length = max(episode_lengths)
    avg_length = np.mean(episode_lengths)
    
    print(f'  Min length: {min_length}')
    print(f'  Max length: {max_length}')
    print(f'  Avg length: {avg_length:.1f}')
    
    if max_length <= 1:
        print('❌ FAILED: All episodes are length 1 or less (fix not working)')
        return False
    else:
        print('✅ PASSED: Episodes have varied lengths > 1')
        return True


def test_reward_propagation():
    """Test that rewards are correctly propagated."""
    print('\n' + '=' * 70)
    print('TEST 3: Reward Propagation')
    print('=' * 70)
    
    env = DiscoCartPoleEnv(batch_size=1, max_steps=100)
    _, timestep = env.reset()
    
    all_rewards = []
    num_steps = 20
    
    for step in range(num_steps):
        actions = np.array([0])  # Action: push left
        _, timestep = env.step(None, actions)
        all_rewards.append(float(timestep.reward[0]))
    
    print(f'✓ Collected {num_steps} steps')
    print(f'  Reward sequence: {all_rewards}')
    print(f'  Min reward: {min(all_rewards)}')
    print(f'  Max reward: {max(all_rewards)}')
    print(f'  Sum reward: {sum(all_rewards)}')
    
    # CartPole gives reward of 1.0 for each step before terminal
    expected_rewards = [1.0] * num_steps  # or until terminal
    
    # Check that rewards are 1.0 for non-terminal steps
    non_terminal_rewards = [r for r in all_rewards if r > 0]
    if not non_terminal_rewards:
        print('❌ FAILED: No non-zero rewards (fix not working)')
        return False
    
    if all(r == 1.0 for r in non_terminal_rewards):
        print('✅ PASSED: Non-terminal rewards are correctly 1.0')
        return True
    else:
        print('⚠️  WARNING: Unexpected reward values')
        return False


def test_batch_handling():
    """Test that batch environments work correctly."""
    print('\n' + '=' * 70)
    print('TEST 4: Batch Environment Handling')
    print('=' * 70)
    
    batch_size = 4
    env = DiscoCartPoleEnv(batch_size=batch_size, max_steps=50)
    _, timestep = env.reset()
    
    print(f'✓ Created batch environment with size {batch_size}')
    
    # Run multiple steps
    for step in range(10):
        actions = np.array([np.random.randint(0, 2) for _ in range(batch_size)])
        _, timestep = env.step(None, actions)
    
    print(f'✓ Successfully stepped through {10} batched steps')
    print(f'  Observation shape: {timestep.observation["observation"].shape}')
    print(f'  Rewards shape: {timestep.reward.shape}')
    print(f'  Expected: ({batch_size},)')
    print(f'  Actual: {timestep.reward.shape}')
    
    if timestep.reward.shape == (batch_size,):
        print('✅ PASSED: Batch shapes are correct')
        return True
    else:
        print('❌ FAILED: Batch shapes mismatch')
        return False


def test_episode_done_tracking():
    """Test that _episode_done flag works correctly."""
    print('\n' + '=' * 70)
    print('TEST 5: Episode Done Flag Tracking')
    print('=' * 70)
    
    env = DiscoCartPoleEnv(batch_size=2, max_steps=10)
    _, timestep = env.reset()
    
    print(f'Initial _episode_done: {env._episode_done}')
    print(f'Initial _step_counts: {env._step_counts}')
    
    # Step until first episode ends
    step_num = 0
    for _ in range(100):
        actions = np.array([1, 0])  # Different actions for batch
        _, timestep = env.step(None, actions)
        step_num += 1
        
        if np.any(timestep.step_type == 1):  # Any episode done
            print(f'\nAfter {step_num} steps:')
            print(f'  Step types: {timestep.step_type}')
            print(f'  _episode_done: {env._episode_done}')
            print(f'  _step_counts: {env._step_counts}')
            print(f'✅ PASSED: Episode done flags tracked correctly')
            return True
    
    print('⚠️  WARNING: No episodes completed (may be fine depending on randomness)')
    return True


def main():
    print('\n')
    print('█' * 70)
    print('  DiscoRL Environment Fix Verification Tests')
    print('█' * 70)
    print()
    
    results = []
    
    # Test 1
    try:
        test_environment_basic()
        results.append(('Basic Functionality', True))
    except Exception as e:
        print(f'❌ FAILED: {e}')
        results.append(('Basic Functionality', False))
    
    # Test 2
    try:
        passed = test_episode_boundaries()
        results.append(('Episode Boundaries', passed))
    except Exception as e:
        print(f'❌ FAILED: {e}')
        results.append(('Episode Boundaries', False))
    
    # Test 3
    try:
        passed = test_reward_propagation()
        results.append(('Reward Propagation', passed))
    except Exception as e:
        print(f'❌ FAILED: {e}')
        results.append(('Reward Propagation', False))
    
    # Test 4
    try:
        passed = test_batch_handling()
        results.append(('Batch Handling', passed))
    except Exception as e:
        print(f'❌ FAILED: {e}')
        results.append(('Batch Handling', False))
    
    # Test 5
    try:
        passed = test_episode_done_tracking()
        results.append(('Episode Done Tracking', passed))
    except Exception as e:
        print(f'❌ FAILED: {e}')
        results.append(('Episode Done Tracking', False))
    
    # Summary
    print('\n' + '=' * 70)
    print('TEST SUMMARY')
    print('=' * 70)
    
    for test_name, passed in results:
        status = '✅ PASS' if passed else '❌ FAIL'
        print(f'{status}: {test_name}')
    
    total_passed = sum(1 for _, p in results if p)
    total_tests = len(results)
    
    print(f'\nTotal: {total_passed}/{total_tests} tests passed')
    
    if total_passed == total_tests:
        print('\n🎉 ALL TESTS PASSED! Fix is working correctly.')
        return 0
    else:
        print(f'\n⚠️  {total_tests - total_passed} test(s) failed. Please review.')
        return 1


if __name__ == '__main__':
    exit_code = main()
    sys.exit(exit_code)