import sys
sys.path.append('../../')

import DDPG
import OurDDPG
import TD3
import utils
import numpy as np
import torch
import argparse
import os
import random

sys.path.append('../')
from common import make_env, get_frame_skip_and_timestep, perform_action, random_jitter_force, random_disturb
from evals import *

default_timestep = 0.02
default_frame_skip = 2


# Main function of the policy. Model is trained and evaluated inside.
def train(policy='TD3', seed=0, start_timesteps=25e3, eval_freq=5e3, max_timesteps=1e5,
          expl_noise=0.1, batch_size=256, discount=0.99, tau=0.005, policy_freq=2, policy_noise=2, noise_clip=0.5,
          save_model=False, load_model="", jit_duration=0.02, g_ratio=1, response_rate=0.08, std_eval=False,
          catastrophe_frequency=1, reflex_response_rate=0.02, reflex_threshold=0.15, reflex_force_scale=1.0):

    delayed_env = True
    max_force = g_ratio * 9.81
    env_name = 'InvertedPendulum-v2'
    eval_policy = eval_policy_std if std_eval else eval_policy_ori
    arguments = [policy, 'reflex', env_name, seed, jit_duration, g_ratio, response_rate, catastrophe_frequency,
                 reflex_response_rate, reflex_threshold, reflex_force_scale]
    file_name = '_'.join([str(x) for x in arguments])
    print("---------------------------------------")
    print(f"Policy: {policy}, Env: {env_name}, Seed: {seed}")
    print("---------------------------------------")

    if not os.path.exists("./results"):
        os.makedirs("./results")

    if save_model and not os.path.exists("./models"):
        os.makedirs("./models")

    frame_skip, timestep, jit_frames = get_frame_skip_and_timestep(jit_duration, response_rate, reflex_response_rate)
    reflex_frames = int(reflex_response_rate/timestep)

    print('timestep:', timestep)  # How long does it take before two frames
    # How many frames to skip before return the state, 1 by default
    print('frameskip:', frame_skip)

    # The ratio of the default time consumption between two states returned and reset version.
    # Used to reset the max episode number to guarantee the actual max time is always the same.
    time_change_factor = (default_timestep * default_frame_skip) / (timestep * frame_skip)
    env = make_env(env_name, seed, time_change_factor, timestep, frame_skip, delayed_env)

    print('time change factor', time_change_factor)
    # Set seeds
    torch.manual_seed(seed)
    np.random.seed(seed)

    max_timesteps = max_timesteps * time_change_factor
    eval_freq = int(eval_freq * time_change_factor)
    start_timesteps = start_timesteps * time_change_factor

    state_dim = env.observation_space[0].shape[0]
    action_dim = env.action_space.shape[0]
    max_action = float(env.action_space.high[0])

    kwargs = {
        "state_dim": state_dim,
        "action_dim": action_dim,
        "max_action": max_action,
        "discount": discount,
        "tau": tau,
        "observation_space": env.observation_space,
        "delayed_env": delayed_env,
        "reflex": True,
        "threshold": reflex_threshold,
        "reflex_force_scale": reflex_force_scale
    }

    # Initialize policy
    if policy == "TD3":
        # Target policy smoothing is scaled wrt the action scale
        kwargs["policy_noise"] = policy_noise * max_action
        kwargs["noise_clip"] = noise_clip * max_action
        kwargs["policy_freq"] = policy_freq
        policy = TD3.TD3(**kwargs)
    elif policy == "OurDDPG":
        policy = OurDDPG.DDPG(**kwargs)
    elif policy == "DDPG":
        policy = DDPG.DDPG(**kwargs)

    if load_model != "":
        policy_file = file_name if load_model == "default" else load_model
        policy.load(f"./models/{policy_file}")

    replay_buffer = utils.ReplayBuffer(state_dim+action_dim, action_dim)

    # Evaluate untrained policy
    evaluations = [eval_policy(policy, env_name, eval_episodes=10, time_change_factor=time_change_factor,
                               jit_duration=jit_duration, env_timestep=timestep, max_force=max_force, frame_skip=frame_skip,
                               jit_frames=jit_frames, delayed_env=delayed_env, reflex_frames=reflex_frames)]

    state, done = env.reset(), False
    episode_reward = 0
    episode_timesteps = 0
    episode_num = 0
    jittered_frames = 0
    disturb = round(random.randint(50, 100) * 0.04 * (1/catastrophe_frequency), 3)

    jittering = False
    jitter_force = 0
    best_performance = 0
    counter = 0
    for t in range(int(max_timesteps)):
        episode_timesteps += 1

        # Select action randomly or according to policy
        if t < start_timesteps:
            reflex, action = policy.select_action(state)
        else:
            reflex, a = policy.select_action(state)
            action = (a + np.random.normal(0, max_action * expl_noise, size=action_dim)).clip(-max_action, max_action)

        # Perform action
        jittering, disturb, counter, jittered_frames, jitter_force, max_force, next_state, reward, done = perform_action(jittering, disturb, counter, response_rate, env, reflex, action, reflex_frames, frame_skip, random_jitter_force, max_force, timestep, jit_frames, jittered_frames, random_disturb, jitter_force, catastrophe_frequency, delayed_env)

        done_bool = float(done) if episode_timesteps < env.env.env._max_episode_steps else 0

        # Store data in replay buffer
        replay_buffer.add(state, action, next_state, reward, done_bool)

        state = next_state
        episode_reward += reward
        counter = round(counter, 3)

        # Train agent after collecting sufficient data
        if t >= start_timesteps:
            policy.train(replay_buffer, batch_size)

        if done:
            # +1 to account for 0 indexing. +0 on ep_timesteps since it will increment +1 even if done=True
            print(
                f"Total T: {t+1} Episode Num: {episode_num+1} Episode T: {episode_timesteps} Reward: {episode_reward:.3f}")
            # Reset environment
            state, done = env.reset(), False
            episode_reward = 0
            episode_timesteps = 0
            episode_num += 1
            counter = 0
            env.model.opt.gravity[0] = 0
            disturb = round(random.randint(50, 100) * 0.04 * (1 / catastrophe_frequency), 3)
            jittering = False
            jittered_frames = 0

        # Evaluate episode
        if (t + 1) % eval_freq == 0:
            avg_reward = eval_policy(policy, env_name, eval_episodes=10, time_change_factor=time_change_factor,
                                     jit_duration=jit_duration, env_timestep=timestep, max_force=max_force,
                                     frame_skip=frame_skip, jit_frames=jit_frames, delayed_env=delayed_env,
                                     reflex_frames=reflex_frames)
            evaluations.append(avg_reward)
            np.save(f"./results/{file_name}", evaluations)
            if best_performance < avg_reward:
                best_performance = avg_reward
                if save_model:
                    policy.save(f"./models/{file_name}_best")

        if counter == disturb:  # Execute adding jitter horizontal force here
            jitter_force, _ = random_jitter_force(max_force)  # Jitter force strength w/ direction
            env.model.opt.gravity[0] = jitter_force
            jittering = True
            jittered_frames = 0

    if save_model:
        policy.save(f"./models/{file_name}_final")
        # if t >= 25000:
        #     env.render()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--policy", default="TD3", help="Policy name (TD3, DDPG or OurDDPG)")
    parser.add_argument("--seed", default=0, type=int, help="Sets Gym, PyTorch and Numpy seeds")
    parser.add_argument("--start_timesteps", default=25e3, type=int, help="Time steps initial random policy is used")
    parser.add_argument("--eval_freq", default=10000, type=int, help="How often (time steps) we evaluate")
    parser.add_argument("--max_timesteps", default=4e5, type=int, help="Max time steps to run environment")
    parser.add_argument("--expl_noise", default=0.1, help="Std of Gaussian exploration noise")
    parser.add_argument("--batch_size", default=256, type=int, help="Batch size for both actor and critic")
    parser.add_argument("--discount", default=0.99, help="Discount factor")
    parser.add_argument("--tau", default=0.005, help="Target network update rate")
    parser.add_argument("--policy_noise", default=0.2, help="Noise added to target policy during critic update")
    parser.add_argument("--noise_clip", default=0.5, help="Range to clip target policy noise")
    parser.add_argument("--policy_freq", default=2, type=int, help="Frequency of delayed policy updates")
    parser.add_argument("--save_model", action="store_true", help="Save model and optimizer parameters")
    parser.add_argument("--load_model", default="", help="Model load file name, `` doesn't load, `default` uses file_name")
    parser.add_argument("--jit_duration", default=0.02, type=float, help="Duration in seconds for the horizontal force")
    parser.add_argument("--g_ratio", default=0, type=float, help='Maximum horizontal force g ratio')
    parser.add_argument("--response_rate", default=0.04, type=float, help="Response time of the agent in seconds")
    parser.add_argument("--std_eval", action="store_true", help="Use standard evaluation or original evaluation policy")
    parser.add_argument("--catastrophe_frequency", default=1.0, type=float, help="Modify how often to apply catastrophe")
    parser.add_argument("--reflex_response_rate", default=0.02, type=float, help="reflex Response time of the agent in seconds")
    parser.add_argument("--reflex_threshold", default=0.15, type=float, help="Threhsold at which the reflex activates")
    parser.add_argument("--reflex_force_scale", default=1.0, type=float, help="Multiply reflex with this value")


    args = parser.parse_args()
    args = vars(args)
    print()
    print('Command-line argument values:')
    for key, value in args.items():
        print('-', key, ':', value)

    print()

    train(**args)

