# -- Code Cell --
from collections import defaultdict
from tqdm import tqdm

# -- 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()
env = gym.make("MountainCar-v0", render_mode="human")

# -- Code Cell --
class TabularQLearning:
    def __init__(self, n_actions, alpha=0.1, gamma=0.99, epsilon=1.0, epsilon_min=0.01):
        self.Q = defaultdict(lambda: np.zeros(n_actions))
        self.alpha = alpha
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_min = epsilon_min
        self.n_actions = n_actions
        self.pos_bins = np.linspace(-1.2, 0.6, 20)
        self.vel_bins = np.linspace(-0.07, 0.07, 20)

    def discretize(self, state):
        position, velocity = state
        p = np.digitize(position, self.pos_bins)
        v = np.digitize(velocity, self.vel_bins)
        return (p, v)

    def choose_action(self, state):
        if np.random.rand() < self.epsilon:
            return np.random.randint(self.n_actions)

        state_key = self.discretize(state)
        return np.argmax(self.Q[state_key])

    def update(self, state, action, reward, next_state):
        state_key = self.discretize(state)
        next_state_key = self.discretize(next_state)

        best_next_q = np.max(self.Q[next_state_key])
        td_target = reward + self.gamma * best_next_q
        td_error = td_target - self.Q[state_key][action]

        self.Q[state_key][action] += self.alpha * td_error

    def decay(self, eps_decay=0.9999):
        self.epsilon = max(self.epsilon_min, self.epsilon * eps_decay)

# -- Code Cell --
agent = TabularQLearning(
    n_actions=env.action_space.n,
    alpha=0.1,
    gamma=0.99,
    epsilon=1.0
)

# -- Code Cell --
episodes = 50000
reward_window = []

for ep in tqdm(range(episodes)):
    state, _ = env.reset()
    done = False
    total_reward = 0

    while not done:
        action = agent.choose_action(state)
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated

        agent.update(state, action, reward, next_state)

        state = next_state
        total_reward += reward

    agent.decay()
    reward_window.append(total_reward)

    if ep % 1000 == 0:
        avg_reward = np.mean(reward_window[-100:]) if reward_window else 0
        print(f"Ep {ep} | States Discovered: {len(agent.Q)} | Avg Rew: {avg_reward:.1f} | ε: {agent.epsilon:.3f}")

states = []
for position in np.linspace(-1.2, 0.6, 10):
    for velocity in np.linspace(-0.07, 0.07, 10):
        states.append([position, velocity])

episode_returns = []

for i, init_state in enumerate(states):
    state, _ = env.reset(seed=i)
    env.unwrapped.state = np.array(init_state)
    state = np.array(init_state)

    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}")

# -- Code Cell --
import pandas as pd

# -- Code Cell --
df_test = pd.read_csv("test.csv")

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

    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

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

        state = next_state

    total_rewards.append(total_reward)

avg_reward = np.mean(total_rewards)
print("Average reward:", avg_reward)

df_logs = pd.DataFrame(logs)
df_logs.to_csv("subi.csv", index=False)