import numpy as np
import random

from .half_cheetah import HalfCheetahEnv


class HalfCheetahDirEnv(HalfCheetahEnv):
    """Half-cheetah environment with target direction, as described in [1]. The
    code is adapted from
    https://github.com/cbfinn/maml_rl/blob/9c8e2ebd741cb0c7b8bf2d040c4caeeb8e06cc95/rllab/envs/mujoco/half_cheetah_env_rand_direc.py

    The half-cheetah follows the dynamics from MuJoCo [2], and receives at each
    time step a reward composed of a control cost and a reward equal to its
    velocity in the target direction. The tasks are generated by sampling the
    target directions from a Bernoulli distribution on {-1, 1} with parameter
    0.5 (-1: backward, +1: forward).

    [1] Chelsea Finn, Pieter Abbeel, Sergey Levine, "Model-Agnostic
        Meta-Learning for Fast Adaptation of Deep Networks", 2017
        (https://arxiv.org/abs/1703.03400)
    [2] Emanuel Todorov, Tom Erez, Yuval Tassa, "MuJoCo: A physics engine for
        model-based control", 2012
        (https://homes.cs.washington.edu/~todorov/papers/TodorovIROS12.pdf)
    """

    def __init__(self, task={}, n_tasks=2, max_episode_steps=200, **kwargs):
        self._task = task
        self.num_tasks = n_tasks
        self.tasks = self.sample_tasks(n_tasks)
        # self.set_task(self.sample_tasks(1)[0])
        self._goal = np.float64(self.tasks[0].get('goal', 0))
        self._max_episode_steps = max_episode_steps
        super(HalfCheetahDirEnv, self).__init__()

    def step(self, action):
        xposbefore = self.sim.data.qpos[0]
        self.do_simulation(action, self.frame_skip)
        xposafter = self.sim.data.qpos[0]

        forward_vel = (xposafter - xposbefore) / self.dt
        # forward_reward = self.goal_direction * forward_vel
        forward_reward = self._goal * forward_vel
        ctrl_cost = 0.5 * 1e-1 * np.sum(np.square(action))

        observation = self._get_obs()
        reward = forward_reward - ctrl_cost
        done = False
        infos = dict(reward_forward=forward_reward,
                     reward_ctrl=-ctrl_cost,
                     task=self.get_task())
        return observation, reward, done, infos

    def sample_step(self, obs, action, replace_obs = True, reservse_obs=True):
        if replace_obs:
            old_state = self.sim.get_state()
            old_qpos = old_state[1]
            old_qvel = old_state[2]
            obs = obs.cpu().numpy()
            qpos = np.concatenate([np.array([0.]), obs[0][:8]])
            qvel = obs[0][8:17]
            self.set_state(qpos, qvel)

        xposbefore = self.sim.data.qpos[0]
        self.do_simulation(action, self.frame_skip)
        xposafter = self.sim.data.qpos[0]

        forward_vel = (xposafter - xposbefore) / self.dt
        forward_reward = self._goal * forward_vel
        ctrl_cost = 0.5 * 1e-1 * np.sum(np.square(action))

        observation = self._get_obs()
        reward = forward_reward - ctrl_cost
        if reservse_obs:
            self.set_state(old_qpos, old_qvel)
        
        return observation, reward

    def sample_tasks(self, n_tasks):
        # for fwd/bwd env, goal direc is backwards if - 1.0, forwards if + 1.0
        return [{'goal': random.choice([-1.0, 1.0])} for _ in range(n_tasks,)]

    def set_task(self, task):
        self._goal = np.float64(task)
        # self.goal_direction = task
        
    def set_goal(self, goal):
        self._goal = np.float64(goal)

    def get_task(self):
        # return self.goal_direction
        return self._goal

    def set_all_goals(self, goals):
        assert self.num_tasks == len(goals)
        self.tasks = [{'goal': dir[0]} for dir in goals]
        self.reset_task(0)

    def get_all_task_idx(self):
        return range(len(self.tasks))

    def reset_task(self, idx=None):
        if idx is not None:
            # idx = self.sample_tasks(1)[0]
            self._task = self.tasks[idx]
            self._goal = self._task['goal']
        # self.set_task(idx)
        self.reset()

    def reward(self, state, action):
        ' Here, state is previous state! r_t = r(s_{t-1}, a_t) '
        qpos = np.concatenate([np.array([0.]), state[:8]])
        qvel = state[8:17]
        # set state to replay it (in hindsight)
        self.set_state(qpos, qvel)

        xposbefore = self.sim.data.qpos[0]
        self.do_simulation(action, self.frame_skip)
        xposafter = self.sim.data.qpos[0]

        forward_vel = (xposafter - xposbefore) / self.dt
        forward_reward = self._goal * forward_vel
        ctrl_cost = 0.5 * 1e-1 * np.sum(np.square(action))
        reward = forward_reward - ctrl_cost
        return reward


class HalfCheetahRandDirOracleEnv(HalfCheetahDirEnv):

    def _get_obs(self):
        return np.concatenate([
            self.sim.data.qpos.flat[1:],
            self.sim.data.qvel.flat,
            self.get_body_com("torso").flat,
            [self.goal_direction]
        ]).astype(np.float32).flatten()
