"""Example: run DiscoRL agent inference on a Gym env (PoC).

This script demonstrates how to:
 - discretize a continuous Gym action space (quick PoC),
 - adapt the Gym env to DiscoRL's Environment interface using
   `disco_gym_adapter.GymToDiscoEnv`, and
 - run a few actor steps with the DiscoRL Agent (inference only).

Important constraints and notes:
 - The provided DiscoRL code expects a scalar discrete action space by
   design (see Agent.__init__). To avoid large code changes we discretize
   your continuous action space into a small action set.
 - This example performs inference only using the freshly initialised
   agent parameters (no meta-training). If you want to run meta-training or
   reproduce the paper, keep using the JAX harness and the original
   training notebooks / scripts.
"""

import os
import sys

# ============================================================================
# JAX/CUDA Configuration: Prevent GPU conflicts between JAX and LBM simulator.
# If you encounter segmentation faults, try setting JAX_PLATFORM_NAME=cpu
# in your shell environment before running this script.
# ============================================================================
# Uncomment below to force JAX to CPU-only mode (safest for PoC):
os.environ['JAX_PLATFORMS'] = 'cpu'

import time
import numpy as np
import jax

# Make sure DiscoRL package is importable. If you installed it with
# `pip install -e ./disco_rl` then the plain import below will work. If not,
# try adding the repo-local path to sys.path.
current_dir = os.path.dirname(os.path.abspath(__file__))
repo_root = os.path.abspath(os.path.join(current_dir, os.pardir))
sys.path.append(os.path.join(repo_root, "disco_rl"))

from disco_gym_adapter import GymToDiscoEnv
from gym_env_250326_erase import CustomEnv

from disco_rl import agent as disco_agent


def main():
  # Choose a small discrete action set (example: grid over each continuous dim).
  # Here the original env has action shape (3,) with each element in [-1,1].
  # We build a simple 5-action set: no-op + ±1 on each axis (sparse).
  discrete_actions = np.array([
      [0.0, 0.0, 0.0],
      [1.0, 0.0, 0.0],
      [-1.0, 0.0, 0.0],
      [0.0, 1.0, 0.0],
      [0.0, -1.0, 0.0],
  ], dtype=np.float32)

  # Build the Disco adapter with discrete actions (batch_size=1).
  base_kwargs = dict(device_id=2)
  disco_env = GymToDiscoEnv(
      gym_env_cls=lambda **kw: CustomEnv(**base_kwargs),
      batch_size=1,
      env_settings=None,
      discrete_actions=discrete_actions,
  )

  # Create agent and settings.
  settings = disco_agent.get_settings_disco()

  # Get action/observation specs from the adapted env.
  single_obs_spec = disco_env.single_observation_spec()
  single_act_spec = disco_env.single_action_spec()
  
  print('Action spec:', single_act_spec)
  print('Observation spec:', single_obs_spec)

  # Create the DiscoRL agent.
  agent = disco_agent.Agent(
      single_observation_spec=single_obs_spec,
      single_action_spec=single_act_spec,
      agent_settings=settings,
      batch_axis_name=None,
  )

  rng = jax.random.PRNGKey(0)
  learner_state = agent.initial_learner_state(rng)
  actor_state = agent.initial_actor_state(rng)

  # Reset environment and run a small rollout.
  _, env_t = disco_env.reset()
  print('\nInitial observation shape:', env_t.observation['observation'].shape)
  print('Initial observation:', env_t.observation)

  print('\nRunning 10 inference steps...')
  for step in range(10):
    rng, subkey = jax.random.split(rng)
    actor_timestep, actor_state = agent.actor_step(
        learner_state.params, subkey, env_t, actor_state
    )
    # actor_timestep.actions is a JAX array (discrete indices). Convert to numpy.
    action = np.asarray(actor_timestep.actions)[0]
    # Step the Gym env via the adapter: pass the discrete action index.
    _, env_t = disco_env.step(None, np.array([action]))
    reward = np.asarray(env_t.reward)
    print(f'  step {step:2d}  action={int(action)}  reward={reward[0] if hasattr(reward, "__len__") else reward:.4f}')
  
  print('\nPoC completed successfully!')

if __name__ == '__main__':
  try:
    main()
  except Exception as e:
    print(f'Error during execution: {e}')
    import traceback
    traceback.print_exc()