import torch
import numpy as np
import cv2
from PIL import Image
from scipy.spatial.transform import Rotation as R
import torchvision.transforms as T
from hydra import initialize, compose
import hydra

P = np.array(
    [[-1, 0, 0, 0], [0, 0, -1, 0], [0, -1, 0, 0], [0, 0, 0, 1]], dtype=np.float32
)

def init_model_loss_fn(cfg):
    device = cfg.device if torch.cuda.is_available() else "cpu"
    model = hydra.utils.instantiate(cfg.model).to(device)
    model_weight_pth = cfg.get("model_weight_pth")

    if model_weight_pth is None:
        raise ValueError("Model weight path is not specified in the config.")

    checkpoint = torch.load(
        model_weight_pth, map_location=device, weights_only=False
    )

    try:
        model.load_state_dict(checkpoint["model"])
    except RuntimeError:
        checkpoint["model"] = {
            k.replace("_orig_mod.", ""): v for k, v in checkpoint["model"].items()
        }
        model.load_state_dict(checkpoint["model"])
    loss_fn = hydra.utils.instantiate(cfg.loss_fn, model=model)
    loss_fn.load_state_dict(checkpoint["loss_fn"])
    loss_fn = loss_fn.to(device)

    return model, loss_fn

class VectorizedBuffer:
    def __init__(self, batch_size, buffer_size, act_dim, device):
        self.batch_size = batch_size
        self.buffer_size = buffer_size
        self.act_dim = act_dim
        self.device = device
        self.image_buffer = None
        self.goal_buffer = None
        self.action_buffer = None
        self.image_buffers_sizes = torch.zeros(batch_size, device=device)
        self.goal_buffer_size = torch.zeros(batch_size, device=device)
        self.action_buffers_sizes = torch.zeros(batch_size, device=device)

    def add_image(self, new_images):
        if self.image_buffer is None:
            self.image_buffer = (
                new_images.unsqueeze(1)
                .repeat(1, self.buffer_size, 1, 1, 1)
                .to(self.device)
            )
        else:
            for b in range(self.batch_size):
                if self.image_buffers_sizes[b] == 0:
                    self.image_buffer[b] = (
                        new_images[b].unsqueeze(0).repeat(self.buffer_size, 1, 1, 1)
                    )
                else:
                    self.image_buffer[b] = torch.roll(
                        self.image_buffer[b], shifts=-1, dims=0
                    )
                    self.image_buffer[b, -1] = new_images[b]
        self.image_buffers_sizes += 1
        self.image_buffers_sizes = torch.clamp(
            self.image_buffers_sizes, max=self.buffer_size
        )

    def add_goal(self, new_goals):
        if self.goal_buffer is None:
            self.goal_buffer = (
                new_goals.unsqueeze(1).repeat(1, self.buffer_size, 1).to(self.device)
            )
        else:
            for b in range(self.batch_size):
                if self.goal_buffer_size[b] == 0:
                    self.goal_buffer[b] = (
                        new_goals[b].unsqueeze(0).repeat(self.buffer_size, 1)
                    )
                else:
                    self.goal_buffer[b] = torch.roll(
                        self.goal_buffer[b], shifts=-1, dims=0
                    )
                    self.goal_buffer[b, -1] = new_goals[b]
        self.goal_buffer_size += 1
        self.goal_buffer_size = torch.clamp(self.goal_buffer_size, max=self.buffer_size)

    def reset(self, batch_indices):
        self.image_buffers_sizes[batch_indices] = 0
        self.goal_buffer_size[batch_indices] = 0
        self.action_buffers_sizes[batch_indices] = 0
        if self.image_buffer is not None:
            self.image_buffer[batch_indices] = torch.zeros_like(
                self.image_buffer[batch_indices]
            )
        if self.goal_buffer is not None:
            self.goal_buffer[batch_indices] = torch.zeros_like(
                self.goal_buffer[batch_indices]
            )
        if self.action_buffer is not None:
            self.action_buffer[batch_indices] = torch.zeros_like(
                self.action_buffer[batch_indices]
            )

    def add_action(self, new_actions):
        B = new_actions.shape[0]
        if self.action_buffer is None:
            self.action_buffer = torch.zeros(
                B, self.buffer_size - 1, self.act_dim, device=self.device
            )
            self.action_buffers_sizes = torch.zeros(B, device=self.device)
        for b in range(B):
            if self.action_buffers_sizes[b] != 0:
                self.action_buffer[b] = torch.roll(
                    self.action_buffer[b], shifts=-1, dims=0
                )
            self.action_buffer[b, -1] = new_actions[b]
        self.action_buffers_sizes += 1

    def get_input_sequence(self):
        B = self.image_buffer.shape[0]
        if self.action_buffer is None:
            action_buffer = torch.zeros(
                B, self.buffer_size - 1, self.act_dim, device=self.device
            )
        else:
            action_buffer = self.action_buffer
        base_act = torch.zeros(B, 1, self.act_dim, device=self.device)
        act_seq = torch.cat([action_buffer, base_act], dim=1)
        if self.goal_buffer is None:
            goal_seq = None
        else:
            goal_seq = torch.stack([goal for goal in self.goal_buffer]).to(
                dtype=torch.float32
            )
        return self.image_buffer, goal_seq, act_seq


