import numpy as np
from gym import utils
from gym.envs.mujoco import mujoco_env
from os import path

class CMOEnv:
    def __init__(self, min_episode_cost, max_episode_cost, velocity_threshold):
        self.min_episode_cost, self.max_episode_cost = min_episode_cost, max_episode_cost
        self.target_cost, self.epsilon = self.max_episode_cost, 0 
        self.linear_preference = None
        self.velocity_threshold = velocity_threshold

    def set_linear_preference(self, linear_preference):
        if linear_preference is not None:
            self.linear_preference = np.array(linear_preference) / (np.linalg.norm(linear_preference, ord=1, axis=-1, keepdims=True)+1e-8)
        else:
            self.linear_preference = None

    def set_target_cost(self, target_cost):
        self.target_cost = target_cost
        self.epsilon = 1 if self.target_cost == 0 else 0

    def get_normalized_score(self, rew):
        rew = np.array(rew)
        rew[-1] = (rew[-1] + self.epsilon) / (self.target_cost + self.epsilon)
        return rew