"""Train DiscoRL agent on CartPole using Disco103 discovered update rule.

This script demonstrates:
1. Wrapping CartPole in DiscoRL's Environment interface
2. Loading Disco103 pre-trained weights
3. Running training loop with DiscoRL's agent and update rule
4. Collecting rollouts and computing losses
5. Saving trained agent for later evaluation

Key insight: We use the Disco103 discovered update rule (meta-net) to guide
the policy/value network training, much like how it was used in the original
paper's meta-evaluation phase.

Design notes:
- DiscoRL expects discrete actions; CartPole is continuous but we discretize.
- We do NOT do meta-training (i.e., we don't update the meta-net itself).
  Instead, we use the pre-trained meta-net to generate update targets for
  the policy/value network.
- This is closer to "meta-evaluation" than "meta-training" in the paper's
  terminology.
"""

import os
import sys

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

from typing import Any, Tuple
import numpy as np
import jax
import jax.numpy as jnp
from ml_collections import config_dict

# Ensure DiscoRL is importable
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'))

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


class DiscoTrainingState:
    """Manages training state for DiscoRL agent."""
    
    def __init__(
        self,
        learner_state: Any,
        actor_state: Any,
        update_rule_params: Any,
        rng: jax.random.PRNGKey,
    ):
        self.learner_state = learner_state
        self.actor_state = actor_state
        self.update_rule_params = update_rule_params
        self.rng = rng


def rollout_trajectory(
    agent: disco_agent.Agent,
    env: DiscoCartPoleEnv,
    training_state: DiscoTrainingState,
    trajectory_length: int,
) -> Tuple[types.ActorRollout, DiscoTrainingState]:
    """Collect one trajectory of experience using the agent.
    
    Args:
        agent: DiscoRL Agent instance
        env: environment
        training_state: current training state
        trajectory_length: number of steps to collect
    
    Returns:
        (rollout, updated_training_state)
    
    The rollout is a types.ActorRollout containing observations, actions,
    rewards, agent outputs, etc., collected over the trajectory.
    """
    rng = training_state.rng
    learner_state = training_state.learner_state
    actor_state = training_state.actor_state
    update_rule_params = training_state.update_rule_params
    
    # Reset environment
    _, env_t = env.reset()
    
    # Collect trajectory
    observations = []
    actions = []
    rewards = []
    discounts = []  # 1.0 for non-terminal, 0.0 for terminal
    agent_outs_list = []
    logits_list = []
    states = []
    
    for step in range(trajectory_length):
        rng, subkey = jax.random.split(rng)
        
        # Agent step (policy + value network forward + action sampling)
        actor_timestep, actor_state = agent.actor_step(
            learner_state.params,
            subkey,
            env_t,
            actor_state,
        )
        
        # Record
        states.append(actor_state)
        observations.append(env_t.observation['observation'])
        agent_outs_list.append(actor_timestep.agent_outs)
        logits_list.append(actor_timestep.logits)
        actions.append(actor_timestep.actions)
        
        # Environment step
        rng, subkey = jax.random.split(rng)
        _, env_t = env.step(None, actor_timestep.actions)
        
        rewards.append(env_t.reward)
        # discount: 1.0 if not terminal, 0.0 if terminal
        # LAST = 2 in dm_env StepType convention
        discounts.append(1.0 - (env_t.step_type == 2).astype(jnp.float32))
    
    # Stack into batch dimensions
    # All shapes should be [T, B, ...] where T=trajectory_length, B=batch_size
    observations = jnp.stack(observations, axis=0)
    actions = jnp.stack(actions, axis=0)
    rewards = jnp.stack(rewards, axis=0)
    discounts = jnp.stack(discounts, axis=0)
    
    # For agent_outs, we need to stack each component
    agent_outs_stacked = jax.tree.map(
        lambda *xs: jnp.stack(xs, axis=0),
        *agent_outs_list,
    )
    
    logits = jnp.stack(logits_list, axis=0)
    
    # Construct rollout (behavior = current agent, so behaviour_agent_out = agent_out)
    rollout = types.ActorRollout(
        observations=observations,
        actions=actions,
        rewards=rewards,
        discounts=discounts,
        agent_outs=agent_outs_stacked,
        logits=logits,
        states=actor_state,  # we only keep the final actor state
    )
    
    # Update training state
    training_state = DiscoTrainingState(
        learner_state=learner_state,
        actor_state=actor_state,
        update_rule_params=update_rule_params,
        rng=rng,
    )
    
    return rollout, training_state


