# -- Code Cell --
import gymnasium as gym

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

env = gym.make("MountainCar-v0")

N_BINS = 20
position_bins = np.linspace(-1.2, 0.6, N_BINS)
velocity_bins = np.linspace(-0.07, 0.07, N_BINS)

class DiscretizedQLearning:
    def __init__(self, bins, n_actions, alpha=0.1, gamma=0.99, epsilon=1.0, epsilon_min=0.01):
        self.bins = bins
        self.Q = np.ones([len(bins[0]) + 1, len(bins[1]) + 1, n_actions]) * 5.0
        self.alpha = alpha
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_min = epsilon_min  

    def discretize(self, state):
        return tuple(np.digitize(s, b) for s, b in zip(state, self.bins))

    def choose_action(self, state, n_actions):
        if np.random.rand() < self.epsilon:
            return np.random.randint(n_actions)
        s = self.discretize(state)
        return np.argmax(self.Q[s])

    def update(self, state, action, reward, next_state):
        s = self.discretize(state)
        s_next = self.discretize(next_state)
        best_next = np.max(self.Q[s_next])
        self.Q[s][action] += self.alpha * (
            reward + self.gamma * best_next - self.Q[s][action]
        )

    def decay(self, ep, eps_decay=0.9995, alpha_decay=0.9999):
        self.epsilon = max(self.epsilon_min, self.epsilon * eps_decay)
        self.alpha = max(0.01, self.alpha * alpha_decay)



# -- Code Cell --
agent = DiscretizedQLearning(
    bins=[position_bins, velocity_bins],
    n_actions=env.action_space.n,
    alpha=0.4,
    gamma=0.99,
    epsilon=1.0,
    epsilon_min=0.01,
)
episodes = 80000
reward_window = []

# -- Code Cell --

for ep in tqdm(range(episodes)):
    state, _ = env.reset()
    done = False
    total_reward = 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

        agent.update(state, action, reward, next_state)
        total_reward += reward 
        state = next_state

    agent.decay(ep)  
    reward_window.append(total_reward)
    if len(reward_window) > 100:
        reward_window.pop(0)
    if ep % 500 == 0:
        print(f"Ep {ep:5d} | Avg(100): {np.mean(reward_window):7.1f} | ε={agent.epsilon:.3f} | α={agent.alpha:.4f}")

# -- 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)
df_logs.to_csv("subi.csv", index = False)