"""Quick sanity check: verify Gym->DiscoRL adapter and agent initialization work.

This script runs a minimal test without full training, just to ensure:
1. Environment adapter creates correctly
2. Agent initializes with correct specs
3. Forward pass (reset + a few steps) works
4. Disco103 weights load (if available)
"""

import os
import sys
import numpy as np

# Force JAX to use CPU only (avoid GPU memory issues)
os.environ['JAX_PLATFORMS'] = 'cpu'

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'))

import jax
import jax.numpy as jnp
from disco_rl import agent as disco_agent
from disco_rl import types
from disco_cartpole_env import DiscoCartPoleEnv
from disco_weights import load_disco103_weights


def test_env_creation():
    """Test 1: Create DiscoCartPoleEnv."""
    print('\n' + '='*60)
    print('Test 1: Environment Creation')
    print('='*60)
    
    try:
        env = DiscoCartPoleEnv(batch_size=2)
        print('✓ DiscoCartPoleEnv created successfully')
        
        obs_spec = env.single_observation_spec()
        act_spec = env.single_action_spec()
        
        print(f'  Single obs spec: {obs_spec}')
        print(f'  Single act spec: {act_spec}')
        
        # Check specs
        assert 'observation' in obs_spec, 'obs spec must have "observation" key'
        print('✓ Observation spec has correct structure')
        
        assert act_spec.dtype == np.int32, 'action spec must be int32'
        assert act_spec.shape == (), 'action spec must be scalar'
        print('✓ Action spec is scalar discrete (int32)')
        
        return env
    except Exception as e:
        print(f'✗ Environment creation failed: {e}')
        import traceback
        traceback.print_exc()
        return None


def test_env_reset_step(env):
    """Test 2: Reset and step environment."""
    print('\n' + '='*60)
    print('Test 2: Environment Reset & Step')
    print('='*60)
    
    try:
        _, timestep = env.reset()
        print('✓ Environment reset')
        
        # Check timestep structure
        assert isinstance(timestep, types.EnvironmentTimestep), 'timestep must be EnvironmentTimestep'
        print('✓ Timestep has correct type')
        
        obs = timestep.observation['observation']
        print(f'  Observation shape: {obs.shape}')
        print(f'  Observation dtype: {obs.dtype}')
        print(f'  Observation range: [{jnp.min(obs):.3f}, {jnp.max(obs):.3f}]')
        
        # Step with random actions
        batch_size = env.batch_size
        actions = jnp.array([0, 1], dtype=jnp.int32)[:batch_size]
        _, timestep_next = env.step(None, actions)
        print(f'✓ Environment stepped with actions {actions}')
        
        print(f'  Reward shape: {timestep_next.reward.shape}')
        print(f'  Reward values: {timestep_next.reward}')
        
        return True
    except Exception as e:
        print(f'✗ Reset/step failed: {e}')
        import traceback
        traceback.print_exc()
        return False


def test_agent_creation(env):
    """Test 3: Create DiscoRL agent."""
    print('\n' + '='*60)
    print('Test 3: DiscoRL Agent Creation')
    print('='*60)
    
    try:
        settings = disco_agent.get_settings_disco()
        print('✓ Loaded Disco agent settings')
        
        agent = disco_agent.Agent(
            single_observation_spec=env.single_observation_spec(),
            single_action_spec=env.single_action_spec(),
            agent_settings=settings,
            batch_axis_name=None,
        )
        print('✓ DiscoRL Agent created')
        
        # Initialize states
        rng = jax.random.PRNGKey(0)
        rng, subkey = jax.random.split(rng)
        learner_state = agent.initial_learner_state(subkey)
        print('✓ Learner state initialized')
        
        rng, subkey = jax.random.split(rng)
        actor_state = agent.initial_actor_state(subkey)
        print('✓ Actor state initialized')
        
        return agent, learner_state, actor_state, rng
    except Exception as e:
        print(f'✗ Agent creation failed: {e}')
        import traceback
        traceback.print_exc()
        return None, None, None, None


def test_agent_forward(env, agent, learner_state, actor_state, rng):
    """Test 4: Agent forward pass."""
    print('\n' + '='*60)
    print('Test 4: Agent Forward Pass')
    print('='*60)
    
    try:
        _, env_t = env.reset()
        print('✓ Environment reset for agent test')
        
        rng, subkey = jax.random.split(rng)
        actor_timestep, new_actor_state = agent.actor_step(
            learner_state.params,
            subkey,
            env_t,
            actor_state,
        )
        print('✓ Agent step completed')
        
        print(f'  Action shape: {actor_timestep.actions.shape}')
        print(f'  Action dtype: {actor_timestep.actions.dtype}')
        print(f'  Sample action: {actor_timestep.actions[0]}')
        
        # Verify action is valid discrete index
        action = int(actor_timestep.actions[0])
        num_actions = env.single_action_spec().maximum + 1
        assert 0 <= action < num_actions, f'Action {action} out of range [0, {num_actions})'
        print(f'✓ Action is valid discrete index (0 <= {action} < {num_actions})')
        
        # Try a few more steps
        for step in range(3):
            rng, subkey = jax.random.split(rng)
            actor_timestep, new_actor_state = agent.actor_step(
                learner_state.params,
                subkey,
                env_t,
                new_actor_state,
            )
            _, env_t = env.step(None, actor_timestep.actions)
        
        print(f'✓ Ran 3 more agent steps successfully')
        
        return True
    except Exception as e:
        print(f'✗ Agent forward pass failed: {e}')
        import traceback
        traceback.print_exc()
        return False


def test_weights_loading():
    """Test 5: Load Disco103 weights."""
    print('\n' + '='*60)
    print('Test 5: Disco103 Weights Loading')
    print('='*60)
    
    try:
        weights = load_disco103_weights(
            disco_rl_path=os.path.join(repo_root, 'disco_rl')
        )
        print('✓ Disco103 weights loaded successfully')
        return True
    except FileNotFoundError as e:
        print(f'⚠ Disco103 weights not found (this is OK for PoC): {e}')
        return False
    except Exception as e:
        print(f'✗ Weight loading failed: {e}')
        import traceback
        traceback.print_exc()
        return False


def main():
    print('='*60)
    print('DiscoRL CartPole Sanity Check')
    print('='*60)
    
    # Test 1: Environment
    env = test_env_creation()
    if env is None:
        print('\n✗ Critical: Environment creation failed. Cannot proceed.')
        return False
    
    # Test 2: Reset/Step
    if not test_env_reset_step(env):
        print('\n✗ Critical: Environment reset/step failed.')
        return False
    
    # Test 3: Agent
    agent, learner_state, actor_state, rng = test_agent_creation(env)
    if agent is None:
        print('\n✗ Critical: Agent creation failed.')
        return False
    
    # Test 4: Agent forward
    if not test_agent_forward(env, agent, learner_state, actor_state, rng):
        print('\n✗ Critical: Agent forward pass failed.')
        return False
    
    # Test 5: Weights
    test_weights_loading()
    
    # Summary
    print('\n' + '='*60)
    print('Sanity Check Complete')
    print('='*60)
    print('✓ All core components working!')
    print('\nNext steps:')
    print('  1. Run: python3 train_disco_cartpole.py')
    print('  2. Run: python3 eval_disco_vs_sb3.py')
    
    return True


if __name__ == '__main__':
    success = main()
    sys.exit(0 if success else 1)