"""Debug script to check what actions agent takes."""

import sys
sys.path.insert(0, '/home/frank14f/Frank_LBM/scripts')
sys.path.insert(0, '/home/frank14f/Frank_LBM')

import jax
import jax.numpy as jnp
import numpy as np
from disco_cartpole_env import DiscoCartPoleEnv
import disco_rl.agent as disco_agent
import disco_rl.types as types

# Create environment
env = DiscoCartPoleEnv(batch_size=1, max_steps=500)

# Create agent
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 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)

# Reset environment
_, env_t = env.reset()

print("Manual trajectory collection:")
print(f"Initial env_t.step_type: {env_t.step_type}")
print(f"Initial env_t.reward: {env_t.reward}")

for step in range(20):
    rng, subkey = jax.random.split(rng)
    
    # Get action from agent
    actor_timestep, actor_state = agent.actor_step(
        learner_state.params,
        subkey,
        env_t,
        actor_state,
    )
    
    action = actor_timestep.actions[0]
    
    # Step environment
    rng, subkey = jax.random.split(rng)
    _, env_t = env.step(None, actor_timestep.actions)
    
    print(f"Step {step+1:2d}: action={int(action)}, reward={env_t.reward[0]:.1f}, step_type={env_t.step_type[0]} (0=FIRST, 1=MID, 2=LAST), done={env._episode_done[0]}")
    
    if env._episode_done[0]:
        print(f"  -> Episode done at step {step+1}!")