"""Evaluation & comparison: DiscoRL vs SB3 PPO on CartPole.

This script:
1. Trains a standard SB3 PPO agent on CartPole (baseline)
2. Evaluates the DiscoRL-trained agent
3. Compares performance metrics
4. Provides visualization / reporting
"""

import os
import sys

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

import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List, Tuple

# Ensure imports work
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 gymnasium as gym
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from disco_cartpole_env import CartPoleDiscoWrapper, DiscoCartPoleEnv
import jax
import jax.numpy as jnp
from disco_rl import agent as disco_agent
from disco_weights import load_disco103_weights


def train_sb3_ppo(
    total_timesteps: int = 50000,
    n_steps: int = 2048,
    batch_size: int = 64,
    learning_rate: float = 3e-4,
) -> PPO:
    """Train SB3 PPO agent on CartPole.
    
    Returns:
        trained PPO model
    """
    print('\n' + '='*60)
    print('Training SB3 PPO Baseline')
    print('='*60)
    
    env = make_vec_env('CartPole-v1', n_envs=4)
    
    model = PPO(
        'MlpPolicy',
        env,
        n_steps=n_steps,
        batch_size=batch_size,
        learning_rate=learning_rate,
        verbose=1,
        device='cpu',  # or 'cuda:0' if GPU available
    )
    
    model.learn(total_timesteps=total_timesteps)
    env.close()
    
    print('SB3 PPO training complete')
    return model


def evaluate_sb3_ppo(model: PPO, num_episodes: int = 10) -> Dict[str, float]:
    """Evaluate trained SB3 PPO.
    
    Returns:
        dict with 'mean_reward', 'std_reward', etc.
    """
    env = gym.make('CartPole-v1')
    
    episode_rewards = []
    episode_lengths = []
    
    for ep in range(num_episodes):
        obs, _ = env.reset()
        done = False
        ep_reward = 0
        ep_len = 0
        
        while not done and ep_len < 500:
            action, _ = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            done = terminated or truncated
            ep_reward += reward
            ep_len += 1
        
        episode_rewards.append(ep_reward)
        episode_lengths.append(ep_len)
    
    env.close()
    
    results = {
        'mean_reward': float(np.mean(episode_rewards)),
        'std_reward': float(np.std(episode_rewards)),
        'mean_length': float(np.mean(episode_lengths)),
        'std_length': float(np.std(episode_lengths)),
        'min_reward': float(np.min(episode_rewards)),
        'max_reward': float(np.max(episode_rewards)),
    }
    
    print(f'  Mean reward: {results["mean_reward"]:.1f} ± {results["std_reward"]:.1f}')
    print(f'  Mean length: {results["mean_length"]:.1f} ± {results["std_length"]:.1f}')
    print(f'  Min/Max reward: {results["min_reward"]:.1f} / {results["max_reward"]:.1f}')
    
    return results


def evaluate_disco_agent(
    agent_params: Dict,
    update_rule_params: Dict,
    num_episodes: int = 10,
) -> Dict[str, float]:
    """Evaluate trained DiscoRL agent.
    
    Returns:
        dict with 'mean_reward', 'std_reward', etc.
    """
    # Create env and agent
    env = DiscoCartPoleEnv(batch_size=1)
    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 actor state (same across episodes)
    rng = jax.random.PRNGKey(0)
    actor_state_template = agent.initial_actor_state(rng)
    
    episode_rewards = []
    episode_lengths = []
    
    for ep in range(num_episodes):
        _, env_t = env.reset()
        rng, subkey = jax.random.split(rng)
        actor_state = actor_state_template
        
        ep_reward = 0.0
        ep_len = 0
        done = False
        
        while not done and ep_len < 500:
            rng, subkey = jax.random.split(rng)
            actor_timestep, actor_state = agent.actor_step(
                agent_params,
                subkey,
                env_t,
                actor_state,
            )
            
            action = np.asarray(actor_timestep.actions)[0]
            _, env_t = env.step(None, [action])
            
            done = bool(np.asarray(env_t.step_type)[0] == 1)
            reward = float(np.asarray(env_t.reward)[0])
            ep_reward += reward
            ep_len += 1
        
        episode_rewards.append(ep_reward)
        episode_lengths.append(ep_len)
    
    results = {
        'mean_reward': float(np.mean(episode_rewards)),
        'std_reward': float(np.std(episode_rewards)),
        'mean_length': float(np.mean(episode_lengths)),
        'std_length': float(np.std(episode_lengths)),
        'min_reward': float(np.min(episode_rewards)),
        'max_reward': float(np.max(episode_rewards)),
    }
    
    print(f'  Mean reward: {results["mean_reward"]:.1f} ± {results["std_reward"]:.1f}')
    print(f'  Mean length: {results["mean_length"]:.1f} ± {results["std_length"]:.1f}')
    print(f'  Min/Max reward: {results["min_reward"]:.1f} / {results["max_reward"]:.1f}')
    
    return results