import sys
sys.path.append('../../')

import DDPG
import OurDDPG
import TD3
import utils
import numpy as np
import torch
import argparse
import os
import random

sys.path.append('../')
from common import make_env, get_frame_skip_and_timestep, perform_action, random_jitter_force, random_disturb
from evals import *

default_timestep = 0.02
default_frame_skip = 2


# Main function of the policy. Model is trained and evaluated inside.
def train(policy='TD3', seed=0, start_timesteps=25e3, eval_freq=5e3, max_timesteps=1e5,
          expl_noise=0.1, batch_size=256, discount=0.99, tau=0.005, policy_freq=2, policy_noise=2, noise_clip=0.5,
          save_model=False, load_model="", jit_duration=0.02, g_ratio=1, response_rate=0.08, std_eval=False,
          catastrophe_frequency=1, reflex_response_rate=0.02, reflex_threshold=0.15, reflex_force_scale=1.0):

    delayed_env = True
    max_force = g_ratio * 9.81
    env_name = 'InvertedPendulum-v2'
    eval_policy = eval_policy_std if std_eval else eval_policy_ori
    arguments = [policy, 'reflex', env_name, seed, jit_duration, g_ratio, response_rate, catastrophe_frequency,
                 reflex_response_rate, reflex_threshold, reflex_force_scale]
    file_name = '_'.join([str(x) for x in arguments])
    print("---------------------------------------")
    print(f"Policy: {policy}, Env: {env_name}, Seed: {seed}")
    print("---------------------------------------")

    if not os.path.exists("./results"):
        os.makedirs("./results")

    if save_model and not os.path.exists("./models"):
        os.makedirs("./models")

    frame_skip, timestep, jit_frames = get_frame_skip_and_timestep(jit_duration, response_rate, reflex_response_rate)
    reflex_frames = int(reflex_response_rate/timestep)

    print('timestep:', timestep)  # How long does it take before two frames
    # How many frames to skip before return the state, 1 by default
    print('frameskip:', frame_skip)

    # The ratio of the default time consumption between two states returned and reset version.
    # Used to reset the max episode number to guarantee the actual max time is always the same.
    time_change_factor = (default_timestep * default_frame_skip) / (timestep * frame_skip)
    env = make_env(env_name, seed, time_change_factor, timestep, frame_skip, delayed_env)

    print('time change factor', time_change_factor)
    # Set seeds
    torch.manual_seed(seed)
    np.random.seed(seed)

    max_timesteps = max_timesteps * time_change_factor
    eval_freq = int(eval_freq * time_change_factor)
    start_timesteps = start_timesteps * time_change_factor

    state_dim = env.observation_space[0].shape[0]
    action_dim = env.action_space.shape[0]
    max_action = float(env.action_space.high[0])

    kwargs = {
        "state_dim": state_dim,
        "action_dim": action_dim,
        "max_action": max_action,
        "discount": discount,
        "tau": tau,
        "observation_space": env.observation_space,
        "delayed_env": delayed_env,
        "reflex": True,
        "threshold": reflex_threshold,
        "reflex_force_scale": reflex_force_scale
    }

    # Initialize policy
    if policy == "TD3":
        # Target policy smoothing is scaled wrt the action scale
        kwargs["policy_noise"] = policy_noise * max_action
        kwargs["noise_clip"] = noise_clip * max_action
        kwargs["policy_freq"] = policy_freq
        policy = TD3.TD3(**kwargs)
    elif policy == "OurDDPG":
        policy = OurDDPG.DDPG(**kwargs)
    elif policy == "DDPG":
        policy = DDPG.DDPG(**kwargs)

    if load_model != "":
        policy_file = file_name if load_model == "default" else load_model
        policy.load(f"./models/{policy_file}")

    replay_buffer = utils.ReplayBuffer(state_dim+action_dim, action_dim)

    # Evaluate untrained policy
    evaluations = [eval_policy(policy, env_name, eval_episodes=10, time_change_factor=time_change_factor,
                               jit_duration=jit_duration, env_timestep=timestep, max_force=max_force, frame_skip=frame_skip,
                               jit_frames=jit_frames, delayed_env=delayed_env, reflex_frames=reflex_frames)]

    state, done = env.reset(), False
    episode_reward = 0
    episode_timesteps = 0
    episode_num = 0
    jittered_frames = 0
    disturb = round(random.randint(50, 100) * 0.04 * (1/catastrophe_frequency), 3)

    jittering = False
    jitter_force = 0
    best_performance = 0
    counter = 0
    for t in range(int(max_timesteps)):
        episode_timesteps += 1

        # Select action randomly or according to policy
        if t < start_timesteps:
            reflex, action = policy.select_action(state)
        else:
            reflex, a = policy.select_action(state)
            action = (a + np.random.normal(0, max_action * expl_noise, size=action_dim)).clip(-max_action, max_action)

        # Perform action
        jittering, disturb, counter, jittered_frames, jitter_force, max_force, next_state, reward, done = perform_action(jittering, disturb, counter, response_rate, env, reflex, action, reflex_frames, frame_skip, random_jitter_force, max_force, timestep, jit_frames, jittered_frames, random_disturb, jitter_force, catastrophe_frequency, delayed_env)

        done_bool = float(done) if episode_timesteps < env.env.env._max_episode_steps else 0

        # Store data in replay buffer
        replay_buffer.add(state, action, next_state, reward, done_bool)

        state = next_state
        episode_reward += reward
        counter = round(counter, 3)

        # Train agent after collecting sufficient data
        if t >= start_timesteps:
            policy.train(replay_buffer, batch_size)

        if done:
            # +1 to account for 0 indexing. +0 on ep_timesteps since it will increment +1 even if done=True
            print(
                f"Total T: {t+1} Episode Num: {episode_num+1} Episode T: {episode_timesteps} Reward: {episode_reward:.3f}")
            # Reset environment
            state, done = env.reset(), False
            episode_reward = 0
            episode_timesteps = 0
            episode_num += 1
            counter = 0
            env.model.opt.gravity[0] = 0
            disturb = round(random.randint(50, 100) * 0.04 * (1 / catastrophe_frequency), 3)
            jittering = False
            jittered_frames = 0

        # Evaluate episode
        if (t + 1) % eval_freq == 0:
            avg_reward = eval_policy(policy, env_name, eval_episodes=10, time_change_factor=time_change_factor,
                                     jit_duration=jit_duration, env_timestep=timestep, max_force=max_force,
                                     frame_skip=frame_skip, jit_frames=jit_frames, delayed_env=delayed_env,
                                     reflex_frames=reflex_frames)
            evaluations.append(avg_reward)
            np.save(f"./results/{file_name}", evaluations)
            if best_performance < avg_reward:
                best_performance = avg_reward
                if save_model:
                    policy.save(f"./models/{file_name}_best")

        if counter == disturb:  # Execute adding jitter horizontal force here
            jitter_force, _ = random_jitter_force(max_force)  # Jitter force strength w/ direction
            env.model.opt.gravity[0] = jitter_force
            jittering = True
            jittered_frames = 0

    if save_model:
        policy.save(f"./models/{file_name}_final")
        # if t >= 25000:
        #     env.render()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--policy", default="TD3", help="Policy name (TD3, DDPG or OurDDPG)")
    parser.add_argument("--seed", default=0, type=int, help="Sets Gym, PyTorch and Numpy seeds")
    parser.add_argument("--start_timesteps", default=25e3, type=int, help="Time steps initial random policy is used")
    parser.add_argument("--eval_freq", default=10000, type=int, help="How often (time steps) we evaluate")
    parser.add_argument("--max_timesteps", default=4e5, type=int, help="Max time steps to run environment")
    parser.add_argument("--expl_noise", default=0.1, help="Std of Gaussian exploration noise")
    parser.add_argument("--batch_size", default=256, type=int, help="Batch size for both actor and critic")
    parser.add_argument("--discount", default=0.99, help="Discount factor")
    parser.add_argument("--tau", default=0.005, help="Target network update rate")
    parser.add_argument("--policy_noise", default=0.2, help="Noise added to target policy during critic update")
    parser.add_argument("--noise_clip", default=0.5, help="Range to clip target policy noise")
    parser.add_argument("--policy_freq", default=2, type=int, help="Frequency of delayed policy updates")
    parser.add_argument("--save_model", action="store_true", help="Save model and optimizer parameters")
    parser.add_argument("--load_model", default="", help="Model load file name, `` doesn't load, `default` uses file_name")
    parser.add_argument("--jit_duration", default=0.02, type=float, help="Duration in seconds for the horizontal force")
    parser.add_argument("--g_ratio", default=0, type=float, help='Maximum horizontal force g ratio')
    parser.add_argument("--response_rate", default=0.04, type=float, help="Response time of the agent in seconds")
    parser.add_argument("--std_eval", action="store_true", help="Use standard evaluation or original evaluation policy")
    parser.add_argument("--catastrophe_frequency", default=1.0, type=float, help="Modify how often to apply catastrophe")
    parser.add_argument("--reflex_response_rate", default=0.02, type=float, help="reflex Response time of the agent in seconds")
    parser.add_argument("--reflex_threshold", default=0.15, type=float, help="Threhsold at which the reflex activates")
    parser.add_argument("--reflex_force_scale", default=1.0, type=float, help="Multiply reflex with this value")


    args = parser.parse_args()
    args = vars(args)
    print()
    print('Command-line argument values:')
    for key, value in args.items():
        print('-', key, ':', value)

    print()

    train(**args)

# Main function of the policy. Model is trained and evaluated inside. # How long does it take before two frames # How many frames to skip before return the state, 1 by default # The ratio of the default time consumption between two states returned and reset version. # Used to reset the max episode number to guarantee the actual max time is always the same. # Set seeds # Initialize policy # Target policy smoothing is scaled wrt the action scale # Evaluate untrained policy # Select action randomly or according to policy # Perform action # Store data in replay buffer # Train agent after collecting sufficient data # +1 to account for 0 indexing. +0 on ep_timesteps since it will increment +1 even if done=True # Reset environment # Evaluate episode # Execute adding jitter horizontal force here # Jitter force strength w/ direction # if t >= 25000: #     env.render()