"""
Example of running PyTorch implementation of DDPG on HalfCheetah.
"""
import copy

import sys
import torch

from gym.envs.mujoco import HalfCheetahEnv

from torch.nn import functional as F

from rlkit.data_management.env_replay_buffer import EnvReplayBuffer
from rlkit.envs.wrappers import NormalizedBoxEnv
from rlkit.exploration_strategies.base import (
    PolicyWrappedWithExplorationStrategy
)
from rlkit.exploration_strategies.ou_strategy import OUStrategy
from rlkit.launchers.launcher_util import setup_logger
from rlkit.samplers.data_collector import MdpPathCollector
from rlkit.torch.networks import ConcatMlp, TanhMlpPolicy
from rlkit.torch.ddpg.ddpg import DDPGTrainer
import rlkit.torch.pytorch_util as ptu
from rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm


def experiment(variant):
    eval_env = NormalizedBoxEnv(HalfCheetahEnv())
    expl_env = NormalizedBoxEnv(HalfCheetahEnv())
    if variant["activation_type"].lower() == "gelu":
        hidden_activation = F.gelu
    else:
        hidden_activation = F.relu
    # Or for a specific version:
    # import gym
    # env = NormalizedBoxEnv(gym.make('HalfCheetah-v1'))
    obs_dim = eval_env.observation_space.low.size
    action_dim = eval_env.action_space.low.size
    qf = ConcatMlp(
        input_size=obs_dim + action_dim,
        output_size=1,
        hidden_activation=hidden_activation,
        **variant['qf_kwargs']
    )
    policy = TanhMlpPolicy(
        input_size=obs_dim,
        output_size=action_dim,
        hidden_activation=hidden_activation,
        **variant['policy_kwargs']
    )
    target_qf = copy.deepcopy(qf)
    target_policy = copy.deepcopy(policy)
    eval_path_collector = MdpPathCollector(eval_env, policy)
    exploration_policy = PolicyWrappedWithExplorationStrategy(
        exploration_strategy=OUStrategy(action_space=expl_env.action_space),
        policy=policy,
    )
    expl_path_collector = MdpPathCollector(expl_env, exploration_policy)
    replay_buffer = EnvReplayBuffer(variant['replay_buffer_size'], expl_env)
    trainer = DDPGTrainer(
        qf=qf,
        target_qf=target_qf,
        policy=policy,
        target_policy=target_policy,
        **variant['trainer_kwargs']
    )
    algorithm = TorchBatchRLAlgorithm(
        trainer=trainer,
        exploration_env=expl_env,
        evaluation_env=eval_env,
        exploration_data_collector=expl_path_collector,
        evaluation_data_collector=eval_path_collector,
        replay_buffer=replay_buffer,
        **variant['algorithm_kwargs']
    )
    algorithm.to(ptu.device)
    algorithm.train()


if __name__ == "__main__":
    # noinspection PyTypeChecker
    random_seed = int(sys.argv[1])
    arch_file_path = sys.argv[2]
    width_arch = []
    q_arch = []
    eval_env = NormalizedBoxEnv(HalfCheetahEnv())
    obs_dim = eval_env.observation_space.low.size
    action_dim = eval_env.action_space.low.size
    print(obs_dim, action_dim)
    action_dim = eval_env.action_space.low.size
    with open(arch_file_path, "r") as arch_file:
        for l in arch_file.readlines():
            l_arr = l.split(",")
            num_hidden = len(l_arr)
            rep_layer_size = int(num_hidden/2)
            for i, hidden_width in enumerate(l_arr):
                width_arch.append(int(hidden_width))
                if i == num_hidden:
                    q_arch.append(int(hidden_width) + action_dim)
                else:
                    q_arch.append(int(hidden_width))

    if len(sys.argv) > 3:
        activation_type = sys.argv[3]
    else:
        activation_type = "relu"

    variant = dict(
        algorithm_kwargs=dict(
            num_epochs=1000,
            num_eval_steps_per_epoch=1000,
            num_trains_per_train_loop=1000,
            num_expl_steps_per_train_loop=1000,
            min_num_steps_before_training=10000,
            max_path_length=1000,
            num_train_loops_per_epoch=1,
            batch_size=128,
        ),
        trainer_kwargs=dict(
            use_soft_update=True,
            tau=1e-2,
            discount=0.99,
            qf_learning_rate=1e-3,
            policy_learning_rate=1e-4,
        ),
        qf_kwargs=dict(
            hidden_sizes=width_arch,
        ),
        policy_kwargs=dict(
            hidden_sizes=width_arch,
        ),
        replay_buffer_size=int(1E6),
        activation_type=activation_type
    )
    ptu.set_gpu_mode(True)  # optionally set the GPU (default=False)
    setup_logger('name-of-experiment', variant=variant)
    torch.manual_seed(random_seed)
    experiment(variant)