def plot_comparison(sb3_results: Dict, disco_results: Dict, save_path: str = None):
    """Plot comparison results."""
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    
    methods = ['SB3 PPO', 'DiscoRL']
    means = [sb3_results['mean_reward'], disco_results['mean_reward']]
    stds = [sb3_results['std_reward'], disco_results['std_reward']]
    
    # Reward plot
    axes[0].bar(methods, means, yerr=stds, capsize=5, alpha=0.7, color=['blue', 'orange'])
    axes[0].set_ylabel('Mean Episode Reward')
    axes[0].set_title('Episode Reward Comparison')
    axes[0].grid(True, alpha=0.3)
    
    # Length plot
    lengths = [sb3_results['mean_length'], disco_results['mean_length']]
    length_stds = [sb3_results['std_length'], disco_results['std_length']]
    axes[1].bar(methods, lengths, yerr=length_stds, capsize=5, alpha=0.7, color=['blue', 'orange'])
    axes[1].set_ylabel('Mean Episode Length')
    axes[1].set_title('Episode Length Comparison')
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=100, bbox_inches='tight')
        print(f'\nSaved comparison plot to {save_path}')
    else:
        plt.show()


def main():
    print('='*60)
    print('DiscoRL vs SB3 PPO on CartPole')
    print('='*60)
    
    # Train SB3 baseline
    print('\n[1/4] Training SB3 PPO...')
    sb3_model = train_sb3_ppo(total_timesteps=50000)
    
    # Evaluate SB3
    print('\n[2/4] Evaluating SB3 PPO...')
    sb3_results = evaluate_sb3_ppo(num_episodes=20)
    
    # Try to load DiscoRL agent from checkpoint
    print('\n[3/4] Loading DiscoRL agent...')
    checkpoint_path = os.path.join(repo_root, 'models', 'disco_cartpole', 'final_agent.npz')
    
    if os.path.exists(checkpoint_path):
        print(f'  Found checkpoint at {checkpoint_path}')
        data = np.load(checkpoint_path, allow_pickle=True)
        agent_params = jax.tree.map(jnp.asarray, data['params'].item())
    else:
        print(f'  Checkpoint not found at {checkpoint_path}')
        print('  Using randomly initialized agent params (will not be competitive).')
        env = DiscoCartPoleEnv(batch_size=1)
        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,
        )
        rng = jax.random.PRNGKey(0)
        learner_state = agent.initial_learner_state(rng)
        agent_params = learner_state.params
    
    # Load meta params
    try:
        disco_weights = load_disco103_weights(
            disco_rl_path=os.path.join(repo_root, 'disco_rl')
        )
        update_rule_params = disco_weights
    except FileNotFoundError:
        print('  Warning: Could not load Disco103 weights; using random initialization.')
        update_rule_params = None
    
    # Evaluate DiscoRL
    print('\n[4/4] Evaluating DiscoRL agent...')
    disco_results = evaluate_disco_agent(agent_params, update_rule_params, num_episodes=20)
    
    # Comparison summary
    print('\n' + '='*60)
    print('Comparison Summary')
    print('='*60)
    print(f'{"Method":<15} {"Mean Reward":<20} {"Mean Length":<20}')
    print('-'*55)
    print(f'{"SB3 PPO":<15} {sb3_results["mean_reward"]:<20.1f} {sb3_results["mean_length"]:<20.1f}')
    print(f'{"DiscoRL":<15} {disco_results["mean_reward"]:<20.1f} {disco_results["mean_length"]:<20.1f}')
    
    # Plot comparison
    plot_save_path = os.path.join(repo_root, 'output', 'disco_vs_sb3_comparison.png')
    os.makedirs(os.path.dirname(plot_save_path), exist_ok=True)
    plot_comparison(sb3_results, disco_results, save_path=plot_save_path)


if __name__ == '__main__':
    main()