import numpy as np
import time
from egogym.policies.base_policy import BasePolicy
from egogym.utils import resize_with_pad

class PIPolicy(BasePolicy):
    
    def __init__(self, config=None, **kwargs):
        super().__init__()
        if config is None:
            config = get_config()
        
        for key, value in kwargs.items():
            config[key] = value
            
        self.config = config
        self.name = "pi"
        self.host = self.config["host"]
        self.port = self.config["port"]
        self.grasping_style = self.config["grasping_style"]
        self.buffer_length = self.config["buffer_length"]
        self.batch_size = 0

        try:
            from openpi_client import websocket_client_policy
        except ImportError:
            raise ImportError("Please install the openpi-client.")
        
        max_retries = 5
        for attempt in range(max_retries):
            try:
                self.model = websocket_client_policy.WebsocketClientPolicy(
                    host=self.host,
                    port=self.port,
                )
                print(f"Connected to PI model at {self.host}:{self.port}")
                break
            except Exception as e:
                if attempt < max_retries - 1:
                    print(f"Connection attempt {attempt + 1} failed: {e}. Retrying...")
                    time.sleep(1)
                else:
                    print(f"Failed to connect to remote model after {max_retries} attempts")
                    raise
                
    def get_action(self, obs):

        if self.batch_size == 0:
            if obs["rgb_exo"].ndim == 3:
                self.batch_size = 1
            else:
                self.batch_size = obs["rgb_exo"].shape[0]
            self.action_buffer = [None] * self.batch_size
            self.current_buffer_index = [0] * self.batch_size

        if self.batch_size == 1:
            obs = {key: np.expand_dims(value, axis=0) if isinstance(value, np.ndarray) else [value] for key, value in obs.items()}

        
        actions = []
        for i in range(self.batch_size):

            if self.action_buffer[i] is None or self.current_buffer_index[i] >= self.buffer_length:
                model_input = {
                    "observation/exterior_image_1_left": resize_with_pad(obs["rgb_exo"][i], 224, 224),
                    "observation/wrist_image_left": resize_with_pad(obs["rgb_ego"][i], 224, 224),
                    "observation/joint_position": np.array(obs["joint_positions"][i][:7]).reshape(7,),
                    "observation/gripper_position": np.array(obs["grasp"][i]).reshape(1,),
                    "prompt": f'pick up the {obs["object_name"][i]} and place it on a plate.',
                }
                self.action_buffer[i] = self.model.infer(model_input)["actions"]  
                self.current_buffer_index[i] = 0    
            current_action = self.action_buffer[i][self.current_buffer_index[i]]
            self.current_buffer_index[i] += 1

            gripper_pos = np.array([np.clip(current_action[7], 0.0, 1.0)])

            if self.grasping_style == "binary":
                gripper_pos = np.array([1.0]) if gripper_pos >= self.config["gripper_threshold"] else np.array([0.0])
            elif self.grasping_style == "semi_binary":
                gripper_pos = gripper_pos if gripper_pos <= self.config["gripper_threshold"] else np.array([1.0])
            elif not self.grasping_style == "continuous":
                raise ValueError(f"Invalid grasping style: {self.grasping_style}")
                
                
            arm_output = current_action[:7].reshape(7)
            actions.append(np.concatenate([arm_output, gripper_pos], axis=0).reshape(1, -1))

        return np.array(actions).reshape(self.batch_size, 8)
    
    def reset(self, indicies=None):
        if indicies is not None:
            for i in indicies:
                self.action_buffer[i] = None
                self.current_buffer_index[i] = 0
        else:
            self.batch_size = 0
            self.action_buffer = None
            self.current_buffer_index = None


def get_config():
    config = {
        "host": "localhost",
        "port": 8000,
        "gripper_threshold": 0.5,
        "grasping_style": "binary",
        "buffer_length": 8,
    }
    return config
