# -- Markdown Cell --
# # Load environment and how to play with it

# -- Code Cell --
import gymnasium as gym
import numpy as np

env = gym.make("MountainCar-v0", render_mode="human")

print("=== ENVIRONMENT INFO ===")
print("Action space:", env.action_space)     # Discrete(3)
print("Observation space:", env.observation_space)  # Box(2,)

print("\n=== ACTIONS ===")
print("0 -> Push Left")
print("1 -> No Push")
print("2 -> Push Right")

print("\n=== STARTING EPISODE ===")

obs, info = env.reset()
print("Initial state:", obs)

total_reward = 0
done = False
step = 0

while not done:
    # 🔹 Random action (you can replace with your agent)
    action = env.action_space.sample()

    # 🔹 Take step
    next_obs, reward, terminated, truncated, info = env.step(action)

    done = terminated or truncated

    print(f"\nStep {step}")
    print("Current state:", obs)
    print("Action taken:", action)
    print("Reward:", reward)
    print("Next state:", next_obs)

    total_reward += reward
    obs = next_obs
    step += 1

print("\n=== EPISODE FINISHED ===")
print("Total reward:", total_reward)

env.close()

# -- Markdown Cell --
# # Evaluate your agent

# -- Code Cell --
import numpy as np

episode_returns = []

for i, init_state in enumerate(states):
    state, _ = env.reset(seed=i)
    
    total_reward = 0
    done = False
    
    while not done:
        s = agent.discretize(state)
        action = np.argmax(agent.Q[s])
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        total_reward += reward
        state = next_state
    
    episode_returns.append(total_reward)
episode_returns = np.array(episode_returns)
print(f"Mean return:   {episode_returns.mean():.2f}")

# -- Markdown Cell --
# # Make a submission 

# -- Code Cell --
import numpy as np
import pandas as pd

total_rewards = []

logs = [] 

for i in range(len(df_test)):
    state_row = df_test.iloc[i]

    custom_state = np.array([
        state_row["start_position"],
        state_row["start_velocity"]
    ])

    obs, info = env.reset()
    env.unwrapped.state = custom_state

    state = custom_state.copy()
    done = False
    total_reward = 0
    step_id = 0 

    while not done:
        action = agent.choose_action(state, env.action_space.n)

        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated

        total_reward += reward

        logs.append({
            "episode_id": i,
            "action": action,
        })

        state = next_state
        step_id += 1 

    total_rewards.append(total_reward)

avg_reward = np.mean(total_rewards)

print("Average reward:", avg_reward)

df_logs = pd.DataFrame(logs)

# -- Markdown Cell --
# ### Note: For you to run correctly eval.py, you need to name your .csv as subi.csv