import numpy as np
import torch
import sys
from roboimi.assets.robots.bidiana_med_vs import BiDianaMedVS
from roboimi.envs.visual_servo.double_vs_base import DualVisualServoEnv

#base_task
class VisualServo_BaseTask(DualVisualServoEnv):
    def __init__(self, cfg, device, headless):
        self.headless = headless
        super().__init__(
            robot=BiDianaMedVS(),
            is_render=self.headless,
            renderer="viewer",
            control_freq=200
        )
        if device == 'cpu' or device == 'cuda:0':
            self.device = device
        self.num_envs = 1
        self.num_obs = cfg.env.num_observations
        self.num_privileged_obs = cfg.env.num_privileged_obs
        self.num_actions = cfg.env.num_actions
        torch._C._jit_set_profiling_mode(False)
        torch._C._jit_set_profiling_executor(False)
        self.obs_buf = torch.zeros(self.num_envs, self.num_obs, device=self.device, dtype=torch.float)
        self.rew_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.float)
        self.reset_buf = torch.ones(self.num_envs, device=self.device, dtype=torch.long)
        self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long)
        self.time_out_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool)
        if self.num_privileged_obs is not None:
            self.privileged_obs_buf = torch.zeros(self.num_envs, self.num_privileged_obs, device=self.device, dtype=torch.float)
        else: 
            self.privileged_obs_buf = None
            # self.num_privileged_obs = self.num_obs

        # self.extras = {}
        
    def get_observations(self):
        return self.obs_buf
    
    def get_privileged_observations(self):
        return self.privileged_obs_buf

    def reset_idx(self, env_ids):
        """Reset selected robots"""
        pass

    def reset(self):
        """ Reset all robots"""
        self.reset_idx(torch.arange(self.num_envs, device=self.device))
        obs, privileged_obs, _, _, _ = self.step(torch.zeros(self.num_envs, self.num_actions, device=self.device, requires_grad=False))
        return obs, privileged_obs

    def step(self, actions):
        print("PPO actions =",actions)
        actions = 0.5
        super().step(actions)

  