import time
import curses
import numpy as np
import warnings
from environment import FrozenLakeEnvCustom
from agent import QLearningAgent

warnings.filterwarnings("ignore")


def train_agent(
    n_training_episodes,
    min_epsilon,
    max_epsilon,
    decay_rate,
    env,
    max_steps,
    agent,
    learning_rate,
    gamma,
    use_frame_delay,
):
    for episode in range(n_training_episodes + 1):

        # Reduce epsilon (because we need less and less exploration)
        epsilon = min_epsilon + (max_epsilon - min_epsilon) * np.exp(
            -decay_rate * episode
        )
        state, info = env.reset()
        done = False
        for step in range(max_steps):

            # Choose the action At using epsilon greedy policy
            action = agent.epsilon_greedy_policy(state, epsilon)

            # Take action At and observe Rt+1 and St+1
            # Take the action (a) and observe the outcome state(s') and reward (r)
            new_state, reward, done, truncated, info = env.step(action)
            agent.update_q_table(state, action, reward, gamma, learning_rate, new_state)

            env.render(
                title=f"Training: {episode}/{n_training_episodes}",
                q_table=agent.q_table,
            )

            if use_frame_delay:
                time.sleep(0.01)

            if done:
                break

            state = new_state
    return agent


def evaluate_agent(env, max_steps, n_eval_episodes, agent, seed, use_frame_delay):
    successful_episodes = []
    episodes_slips = []
    for episode in range(n_eval_episodes + 1):
        if seed:
            state, info = env.reset(seed=seed[episode])
        else:
            state, info = env.reset()
        done = False
        total_rewards_ep = 0

        slips = []
        for step in range(max_steps):

            # Take the action (index) that have the maximum expected future reward given that state
            action = agent.greedy_policy(state)

            expected_new_state = env.get_expected_new_state_for_action(action)
            new_state, reward, done, truncated, info = env.step(action)
            total_rewards_ep += reward

            if expected_new_state != new_state:
                slips.append((step, action, expected_new_state, new_state))

            if reward != 0:
                successful_episodes.append(episode)

            env.render(
                title=f"Evaluating: {episode}/{n_eval_episodes} | Slips: {len(slips)}",
                q_table=agent.q_table,
            )
            episodes_slips.append(len(slips))

            if use_frame_delay:
                time.sleep(0.01)

            if done:
                break
            state = new_state

    mean_slips = np.mean(episodes_slips)
    return successful_episodes, mean_slips


def main(screen):

    # Training parameters
    n_training_episodes = 2000  # Total training episodes
    learning_rate = 0.1  # Learning rate

    # Evaluation parameters
    n_eval_episodes = 100  # Total number of test episodes

    # Environment parameters
    max_steps = 99  # Max steps per episode
    gamma = 0.99  # Discounting rate
    eval_seed = []  # The evaluation seed of the environment

    # Exploration parameters
    max_epsilon = 1.0  # Exploration probability at start
    min_epsilon = 0.05  # Minimum exploration probability
    decay_rate = 0.0005  # Exponential decay rate for exploration prob

    use_frame_delay = False

    env = FrozenLakeEnvCustom(map_name="4x4", is_slippery=True, render_mode="curses")

    agent = QLearningAgent(env)
    agent = train_agent(
        n_training_episodes,
        min_epsilon,
        max_epsilon,
        decay_rate,
        env,
        max_steps,
        agent,
        learning_rate,
        gamma,
        use_frame_delay,
    )

    successful_episodes, mean_slips = evaluate_agent(
        env, max_steps, n_eval_episodes, agent, eval_seed, use_frame_delay
    )

    env_curses_screen = env.curses_screen
    env_curses_screen.addstr(
        5,
        2,
        f"Successful episodes: {len(successful_episodes)}/{n_eval_episodes} | Avg slips: {mean_slips:.2f}\n\n",
    )
    env_curses_screen.noutrefresh()
    curses.doupdate()
    time.sleep(10)


if __name__ == "__main__":

    # Reset the terminal state after using curses
    # Call main("") instead if you want to leave the final state of the environment
    # on the terminal
    curses.wrapper(main)