from __future__ import division
from __future__ import absolute_import
from rllab2.envs.base import Step
from rllab2.misc.overrides import overrides
from .mujoco_env import MujocoEnv
import numpy as np
from rllab2.core.serializable import Serializable
from rllab2.misc import logger
from rllab2.misc import autoargs


class SwimmerEnv(MujocoEnv, Serializable):

    FILE = u'swimmer.xml'
    ORI_IND = 2

    @autoargs.arg(u'ctrl_cost_coeff', type=float,
                  help=u'cost coefficient for controls')
    def __init__(
            self,
            ctrl_cost_coeff=1e-2,
            *args, **kwargs):
        self.ctrl_cost_coeff = ctrl_cost_coeff
        super(SwimmerEnv, self).__init__(*args, **kwargs)
        Serializable.quick_init(self, locals())

    def get_current_obs(self):
        return np.concatenate([
            self.model.data.qpos.flat,
            self.model.data.qvel.flat,
            self.get_body_com(u"torso").flat,
        ]).reshape(-1)

    def get_ori(self):
        return self.model.data.qpos[self.__class__.ORI_IND]

    def step(self, action):
        self.forward_dynamics(action)
        next_obs = self.get_current_obs()
        lb, ub = self.action_bounds
        scaling = (ub - lb) * 0.5
        ctrl_cost = 0.5 * self.ctrl_cost_coeff * np.sum(
            np.square(action / scaling))
        forward_reward = self.get_body_comvel(u"torso")[0]
        reward = forward_reward - ctrl_cost
        done = False
        return Step(next_obs, reward, done)

    @overrides
    def log_diagnostics(self, paths):
        if len(paths) > 0:
            progs = [
                path[u"observations"][-1][-3] - path[u"observations"][0][-3]
                for path in paths
            ]
            logger.record_tabular(u'AverageForwardProgress', np.mean(progs))
            logger.record_tabular(u'MaxForwardProgress', np.max(progs))
            logger.record_tabular(u'MinForwardProgress', np.min(progs))
            logger.record_tabular(u'StdForwardProgress', np.std(progs))
        else:
            logger.record_tabular(u'AverageForwardProgress', np.nan)
            logger.record_tabular(u'MaxForwardProgress', np.nan)
            logger.record_tabular(u'MinForwardProgress', np.nan)
            logger.record_tabular(u'StdForwardProgress', np.nan)
