import pandas  as pd
import numpy as np
import gymnasium as gym
import numpy as np
from tqdm import tqdm

env = gym.make("MountainCar-v0")
df_logs = pd.read_csv("subi.csv")
df_test = pd.read_csv("test.csv")

episode_rewards = []

for episode_id in range(len(df_test)):

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

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

    total_reward = 0

    episode_actions = df_logs[df_logs["episode_id"] == episode_id].sort_values("episode_id")

    state = custom_state.copy()

    for _, row in episode_actions.iterrows():
        action = int(row["action"])

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

        total_reward += reward
        state = next_state

        if terminated or truncated:
            break

    episode_rewards.append(total_reward)

avg_reward = np.mean(episode_rewards)
def scale_reward(avg_reward):
    if avg_reward >= -130:
        return 100.0
    elif avg_reward <= -200:
        return 0.0
    else:
        return (avg_reward + 200) / (70) * 100 
print(f"YOUR SCORE IS: {scale_reward(avg_reward)}. CONGRATS!")