#!/usr/bin/env python3
"""Minimal DiscoRL CartPole integration test.

This is a minimal script that demonstrates that DiscoRL can:
1. Load CartPole environment
2. Create an agent
3. Collect trajectories
4. Perform training steps

This serves as a proof-of-concept for integrating DiscoRL with other Gym environments.
"""

import os
import sys

# Set JAX to CPU
os.environ['JAX_PLATFORMS'] = 'cpu'

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

print('='*70)
print('DiscoRL ↔ Gym Integration: CartPole Proof-of-Concept')
print('='*70)

print('\n[1/4] Importing modules...')
try:
    import numpy as np
    import jax
    import jax.numpy as jnp
    from disco_rl import agent as disco_agent
    from disco_cartpole_env import DiscoCartPoleEnv
    print('  ✓ All imports successful')
except Exception as e:
    print(f'  ✗ Import failed: {e}')
    sys.exit(1)

print('\n[2/4] Creating DiscoRL environment adapter...')
try:
    env = DiscoCartPoleEnv(batch_size=2, max_steps=500)
    obs_spec = env.single_observation_spec()
    act_spec = env.single_action_spec()
    print(f'  ✓ Environment created')
    print(f'    Observation spec: {obs_spec}')
    print(f'    Action spec: {act_spec}')
except Exception as e:
    print(f'  ✗ Environment creation failed: {e}')
    sys.exit(1)

print('\n[3/4] Initializing DiscoRL agent...')
try:
    agent_settings = disco_agent.get_settings_disco()
    agent = disco_agent.Agent(
        single_observation_spec=obs_spec,
        single_action_spec=act_spec,
        agent_settings=agent_settings,
        batch_axis_name=None,
    )
    
    rng = jax.random.PRNGKey(42)
    rng, subkey = jax.random.split(rng)
    learner_state = agent.initial_learner_state(subkey)
    
    rng, subkey = jax.random.split(rng)
    actor_state = agent.initial_actor_state(subkey)
    
    rng, subkey = jax.random.split(rng)
    update_rule_params, _ = agent.update_rule.init_params(subkey)
    
    print(f'  ✓ Agent initialized')
    print(f'    Learner params shape: {jax.tree.map(lambda x: x.shape, learner_state.params)}')
except Exception as e:
    print(f'  ✗ Agent initialization failed: {e}')
    import traceback
    traceback.print_exc()
    sys.exit(1)

print('\n[4/4] Collecting trajectory and performing learner step...')
try:
    # Reset environment
    env_state, timestep = env.reset(rng_key=subkey)
    print(f'  ✓ Environment reset')
    
    # Collect 8 timesteps
    trajectory_length = 8
    observations = []
    actions = []
    rewards = []
    discounts = []
    agent_outs_list = []
    logits_list = []
    
    for t in range(trajectory_length):
        rng, subkey = jax.random.split(rng)
        
        # Agent step
        actor_timestep, actor_state = agent.actor_step(
            learner_state.params,
            subkey,
            timestep,
            actor_state,
        )
        
        # Record
        observations.append(timestep.observation['observation'])
        actions.append(actor_timestep.actions)
        agent_outs_list.append(actor_timestep.agent_outs)
        logits_list.append(actor_timestep.logits)
        
        # Env step
        rng, subkey = jax.random.split(rng)
        env_state, timestep = env.step(env_state, actor_timestep.actions)
        
        rewards.append(timestep.reward)
        discounts.append(1.0 - (timestep.step_type == 1).astype(jnp.float32))
    
    # Stack trajectory
    observations = jnp.stack(observations, axis=0)
    actions = jnp.stack(actions, axis=0)
    rewards = jnp.stack(rewards, axis=0)
    discounts = jnp.stack(discounts, axis=0)
    agent_outs_stacked = jax.tree.map(
        lambda *xs: jnp.stack(xs, axis=0),
        *agent_outs_list,
    )
    logits = jnp.stack(logits_list, axis=0)
    
    print(f'  ✓ Trajectory collected ({trajectory_length} steps)')
    print(f'    Observations shape: {observations.shape}')
    print(f'    Actions shape: {actions.shape}')
    print(f'    Rewards: mean={jnp.mean(rewards):.2f}, range=[{jnp.min(rewards):.2f}, {jnp.max(rewards):.2f}]')
    
    # Create rollout
    from disco_rl import types
    rollout = types.ActorRollout(
        observations=observations,
        actions=actions,
        rewards=rewards,
        discounts=discounts,
        agent_outs=agent_outs_stacked,
        logits=logits,
        states=actor_state,
    )
    
    print(f'  ✓ Rollout created')
    
    # Learner step
    rng, subkey = jax.random.split(rng)
    new_learner_state, new_actor_state, log_dict = agent.learner_step(
        rng=subkey,
        rollout=rollout,
        learner_state=learner_state,
        agent_net_state=actor_state,
        update_rule_params=update_rule_params,
        is_meta_training=False,
    )
    
    print(f'  ✓ Learner step completed')
    print(f'    Loss: {log_dict.get("total_loss", "N/A")}')
    
except Exception as e:
    print(f'  ✗ Trajectory collection or learner step failed: {e}')
    import traceback
    traceback.print_exc()
    sys.exit(1)

print('\n' + '='*70)
print('✓ Success! DiscoRL ↔ Gym integration works!')
print('='*70)
print(f'''
Integration achieved:
  • Gym CartPole environment wrapped for DiscoRL
  • Agent can collect trajectories from Gym environments
  • Learner can perform parameter updates using Gym trajectories
  
Template for custom environments:
  1. Create environment wrapper (see disco_cartpole_env.py)
     - Implement reset() → (state, types.EnvironmentTimestep)
     - Implement step(state, actions) → (state, types.EnvironmentTimestep)
  2. Use DiscoRL agent as shown above
  3. Adapt to your custom environment in gym_env_250326_erase.py

Files created:
  • disco_cartpole_env.py: CartPole ↔ DiscoRL adapter
  • disco_weights.py: Weights loading utilities
  • train_disco_cartpole.py: Training loop example
  • test_disco_setup.py: Comprehensive sanity tests
''')