import torch
import numpy as np

from utils.bbox_utils import create_bbox_geometry_objects_for_frame, build_per_vertex_color_map_for_world_scenario
from utils.camera.pinhole import PinholeCamera
from utils.graphics_utils import render_geometries
from utils.camera.base import CameraBase

from anydata.utils.colmap import qvec2rotmat, qvec2rotmat_batched
from anydata.geometry.camera import Camera


class AnyDataCamera(CameraBase):
    def __init__(self, K, hw, dtype=torch.float32, device=None):
        self.h = int(hw[0])
        self.w = int(hw[1])
        self.cam = Camera(K=K.unsqueeze(0), hw=(self.h, self.w))

        if device is None:
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.device = device
        self.dtype = dtype

    @property
    def width(self) -> int:
        return self.w

    @property
    def height(self) -> int:
        return self.h

    def ray2pixel_np(self, rays: np.ndarray) -> np.ndarray:
        """
        Args:
            rays: (M, 3), camera rays in camera coordinate (opencv convention)
        Returns:
            pixel_coords: (M, 2), pixel coordinates, not normalized to (0, 1)
        """
        rays = torch.tensor(rays).unsqueeze(0).permute(0, 2, 1).float()
        pixel_coords = self.cam.project_points(rays, normalize=False).squeeze(0)
        return pixel_coords


def render_nvidia_world_scenario(data, tc=None, has_batch_dim=False):

    if tc is None:
        return {key: render_nvidia_world_scenario(data, key, has_batch_dim=has_batch_dim) for key in data['bbox3d'].keys()}

    # Precompute bbox color map once per camera/clip
    bbox_per_vertex_color_map = build_per_vertex_color_map_for_world_scenario()

    bboxes3d = data['bbox3d'][tc]
    extrinsics = data['extrinsics'][tc]
    intrinsics = data['intrinsics'][tc]
    hw = data['rgb'][tc].shape[-2:]

    if has_batch_dim:
        intrinsics = intrinsics[0]
        extrinsics = extrinsics[0]

    bboxes3d = bboxes3d[bboxes3d[:, -3] > 0]

    hwl = bboxes3d[:, 0:3]
    lwh = hwl[:, ::-1]
    
    tvec = np.array(bboxes3d[:, 3:6])
    qvec = np.array(bboxes3d[:, 6:10])

    rotmat = np.transpose(qvec2rotmat_batched(qvec), [2, 0, 1])
    Tcw = np.repeat(np.expand_dims(np.eye(4), 0), rotmat.shape[0], axis=0)
    Tcw[:, :3, :3] = rotmat
    Tcw[:, :3, -1] = tvec

    camera_pose = extrinsics.numpy()
    object_to_world = camera_pose @ Tcw
    camera_model = AnyDataCamera(K=intrinsics, hw=hw)

    geometry_objects = []
    for i in range(bboxes3d.shape[0]):

        info = {0: {
            'object_type': 'Car',
            'object_to_world': object_to_world[i],
            'object_lwh': lwh[i],
        }}

        # bounding boxes
        geometry_objects.extend(
            create_bbox_geometry_objects_for_frame(
                info,
                camera_pose,
                camera_model,
                fill_face='all',
                fill_face_style='solid',
                object_type_to_per_vertex_color=bbox_per_vertex_color_map,
                line_width=4,
                edge_color=[200, 200, 200],
            )
        )

    # Render once for this frame
    combined_frame = render_geometries(
        geometry_objects,
        camera_model.height,
        camera_model.width,
        depth_max=200,
        depth_gradient=True,
    )

    return combined_frame