#!/usr/bin/env python3
"""Demo script: DiscoRL agent evaluation on CartPole.

This script loads a trained DiscoRL agent and evaluates it on CartPole.
Serves as a template for adapting DiscoRL to custom environments.
"""

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

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

# Add repo root 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 evaluate_agent(agent, env, num_episodes: int = 10, max_steps: int = 500):
    """Evaluate agent on environment.
    
    Args:
        agent: DiscoRL Agent
        env: DiscoCartPoleEnv
        num_episodes: number of evaluation episodes
        max_steps: max steps per episode
    
    Returns:
        (rewards_per_episode, success_rate)
    """
    rewards_per_episode = []
    successes = 0
    
    for episode in range(num_episodes):
        rng = jax.random.PRNGKey(episode)
        rng, subkey = jax.random.split(rng)
        
        # Reset
        state, timestep = env.reset(rng_key=subkey)
        actor_state = agent.initial_actor_state(subkey)
        
        episode_reward = 0.0
        
        for step in range(max_steps):
            # Agent step
            rng, subkey = jax.random.split(rng)
            actor_output, actor_state = agent.actor_step(
                timestep.observation,
                actor_state,
                is_eval=True,
                training_state=None,
                rng=subkey,
            )
            actions = actor_output.actions
            
            # Env step
            state, timestep = env.step(state, actions)
            
            # Accumulate reward
            episode_reward += float(jnp.mean(timestep.reward))
            
            # Check terminal
            if jnp.any(timestep.step_type == 1):  # StepType.LAST
                break
        
        rewards_per_episode.append(episode_reward)
        if episode_reward > 400:  # CartPole "solved" at 400+ steps
            successes += 1
    
    success_rate = successes / num_episodes
    return rewards_per_episode, success_rate


def main():
    print('='*60)
    print('DiscoRL CartPole Evaluation Demo')
    print('='*60)
    
    # Create environment
    print('\nCreating environment...')
    env = DiscoCartPoleEnv(batch_size=1, max_steps=500)
    single_obs_spec = env.single_observation_spec()
    single_act_spec = env.single_action_spec()
    
    # Create agent
    print('Creating agent...')
    agent_settings = disco_agent.get_settings_disco()
    agent = disco_agent.Agent(
        single_observation_spec=single_obs_spec,
        single_action_spec=single_act_spec,
        agent_settings=agent_settings,
        batch_axis_name=None,
    )
    
    # Initialize state
    print('Initializing agent state...')
    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)
    
    # Try to load saved weights
    saved_path = os.path.join(repo_root, 'models', 'disco_cartpole', 'final_agent.npz')
    if os.path.exists(saved_path):
        print(f'\nLoading saved agent from {saved_path}...')
        try:
            saved = np.load(saved_path)
            # Note: You'd need to implement proper deserialization
            # For now, just note that weights are available
            print(f'  Saved weights available: {list(saved.files)}')
        except Exception as e:
            print(f'  Warning: Could not load weights: {e}')
    else:
        print(f'\nNo saved weights found at {saved_path}')
        print('Using random initialization for this demo.')
    
    # Evaluate
    print('\n' + '='*60)
    print('Evaluating Agent')
    print('='*60)
    
    rewards, success_rate = evaluate_agent(agent, env, num_episodes=10)
    
    print(f'\nResults (10 episodes):')
    print(f'  Mean reward: {np.mean(rewards):.2f}')
    print(f'  Max reward: {np.max(rewards):.2f}')
    print(f'  Min reward: {np.min(rewards):.2f}')
    print(f'  Success rate: {success_rate:.1%}')
    print(f'\nRewards by episode:')
    for i, r in enumerate(rewards):
        print(f'    Episode {i+1:2d}: {r:6.1f}')
    
    print('\n' + '='*60)
    print('Demo Complete')
    print('='*60)


if __name__ == '__main__':
    main()