def unwrap_model(model):
    if isinstance(
        model, (torch.nn.DataParallel, torch.nn.parallel.DistributedDataParallel)
    ):
        return model.module
    return model


class Policy:
    def __init__(self, model_path=None, device="cpu", model=None, loss_fn=None):
        if model_path is None and (model is None or loss_fn is None):
            raise ValueError(
                "Either model_path or both model and loss_fn must be provided."
            )

        if model is None or loss_fn is None:
            with initialize(config_path="configs", version_base=None):
                cfg = compose(config_name="run_vqbet")
            checkpoint = torch.load(model_path, map_location="cpu", weights_only=False)
            ckpt_cfg = checkpoint["cfg"]
            del checkpoint

            cfg.goal_dim = ckpt_cfg["loss_fn"]["goal_dim"]

            cfg.gpt_input_dim = ckpt_cfg["loss_fn"]["gpt_model"]["config"]["input_dim"]
            cfg.loss_fn.gpt_model.config.n_layer = ckpt_cfg["loss_fn"]["gpt_model"][
                "config"
            ]["n_layer"]
            cfg.loss_fn.gpt_model.config.n_head = ckpt_cfg["loss_fn"]["gpt_model"][
                "config"
            ]["n_head"]
            cfg.loss_fn.gpt_model.config.n_embd = ckpt_cfg["loss_fn"]["gpt_model"][
                "config"
            ]["n_embd"]

            cfg.vqvae_n_embed = ckpt_cfg["vqvae_n_embed"]
            cfg.model_weight_pth = model_path

            cfg.device = device
            model, loss_fn = init_model_loss_fn(cfg)

        self.to_tensor = T.ToTensor()
        self.model = unwrap_model(model)
        self.loss_fn = unwrap_model(loss_fn)
        self.buffer_size = self.loss_fn._vqbet.obs_window_size
        self.device = device

        goal_dim = self.loss_fn.goal_dim
        self.condition = f"{goal_dim}d"
        if goal_dim == 0:
            self.condition = None

        valid_conditions = ("4d", "3d", "2d")
        if self.condition is not None and self.condition not in valid_conditions:
            raise ValueError(
                f"'condition' must be one of {valid_conditions}, got '{self.condition}'"
            )

        self.model.eval()
        self.loss_fn.eval()

        self.vectorized_buffer = None
        self.act_dim = 7

        self.rot_yx_90 = (
            R.from_euler("y", 90, degrees=True).as_matrix()
            @ R.from_euler("x", 90, degrees=True).as_matrix()
        )
        self.Tyx = np.eye(4, dtype=np.float32)
        self.Tyx[:3, :3] = self.rot_yx_90

        rot_z_90 = R.from_euler("z", 90, degrees=True).as_matrix()
        Tz = np.eye(4, dtype=np.float32)
        Tz[:3, :3] = rot_z_90

        M = self.Tyx @ Tz @ P.T
        M_inv = M.T

        self.M_t = torch.from_numpy(M).to(device)
        self.M_inv_t = torch.from_numpy(M_inv).to(device)

    def reset(self, indicies=None):
        if indicies is None:
            self.vectorized_buffer = None
        else:
            self.vectorized_buffer.reset(indicies)

    def process_image(self, img):
        if isinstance(img, np.ndarray):
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = Image.fromarray(img)
        return self.to_tensor(img)
    
    def process_image_batch(self, imgs):
        if isinstance(imgs, np.ndarray):
            if imgs.dtype == np.uint8:
                imgs_tensor = torch.from_numpy(imgs).permute(0, 3, 1, 2).float() / 255.0
            else:
                imgs_tensor = torch.from_numpy(imgs).permute(0, 3, 1, 2).float()
            return imgs_tensor.to(self.device)
        else:
            return torch.stack([self.process_image(imgs[i]) for i in range(len(imgs))]).to(self.device)

    def infer(self, observations):
        obs = observations["rgb_ego"]
        if self.condition == "2d":
            goal = observations["object_2d_position"]
        elif self.condition == "3d":
            goal = observations["object_3d_position"]
        else:
            goal = None

        if len(obs.shape) == 3:
            obs = np.expand_dims(obs, axis=0)
            if goal is not None:
                goal = np.expand_dims(goal, axis=0)

        B = obs.shape[0]
        processed_images = self.process_image_batch(obs)

        if self.vectorized_buffer is None:
            image_shape = processed_images.shape[1:]
            self.vectorized_buffer = VectorizedBuffer(
                batch_size=B,
                buffer_size=self.buffer_size,
                act_dim=self.act_dim,
                device=self.device,
            )

        self.vectorized_buffer.add_image(processed_images)

        if self.condition is not None:
            processed_goals = torch.from_numpy(goal).to(self.device, dtype=torch.float32).view(B, -1)
            self.vectorized_buffer.add_goal(processed_goals)

        img_seq, goal_seq, act_seq = self.vectorized_buffer.get_input_sequence()

        with torch.no_grad():
            model_input = (img_seq, goal_seq, act_seq)
            model_output = self.model(model_input)
            action_tensors, logs = self.loss_fn.step(
                model_input, model_output, return_all=True
            )

        action_tensors = action_tensors.squeeze(1).to(self.device)
        self.vectorized_buffer.add_action(action_tensors)
        action_tensors = action_tensors.cpu().numpy()

        return action_tensors
