#!/usr/bin/env python3
"""Simple DiscoRL CartPole inference example.

Shows how to use a trained DiscoRL agent for policy inference on CartPole.
"""

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

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

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

from disco_rl import agent as disco_agent
from disco_cartpole_env import DiscoCartPoleEnv


def rollout_policy(agent, learner_state, env, num_steps: int = 100):
    """Roll out policy to collect trajectory.
    
    Args:
        agent: DiscoRL Agent
        learner_state: learned parameters
        env: DiscoCartPoleEnv
        num_steps: number of steps to collect
    
    Returns:
        (total_reward, trajectory_length)
    """
    rng = jax.random.PRNGKey(0)
    rng, subkey = jax.random.split(rng)
    
    # Reset environment
    state, timestep = env.reset(rng_key=subkey)
    
    # Initialize actor state
    rng, subkey = jax.random.split(rng)
    actor_state = agent.initial_actor_state(subkey)
    
    total_reward = 0.0
    for step in range(num_steps):
        # Get action from agent using learned params
        rng, subkey = jax.random.split(rng)
        actor_timestep, actor_state = agent.actor_step(
            learner_state.params,
            subkey,
            timestep,
            actor_state,
        )
        
        # Step environment
        state, timestep = env.step(state, actor_timestep.actions)
        
        # Accumulate reward
        total_reward += float(jnp.mean(timestep.reward))
        
        # Terminal check
        if jnp.any(timestep.step_type == 1):
            break
    
    return total_reward, step + 1


def main():
    print('='*60)
    print('DiscoRL CartPole Inference Example')
    print('='*60)
    
    # Setup
    print('\nSetting up...')
    env = DiscoCartPoleEnv(batch_size=1, max_steps=500)
    
    agent_settings = disco_agent.get_settings_disco()
    agent = disco_agent.Agent(
        single_observation_spec=env.single_observation_spec(),
        single_action_spec=env.single_action_spec(),
        agent_settings=agent_settings,
        batch_axis_name=None,
    )
    
    # Initialize learner state
    rng = jax.random.PRNGKey(42)
    rng, subkey = jax.random.split(rng)
    learner_state = agent.initial_learner_state(subkey)
    
    print('\nRunning policy rollouts...')
    
    # Run 5 rollouts
    results = []
    for i in range(5):
        reward, steps = rollout_policy(agent, learner_state, env, num_steps=500)
        results.append((reward, steps))
        print(f'  Rollout {i+1}: reward={reward:7.1f}, steps={steps:3d}')
    
    # Summary
    rewards = [r for r, _ in results]
    print(f'\nSummary:')
    print(f'  Mean reward: {np.mean(rewards):.1f}')
    print(f'  Max reward:  {np.max(rewards):.1f}')
    print(f'  Min reward:  {np.min(rewards):.1f}')
    print(f'  Std:         {np.std(rewards):.1f}')
    
    print('\n✓ Inference example complete!')


if __name__ == '__main__':
    main()