import os
import time

import cv2
import torch
import numpy as np 
import imageio
import wandb
from datetime import datetime

from ddqn_agent import Agent
from atari_env import ALEEnvironment
from args import HIntArgumentParser
from metacontroller import MetaController


GPU_DEVICE = 4
VERBOSE = 2
DEBUG = False
TEST = False

STOP_TRAINING_THRESHOLD = 0.90
EPISODES = 1000000
MAX_EPISODE_STEPS = 1000
RANDOM_PLAY_STEPS = 20000

# Use subgoals from gaze analysis 
GOAL_FEATURE = True
GOAL_MODE = 'single_agent' 
ALL_SUBGOALS = np.loadtxt('subgoals.txt', dtype=int, delimiter=',')
SUBGOAL_ORDER = [8, 6, 1, 0, 2, 7, 2, 0, 1, 6, 8, 9] 
GOALS = len(SUBGOAL_ORDER) #if GOAL_MODE == 'full_sequence' else len(np.unique(SUBGOAL_ORDER))
SUBGOALs = SUBGOAL_ORDER #if GOAL_MODE == 'full_sequence' else np.array(SUBGOAL_ORDER)[np.sort(np.unique(SUBGOAL_ORDER, return_index=True)[1])] # unsorted unique -> keep subgoal appearance order 
TRAINED_GOALS = [False] * GOALS