def main():
    # ===== Setup =====
    print('='*60)
    print('DiscoRL Training on CartPole')
    print('='*60)
    
    # Configuration
    batch_size = 4
    trajectory_length = 32
    num_iterations = 50  # number of training steps
    learning_rate = 1e-4
    
    # Create environment
    print('\nCreating environment...')
    env = DiscoCartPoleEnv(
        batch_size=batch_size,
        max_steps=500,
    )
    single_obs_spec = env.single_observation_spec()
    single_act_spec = env.single_action_spec()
    print(f'  Obs spec: {single_obs_spec}')
    print(f'  Act spec: {single_act_spec}')
    
    # Create agent with Disco103 settings
    print('\nCreating DiscoRL 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 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)
    
    # Initialize meta params from update rule (random initialization)
    print('\nInitializing meta-params...')
    rng, subkey = jax.random.split(rng)
    update_rule_params, _ = agent.update_rule.init_params(subkey)
    
    # ===== Training Loop =====
    print('\n' + '='*60)
    print('Starting Training')
    print('='*60)
    
    training_state = DiscoTrainingState(
        learner_state=learner_state,
        actor_state=actor_state,
        update_rule_params=update_rule_params,
        rng=rng,
    )
    
    cumulative_rewards = []
    
    for iteration in range(num_iterations):
        # Collect trajectory
        rollout, training_state = rollout_trajectory(
            agent=agent,
            env=env,
            training_state=training_state,
            trajectory_length=trajectory_length,
        )
        
        # Compute average reward in this rollout
        avg_reward = jnp.mean(rollout.rewards)
        cumulative_rewards.append(float(avg_reward))
        
        # Learner step (update network parameters using meta-net guidance)
        rng, subkey = jax.random.split(training_state.rng)
        new_learner_state, new_actor_state, log_dict = agent.learner_step(
            rng=subkey,
            rollout=rollout,
            learner_state=training_state.learner_state,
            agent_net_state=training_state.actor_state,
            update_rule_params=training_state.update_rule_params,
            is_meta_training=False,  # We use pre-trained meta-net, no meta-training
        )
        
        training_state = DiscoTrainingState(
            learner_state=new_learner_state,
            actor_state=new_actor_state,
            update_rule_params=training_state.update_rule_params,
            rng=rng,
        )
        
        # Logging
        if (iteration + 1) % 10 == 0:
            print(f'Iter {iteration+1:4d}: '
                  f'  avg_reward={avg_reward:7.3f}  '
                  f'  total_loss={log_dict.get("total_loss", 0):7.5f}')
        
        # Checkpointing (save every 50 iterations)
        if (iteration + 1) % 50 == 0:
            checkpoint_path = os.path.join(
                repo_root, 'models', 'disco_cartpole', f'checkpoint_{iteration+1}.npz'
            )
            os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
            # Save learner params (simplified; real implementation would also save state)
            np.savez(
                checkpoint_path,
                params=jax.tree.map(np.asarray, training_state.learner_state.params),
            )
            print(f'  Saved checkpoint to {checkpoint_path}')
    
    # ===== Final Results =====
    print('\n' + '='*60)
    print('Training Complete')
    print('='*60)
    print(f'Final avg reward (last 10 iters): {np.mean(cumulative_rewards[-10:]):.3f}')
    print(f'Max avg reward seen: {np.max(cumulative_rewards):.3f}')
    
    # Save final model
    final_path = os.path.join(
        repo_root, 'models', 'disco_cartpole', 'final_agent.npz'
    )
    os.makedirs(os.path.dirname(final_path), exist_ok=True)
    np.savez(
        final_path,
        params=jax.tree.map(np.asarray, training_state.learner_state.params),
    )
    print(f'Saved final agent to {final_path}')


if __name__ == '__main__':
    main()