# from roboimi.envs.vs_rl_task.task.vs_base_task import VisualServo_BaseTask
from .vs_base_task import VisualServo_BaseTask
import numpy as np
import os
import torch
from torch import Tensor
from typing import Tuple, Dict

class VisualServoDualArmRobot(VisualServo_BaseTask):
    def __init__(self, cfg, device, headless):
        """ 
        Args:
            cfg (Dict): Task Environment config file
            device (string): 'cuda:0' or 'cpu'
            headless (bool): Run without rendering if True
        """
        self.cfg = cfg
        self.height_samples = None
        self.debug_viz = False
        self.init_done = False
        self._parse_cfg(self.cfg)
        self.mj_dt = 0.001          #attention
        super().__init__(self.cfg, device, headless)
        self._init_buffers()
        self._prepare_reward_function()
        self.is_timeout = False
        self.is_reset = True
        self.reset_buf = torch.zeros(self.num_envs)
        # self.init_done = True



    def reset(self):
        if  not self.is_reset:
            return
        else:
            '''
            env_ids = []
            '''
            self._resample_commands()
            super().reset()
            #only one env for now!
            self.last_actions[0] = 0.
            self.episode_length_buf[0] = 0.
            self.is_timeout = False
            self.reset_buf[0] = True # 0 /1 
            self.extras["episode"] = {}
            for key in self.episode_sums.keys():
                self.extras["episode"]['rew_' + key] = torch.mean(self.episode_sums[key][0]) / self.max_episode_length_s
                self.episode_sums[key][0] = 0.
            if self.cfg.env.send_timeouts:
                self.extras["time_outs"] = self.is_timeout
           
        
            


    def check_termination(self):
        if self.common_step_counter >= self.max_episode_length:
            self.is_timeout = True
        else: #self.is_timeout = False
            self.is_timeout = False
            
        if self.is_detect and (not self.is_timeout):
            self.is_reset = False
        else:
            self.is_reset = True
        self.reset_buf[0] = self.is_reset

        

    def step(self,actions):

        clip_actions = self.cfg.normalization.clip_actions
        self.actions = torch.clip(actions, -clip_actions, clip_actions).to(self.device)
        #---------------------------------------------------------------------------------
        actions_scaled = actions * self.cfg.control.action_scale #adjust action upper/lower limits   
        actions_adj = actions_scaled*actions  
        super().step(actions_adj)
        # super().step(actions)  
        self.episode_length_buf += 1
        self.common_step_counter += 1
        self.check_termination()
        self.compute_reward()
        self.reset()
        self.compute_observations()
        self.last_actions[:] = self.actions[:]
        #---------------------------------------------------------------------------------
        clip_obs = self.cfg.normalization.clip_observations
        self.obs_buf = torch.clip(self.obs_buf, -clip_obs, clip_obs)
        if self.privileged_obs_buf is not None:
            self.privileged_obs_buf = torch.clip(self.privileged_obs_buf, -clip_obs, clip_obs)

        self.last_actions[:] = self.actions[:]
        return self.obs_buf, self.privileged_obs_buf, self.rew_buf, self.reset_buf, self.extras

    def _parse_cfg(self, cfg):
        self.dt = self.cfg.control.decimation * self.mj_dt
        self.obs_scales = self.cfg.normalization.obs_scales
        self.reward_scales = self.class_to_dict(self.cfg.rewards.scales)
        self.command_ranges = self.class_to_dict(self.cfg.commands.ranges)
        self.max_episode_length_s = self.cfg.env.episode_length_s
        self.max_episode_length = np.ceil(self.max_episode_length_s / self.dt)

       

    def _init_buffers(self):
        """ Initialize torch tensors which will contain simulation states and processed quantities
        """
         # get gym GPU state tensors

        # actor_root_state = self.gym.acquire_actor_root_state_tensor(self.sim)
        # dof_state_tensor = self.gym.acquire_dof_state_tensor(self.sim)
        # net_contact_forces = self.gym.acquire_net_contact_force_tensor(self.sim)
        # self.gym.refresh_dof_state_tensor(self.sim)
        # self.gym.refresh_actor_root_state_tensor(self.sim)
        # self.gym.refresh_net_contact_force_tensor(self.sim)

        # initialize some data used later on
        self.common_step_counter = 0
        self.extras = {}
        self.noise_scale_vec = self._get_noise_scale_vec(self.cfg)
        self.actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False)
        self.last_actions = torch.zeros(self.num_envs, self.num_actions, dtype=torch.float, device=self.device, requires_grad=False)
        

    def _prepare_reward_function(self):
        """ Prepares a list of reward functions, whcih will be called to compute the total reward.
            Looks for self._reward_<REWARD_NAME>, where <REWARD_NAME> are names of all non zero reward scales in the cfg.
        """
        # remove zero scales + multiply non-zero ones by dt
        for key in list(self.reward_scales.keys()):
            scale = self.reward_scales[key]
            if scale==0:
                self.reward_scales.pop(key) 
            else:
                self.reward_scales[key] *= self.dt
        # prepare list of functions
        self.reward_functions = []
        self.reward_names = []
        for name, scale in self.reward_scales.items():
            if name=="termination":
                continue
            self.reward_names.append(name)
            name = '_reward_' + name
            self.reward_functions.append(getattr(self, name))

        # reward episode sums
        self.episode_sums = {name: torch.zeros(self.num_envs, dtype=torch.float, device=self.device, requires_grad=False)
                             for name in self.reward_scales.keys()}


    # def _get_noise_scale_vec(self, cfg):
    #     """ Sets a vector used to scale the noise added to the observations.
    #         [NOTE]: Must be adapted when changing the observations structure

    #     Args:
    #         cfg (Dict): Environment config file

    #     Returns:
    #         [torch.Tensor]: Vector of scales used to multiply a uniform distribution in [-1, 1]
    #     """
    #     noise_vec = torch.zeros_like(self.obs_buf[0])
    #     self.add_noise = self.cfg.noise.add_noise
    #     noise_scales = self.cfg.noise.noise_scales
    #     noise_level = self.cfg.noise.noise_level
    #     noise_vec[:3] = noise_scales.lin_vel * noise_level * self.obs_scales.lin_vel
    #     noise_vec[3:6] = noise_scales.ang_vel * noise_level * self.obs_scales.ang_vel
    #     noise_vec[6:9] = noise_scales.gravity * noise_level
    #     noise_vec[9:12] = 0. # commands
    #     noise_vec[12:24] = noise_scales.dof_pos * noise_level * self.obs_scales.dof_pos
    #     noise_vec[24:36] = noise_scales.dof_vel * noise_level * self.obs_scales.dof_vel
    #     noise_vec[36:48] = 0. # previous actions
    #     if self.cfg.terrain.measure_heights:
    #         noise_vec[48:235] = noise_scales.height_measurements* noise_level * self.obs_scales.height_measurements
    #     return noise_vec
    

    def _resample_commands(self):
        """ Randommly select commands of some environments

        Args:
            env_ids (List[int]): Environments ids for which new commands are needed
        """
        pass


    def compute_reward(self):
        """ Compute rewards
            
        """
        self.rew_buf[:] = 0.
        for i in range(len(self.reward_functions)):
            name = self.reward_names[i]
            rew = self.reward_functions[i]() * self.reward_scales[name]
            self.rew_buf += rew
            self.episode_sums[name] += rew
        if self.cfg.rewards.only_positive_rewards:
            self.rew_buf[:] = torch.clip(self.rew_buf[:], min=0.)
        # add termination reward after clipping
        # if "termination" in self.reward_scales:
        #     rew = self._reward_termination() * self.reward_scales["termination"]
        #     self.rew_buf += rew
        #     self.episode_sums["termination"] += rew
    
    def compute_observations(self):
        """ Computes observations
        """
        # self.obs_buf = torch.cat((  self.base_lin_vel * self.obs_scales.lin_vel,
        #                             self.base_ang_vel  * self.obs_scales.ang_vel,
        #                             self.projected_gravity,
        #                             self.commands[:, :3] * self.commands_scale,
        #                             (self.dof_pos - self.default_dof_pos) * self.obs_scales.dof_pos,
        #                             self.dof_vel * self.obs_scales.dof_vel,
        #                             self.actions
        #                             ),dim=-1)
        self.obs_buf = torch.zeros(1,10)
        # add perceptive inputs if not blind
        # add noise if needed
        # if self.add_noise:
        #     self.obs_buf += (2 * torch.rand_like(self.obs_buf) - 1) * self.noise_scale_vec


    def class_to_dict(self,obj) -> dict:
        if not  hasattr(obj,"__dict__"):
            return obj
        result = {}
        for key in dir(obj):
            if key.startswith("_"):
                continue
            element = []
            val = getattr(obj, key)
            if isinstance(val, list):
                for item in val:
                    element.append(self.class_to_dict(item))
            else:
                element = self.class_to_dict(val)
            result[key] = element
        return result


    # cfg 中必须和_reward_ 后面的部分写的一样
    def _reward_action_rate(self):
        # Penalize changes in actions
        return torch.sum(torch.square(self.last_actions - self.actions), dim=1)
    
    def _reward_lin_vel_xyz(self):
        return torch.zeros(1,1)
        # return torch.square(self.ee_lin_vel[:, 3])