import cv2
import torch
import torch.nn.functional as F
import numpy as np
from scipy.spatial.transform import Rotation as R
from egogym.policies.base_policy import BasePolicy
from baselines.rum.policy import Policy

class RUMPolicy(BasePolicy):
    
    def __init__(self, config=None, policy=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 = "rum"
        self.device = config.get("device", "cuda" if torch.cuda.is_available() else "cpu")
        if policy is not None:
            self.policy = policy
        else:
            self.policy = Policy(model_path=config["model_path"], device=self.device)
        self.gripper_threshold = config["gripper_threshold"]
        self.grasping_style = config["grasping_style"]
        self.num_envs = None
        self.grasped = None
        self.desired_axis_t = torch.tensor([
            [-1,  0,  0],
            [ 0,  0, -1],
            [ 0, -1,  0]
        ], dtype=torch.float32, device=self.device)
        self.grasp_override_t = torch.tensor([0.00, 0.18, 0.04], dtype=torch.float32, device=self.device)

    def get_action(self, obs):

        if self.num_envs is None:
            if len(obs["rgb_ego"].shape) == 4:
                self.num_envs = obs["rgb_ego"].shape[0]
            else:
                self.num_envs = 1
            self.grasped = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)
        
        if self.num_envs == 1:
            obs = {key: obs[key][np.newaxis, ...] if isinstance(obs[key], np.ndarray) else obs[key] for key in obs}

        T_world_camera = torch.from_numpy(obs["camera_pose"]).view(self.num_envs, 4, 4).to(self.device)
        if "object_pose" in obs:
            T_world_object = torch.from_numpy(obs["object_pose"]).view(self.num_envs, 4, 4).to(self.device)
        else:
            T_world_object = torch.from_numpy(obs["handle_pose"]).view(self.num_envs, 4, 4).to(self.device)
        
        R_world_camera = T_world_camera[:, :3, :3]
        t_world_camera = T_world_camera[:, :3, 3]
        p_world = T_world_object[:, :3, 3]
        R_camera_world = R_world_camera.transpose(1, 2)
        p_camera = torch.bmm(R_camera_world, (p_world - t_world_camera).unsqueeze(-1)).squeeze(-1)
        object_3d_position = p_camera @ self.desired_axis_t.T
        object_3d_position[self.grasped] = self.grasp_override_t
        
        rgb_tensor = torch.from_numpy(obs["rgb_ego"]).permute(0, 3, 1, 2).to(self.device).float()
        rgb_resized = F.interpolate(rgb_tensor, size=(224, 224), mode='bilinear', align_corners=False)
        rgb_ego = rgb_resized.permute(0, 2, 3, 1).byte().cpu().numpy()

        rum_obs = {
            "rgb_ego": rgb_ego,
            "object_3d_position": object_3d_position.cpu().numpy(),
        }
        action_tensors = self.policy.infer(rum_obs)
        
        action_tensors_t = torch.from_numpy(action_tensors).to(self.device)
        R_action = R.from_euler("xyz", action_tensors[:, 3:6]).as_matrix()
        R_action_t = torch.from_numpy(R_action).to(self.device, dtype=torch.float32)
        
        action_pose = torch.eye(4, device=self.device).unsqueeze(0).repeat(self.num_envs, 1, 1)
        action_pose[:, :3, :3] = R_action_t
        action_pose[:, :3, 3] = action_tensors_t[:, :3]
        action_pose[:, :3, :3] = self.desired_axis_t.T @ action_pose[:, :3, :3] @ self.desired_axis_t
        action_pose[:, :3, 3] = action_pose[:, :3, 3] @ self.desired_axis_t
        
        self.grasped = torch.maximum(self.grasped, action_tensors_t[:, 6] < self.gripper_threshold)
        grasp_action = self.grasped.unsqueeze(1).float()
        if self.grasping_style == "continuous":
            grasp_action = 1 - action_tensors_t[:, 6].unsqueeze(1)

        actions = torch.cat([action_pose.view(self.num_envs, -1), grasp_action], dim=1)

        return actions.cpu().numpy()
    
    def reset(self, indicies=None):
        self.policy.reset(indicies)
        if indicies is not None and self.num_envs is not None:
            self.grasped[indicies] = False
        elif self.num_envs is not None:
            self.grasped = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device)


def get_config():
    config = {
        "model_path": "checkpoints/rum_pick.pt",
        "gripper_threshold": 0.7,
        "grasping_style": "binary",
        "device": None,
    }
    return config
