import gym

import forks.rlkit.rlkit.torch.pytorch_util as ptu
from forks.rlkit.rlkit.data_management.obs_dict_replay_buffer import ObsDictRelabelingBuffer
from forks.rlkit.rlkit.launchers.launcher_util import setup_logger
from forks.rlkit.rlkit.samplers.data_collector import GoalConditionedPathCollector
from forks.rlkit.rlkit.torch.her.her import HERTrainer
from forks.rlkit.rlkit.torch.networks import FlattenMlp
from forks.rlkit.rlkit.torch.sac.policies import MakeDeterministic, TanhGaussianPolicy
from forks.rlkit.rlkit.torch.sac.sac import SACTrainer
from forks.rlkit.rlkit.torch.torch_rl_algorithm import TorchBatchRLAlgorithm


def experiment(variant):
    eval_env = gym.make("FetchReach-v1")
    expl_env = gym.make("FetchReach-v1")

    observation_key = "observation"
    desired_goal_key = "desired_goal"

    achieved_goal_key = desired_goal_key.replace("desired", "achieved")
    replay_buffer = ObsDictRelabelingBuffer(
        env=eval_env,
        observation_key=observation_key,
        desired_goal_key=desired_goal_key,
        achieved_goal_key=achieved_goal_key,
        **variant["replay_buffer_kwargs"]
    )
    obs_dim = eval_env.observation_space.spaces["observation"].low.size
    action_dim = eval_env.action_space.low.size
    goal_dim = eval_env.observation_space.spaces["desired_goal"].low.size
    qf1 = FlattenMlp(
        input_size=obs_dim + action_dim + goal_dim,
        output_size=1,
        **variant["qf_kwargs"]
    )
    qf2 = FlattenMlp(
        input_size=obs_dim + action_dim + goal_dim,
        output_size=1,
        **variant["qf_kwargs"]
    )
    target_qf1 = FlattenMlp(
        input_size=obs_dim + action_dim + goal_dim,
        output_size=1,
        **variant["qf_kwargs"]
    )
    target_qf2 = FlattenMlp(
        input_size=obs_dim + action_dim + goal_dim,
        output_size=1,
        **variant["qf_kwargs"]
    )
    policy = TanhGaussianPolicy(
        obs_dim=obs_dim + goal_dim, action_dim=action_dim, **variant["policy_kwargs"]
    )
    eval_policy = MakeDeterministic(policy)
    trainer = SACTrainer(
        env=eval_env,
        policy=policy,
        qf1=qf1,
        qf2=qf2,
        target_qf1=target_qf1,
        target_qf2=target_qf2,
        **variant["sac_trainer_kwargs"]
    )
    trainer = HERTrainer(trainer)
    eval_path_collector = GoalConditionedPathCollector(
        eval_env,
        eval_policy,
        observation_key=observation_key,
        desired_goal_key=desired_goal_key,
    )
    expl_path_collector = GoalConditionedPathCollector(
        expl_env,
        policy,
        observation_key=observation_key,
        desired_goal_key=desired_goal_key,
    )
    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["algo_kwargs"]
    )
    algorithm.to(ptu.device)
    algorithm.train()


if __name__ == "__main__":
    variant = dict(
        algorithm="HER-SAC",
        version="normal",
        algo_kwargs=dict(
            batch_size=128,
            num_epochs=100,
            num_eval_steps_per_epoch=5000,
            num_expl_steps_per_train_loop=1000,
            num_trains_per_train_loop=1000,
            min_num_steps_before_training=1000,
            max_path_length=50,
        ),
        sac_trainer_kwargs=dict(
            discount=0.99,
            soft_target_tau=5e-3,
            target_update_period=1,
            policy_lr=3e-4,
            qf_lr=3e-4,
            reward_scale=1,
            use_automatic_entropy_tuning=True,
        ),
        replay_buffer_kwargs=dict(
            max_size=int(1e6),
            fraction_goals_rollout_goals=0.2,  # equal to k = 4 in HER paper
            fraction_goals_env_goals=0,
        ),
        qf_kwargs=dict(hidden_sizes=[400, 300]),
        policy_kwargs=dict(hidden_sizes=[400, 300]),
    )
    setup_logger("her-sac-fetch-experiment", variant=variant)
    experiment(variant)