def main():

    # init random seed 
    RANDOM_SEED = 9 #np.random.randint(100)
    print(f'Setting random seed to {RANDOM_SEED}')
    os.environ['PYTHONHASHSEED']=str(RANDOM_SEED)
    np.random.seed(RANDOM_SEED)
    torch.manual_seed(RANDOM_SEED)

    x = datetime.now()

    TIME_STAMP = x.strftime("%d%b%y-%H%M%S")
    MODEL_NAME = 'hInt-RL-single-dist_' + str(RANDOM_SEED) + '_' + TIME_STAMP

    os.environ['CUDA_VISIBLE_DEVICES']=str(GPU_DEVICE)
    os.environ['CUDA_LAUNCH_BLOCKING'] = str(1.0)

    actionMap = [0, 1, 2, 3, 4, 5, 11, 12]
    actionExplain = ['no action', 'jump', 'up', 'right', 'left', 'down', 'jump right', 'jump left']
    actionVectors = [[0,0], [0.0,-1.0], [0.0,-1.0], [1.0,0.0], [-1.0,0.0], [0.0,1.0], [0.7071067811865475,-0.7071067811865475], [-0.7071067811865475,-0.7071067811865475]]
    goalExplain = {8: 'start', 6: 'rope', 1: 'lower right ladder', 0: 'danger zone', 2: 'lower left ladder',  7: 'key', 9:'left door'}

    inv_label_mapping = {}    
    for i, l in enumerate(SUBGOALs):  # np.unique without sorting the values 
        inv_label_mapping[l] = i
    print(inv_label_mapping)

    subgoal_success_tracker = [[0] for i in range(GOALS)]
    subgoal_success_steps = [[] for i in range(GOALS)]
    subgoal_trailing_performance = [0.0 for _ in range(GOALS)]
    
    parser = HIntArgumentParser()
    parser.set_common_args()

    args = parser.parse_args(['--model_name', MODEL_NAME, '--verbose', str(VERBOSE), '--random_seed', str(RANDOM_SEED), 
                              '--num_goals', str(GOALS), '--goal_mode', GOAL_MODE, 
                              '--epsilon_start', str(1.0), '--epsilon_end', str(0.1), '--epsilon_steps', str(2000000)])
    
    if not DEBUG:
        wandb.init(project="hInt-RL", config=vars(args), name=MODEL_NAME)

    print(*[f'\n[PARAM] {k}: {v}' for (k,v) in vars(args).items()])
    print()

    env = ALEEnvironment(args, device=f"cuda:{GPU_DEVICE}")

    input_shape = env.getStackedState().shape
    #input_shape = np.array([input_shape[2], input_shape[1], input_shape[0]])
    print(f'Input shape {input_shape}\n')
   
    device = torch.device(f"cuda:{GPU_DEVICE}" if torch.cuda.is_available() else "cpu")

    # Setup metacontroller
    metacontroller = MetaController(device=device, args=args, input_shape=input_shape, n_goals=GOALS, hidden_nodes=args.hidden_nodes)
    
    # Setup single agent 
    agent = Agent(device=device, goal=None, input_shape=input_shape, n_actions=len(actionMap), random_play_steps=RANDOM_PLAY_STEPS, args=args, goal_feature=GOAL_FEATURE)

    total_rewards = []
    wrong_meta_pred = 0 
    total_steps = 0 
    create_video = False

    for episode in range(EPISODES):
        if args.verbose >= 1: 
            print(f'\n[Episode {episode}] Starting new episode...')

        # Initialize environment and get state
        env.restart()
        state = env.getStackedState()

        episode_steps = 0
        subgoal_agent_loss = 0
        meta_labels = []
        wrong_goal = False
        goal_idx = 0
        true_goal = goal_idx # if GOAL_MODE == 'full_sequence' else inv_label_mapping.get(SUBGOAL_ORDER[goal_idx])

        img_array = []
        if episode % 10000 == 0:
            create_video = True

        expert_goal = np.zeros((1, GOALS))
        expert_goal[0, true_goal] = 1.0
        meta_labels.append((state, expert_goal)) # append, but do not collect yet 

        if TRAINED_GOALS[true_goal]:
            goal = true_goal
            # Collect sample for metacontroller because goal will be reached
            metacontroller.collect(state, expert_goal) 
        else:
            # Metacontroller predict goal 
            goal = metacontroller.sample(metacontroller.predict(state))

        if goal!= true_goal:
            wrong_goal = True
            if args.verbose >= 1: 
                print(f"[Episode {episode}] Metacontroller predicted {goal} instead of {true_goal} as goal \U0001F47E ")
            wrong_meta_pred += 1
            if wrong_meta_pred % 100 == 0:
                print(f'[Episode {episode}] 100 wrong meta choices. Resetting metacontroller... ')
                metacontroller.reset() 

        else:
            if args.verbose >= 1: 
                print(f'[Episode {episode}] Metacontroller predicted goal {goal} as new goal...')

        all_step_times = []    
        step_time = time.time()

        while not env.isTerminal() and episode_steps < MAX_EPISODE_STEPS and not wrong_goal:

            # Unroll episode until agent reaches goal
            while not env.trueGoalReached(SUBGOAL_ORDER[goal_idx]) and not env.isTerminal() and episode_steps < MAX_EPISODE_STEPS and not wrong_goal: 
                
                goal_position = torch.tensor(np.array([ALL_SUBGOALS[SUBGOAL_ORDER[goal_idx]]]), device=device)

                if create_video:
                    img = env.getScreenOrig()
                    
                    
                    cv2.putText(img=img, text='goal: ' + goalExplain.get(SUBGOAL_ORDER[goal_idx]), org=(5, 205), fontFace=cv2.FONT_HERSHEY_SIMPLEX, 
                                fontScale=0.3, color=(255, 255, 255),thickness=1)

                    img_array.append(img)

                if GOAL_FEATURE:
                    action = agent.select_action(state.unsqueeze(0), goal_position) 
                else:
                    action = agent.select_action(state.unsqueeze(0)) 
                              
                external_reward = env.act(actionMap[action])
                if (external_reward != 0):
                    external_reward = 1.0

                all_step_times.append(time.time() - step_time)
                step_time = time.time()
                episode_steps += 1
                
                # Calculate intrinsic reward [optionally add distance reward]
                reward = 0.0 # - 1.0 / 1000       # small negative reward for each step 
                if env.trueGoalReached(SUBGOAL_ORDER[goal_idx]):
                    reward += 1.0
                if env.isTerminal() or episode_steps==MAX_EPISODE_STEPS:
                    reward -= 1.0

                """
                # Simple direction reward
                goal_vec = env.get_goal_direction(SUBGOAL_ORDER[goal_idx])
                action_vec = actionVectors[action]
                direction_reward = 0.001 * np.dot(action_vec, goal_vec)
                reward += direction_reward
                """

                # Distance Reward
                last_goal = SUBGOAL_ORDER[goal_idx-1] if goal_idx != 0 else None
                reward += env.distanceReward(lastGoal=last_goal, goal=SUBGOAL_ORDER[goal_idx])
                

                total_rewards.append(reward)
                reward = torch.tensor([reward], device=device)
                
                next_state = env.getStackedState()
                
                # Store transition and update network parameters
                agent.remember(state, action, reward, next_state, env.isTerminal(), goal_position)

                # Move to the next state
                state = next_state

                # Optimize the policy network
                agent_loss = agent.optimize_model()
                if agent_loss is not None:
                    subgoal_agent_loss += agent_loss

            # Update goal
            if episode_steps >= MAX_EPISODE_STEPS:
                subgoal_success_tracker[goal_idx].append(0)
                if args.verbose >= 1: 
                    print(f'[Episode {episode}] Reached maximum epsiode steps: {episode_steps}. Terminate episode.')
                break

            elif env.trueGoalReached(SUBGOAL_ORDER[goal_idx]):
                if args.verbose >= 1: 
                    print(f'[Episode {episode}] Goal reached! \U0001F389 after step #{episode_steps}. ')

                subgoal_success_tracker[goal_idx].append(1)
                subgoal_success_steps[goal_idx].append(episode_steps)
                
                # Predict new goal and continue if its true goal 
                goal_idx += 1

                """
                if subgoal_trailing_performance[goal_idx] == 0: 
                    # goal has not been seen before --> reset TanH schedule to enforce new exploration
                    print(episode, episode_steps, goal, 'resetting TanH schedule...')
                    agent.reset_epsilon_schedule()
                """

                if goal_idx == len(SUBGOAL_ORDER):
                    # Finished all options --> start new episode
                    wrong_goal = True
                    print(f"[Episode {episode}] Reached all goals! \U0001F38A \U0001F38A \U0001F38A")
                    break 

                true_goal = goal_idx # if GOAL_MODE == 'full_sequence' else inv_label_mapping.get(SUBGOAL_ORDER[goal_idx])
                expert_goal = np.zeros((1, GOALS))
                expert_goal[0, true_goal] = 1.0
                meta_labels.append((state, expert_goal)) # append, but do not collect yet 

                # Metacontroller predict goal 
                goal = metacontroller.sample(metacontroller.predict(state))

                if goal!= true_goal:
                    wrong_goal = True
                    if args.verbose >= 1: 
                        print(f"[Episode {episode}] Metacontroller predicted {goal} instead of {true_goal} as goal \U0001F47E ")
                    wrong_meta_pred += 1
                    if wrong_meta_pred % 100 == 0:
                        if args.verbose >= 1: 
                            print(f'[Episode {episode}] 100 wrong meta choices. Resetting metacontroller... ')
                        metacontroller.reset() 
                    break
                else:
                    # Continue with new goal
                    if args.verbose >= 1: 
                        print(f'[Episode {episode}] Metacontroller predicted goal {goal} as new goal...')

            else:
                subgoal_success_tracker[goal_idx].append(0)
                if args.verbose >= 1: 
                    print(f'[Episode {episode}] Agent killed after {episode_steps} steps \U0001F47E Terminate episode.')
                break

        if create_video and img_array:
            out_path = f'runs/video_run{TIME_STAMP}-E{episode}.gif'
            imageio.mimsave(out_path, img_array)
            print('Saved gif to ', out_path)
            create_video = False

        """END OF EPISODE"""
        # At end of episode: aggregate data for metacontroller 
        item = meta_labels[-1]
        metacontroller.collect(item[0], item[1]) 

        total_steps += episode_steps

        # if agent was running 
        if all_step_times:
            # Print average step time
            avg_step_time = sum(all_step_times) / len(all_step_times)
            
            # Log agent losses 
            agent_loss = subgoal_agent_loss / episode_steps

        else: 
            avg_step_time = 0 
            agent_loss = 0 

        
        if not DEBUG: 
            wandb.log({'mean_reward': np.array(total_rewards[-episode_steps:]).mean(), 'avg_step_time': avg_step_time*100, 'agent_loss': agent_loss, 'epsilon': agent.get_epsilon(), 'episode': episode})
            
        for g in range(GOALS):
            if len(subgoal_success_tracker[g]) < episode: 
                    # goal not evaluated yet because other goals not reached
                    subgoal_success_tracker[g].append(None)
            if not DEBUG:
                wandb.log({ f'sub-goal {g}': subgoal_success_tracker[g][-1], 'episode': episode})

        # Train metacontroller
        if metacontroller.check_training_clock():
            if args.verbose >= 2: 
                print('###################################')
                print('###   Training Metacontroller   ###')
                print('###################################')

            meta_loss, meta_acc = metacontroller.train()
            if not DEBUG:
                wandb.log({'meta_loss': meta_loss, 'meta_acc': meta_acc, 'episode': episode})

            if args.verbose >= 2: 
                print(f'Metacontroller loss: {meta_loss:.2f}\n')
                print('###################################')
        
        if len(subgoal_success_tracker[goal_idx]) > 100:
            subgoal_trailing_performance[goal] = sum([v for v in subgoal_success_tracker[goal][-100:] if not v is None])/ 100.0
            
            if args.verbose >= 0: 
                print(f'[Episode {episode}] Subgoal trailing performance for goal {goal_idx} is {subgoal_trailing_performance[goal_idx]:.2f}')
                print(f'[Episode {episode}] Agent {goal_idx} steps done {agent.steps_done}')
                print(f'[Episode {episode}] Agent epsilon value is {agent.get_epsilon():.2f}')
            
            if subgoal_trailing_performance[goal] > STOP_TRAINING_THRESHOLD:
                if args.verbose >= 1: 
                    print(f'[Episode {episode}] Training for goal #{goal} completed...')

                if true_goal == GOALS:
                    # Last goal reached trailing performance --> stop training
                    print(f"\nTraining completed \U0001F38A \U0001F38A \U0001F38A \n")
                    agent.finish_training(TIME_STAMP)
                    torch.save(metacontroller.model.state_dict(), f"run{TIME_STAMP}-metacontroller.pt")
                    break
    
    # After training is completed 
    for goal in range(GOALS):
        if subgoal_success_steps[goal]:
            print(f'\nAgent{goal} performance: {np.array(subgoal_success_tracker[goal]).mean():.2%}')
            print(f'Reached goal {goal} {sum(subgoal_success_tracker[goal])}x with {np.array(subgoal_success_steps[goal]).mean():.0f} average steps')


if __name__ == "__main__":
    main()
