#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

import torch
from torch import nn
import numpy as np
from d3gs.utils.graphics_utils import getWorld2View2, getProjectionMatrix
import time

# class Camera_CPU(nn.Module):
#     def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask,
#                  image_name, uid,
#                  trans=np.array([0.0, 0.0, 0.0]), scale=1.0, data_device="cpu"
#                  ):
#         super(Camera_CPU, self).__init__()
        
#         data_device="cpu"
        
#         self.uid = uid
#         self.colmap_id = colmap_id
#         self.R = R
#         self.T = T
#         self.FoVx = FoVx
#         self.FoVy = FoVy
#         self.image_name = image_name
#         self.data_device = torch.device(data_device)

#         # 保存在 CPU 中，稍后由 to_device 迁移
#         self.original_image = image.clamp(0.0, 1.0)
#         self.gt_alpha_mask = gt_alpha_mask
#         self.trans = trans
#         self.scale = scale
#         self.zfar = 100.0
#         self.znear = 0.01

#         self.image_width = image.shape[2]
#         self.image_height = image.shape[1]

#         self.world_view_transform = None
#         self.projection_matrix = None
#         self.full_proj_transform = None
#         self.camera_center = None

#     def to_device(self, device="cuda"):
#         # 图像及其 mask
#         self.original_image = self.original_image.to(device)
#         if self.gt_alpha_mask is not None:
#             self.original_image *= self.gt_alpha_mask.to(device)
#         else:
#             self.original_image *= torch.ones(
#                 (1, self.image_height, self.image_width), device=device)

#         # 矩阵处理
#         self.world_view_transform = torch.tensor(
#             getWorld2View2(self.R, self.T, self.trans, self.scale), device=device
#         ).transpose(0, 1)

#         self.projection_matrix = getProjectionMatrix(
#             znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy
#         ).transpose(0, 1).to(self.data_device)

#         self.full_proj_transform = self.world_view_transform.unsqueeze(0).bmm(
#             self.projection_matrix.unsqueeze(0)
#         ).squeeze(0)

#         self.camera_center = self.world_view_transform.inverse()[3, :3]


class Camera(nn.Module):
    def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask,
                 image_name, uid,
                 trans=np.array([0.0, 0.0, 0.0]), scale=1.0, data_device = "cuda"
                 ):
        super(Camera, self).__init__()

        self.uid = uid
        self.colmap_id = colmap_id
        self.R = R
        self.T = T
        self.FoVx = FoVx
        self.FoVy = FoVy
        self.image_name = image_name
        # 在这一步很可能涉及到cuda了
        try:
            self.data_device = torch.device(data_device)
        except Exception as e:
            print(e)
            print(f"[Warning] Custom device {data_device} failed, fallback to default cuda device" )
            self.data_device = torch.device("cuda")

        # 在cuda中
        self.original_image = image.clamp(0.0, 1.0).to(self.data_device)  # cuda
        self.image_width = self.original_image.shape[2]
        self.image_height = self.original_image.shape[1]

        if gt_alpha_mask is not None:
            self.original_image *= gt_alpha_mask.to(self.data_device)
        else:
            self.original_image *= torch.ones((1, self.image_height, self.image_width), device=self.data_device)

        self.zfar = 100.0
        self.znear = 0.01

        self.trans = trans
        self.scale = scale
        
        # 以下四个变量都在cuda中
        self.world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose(0, 1).cuda()  # cuda
        self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).cuda()  # cuda
        self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0)  # cuda
        self.camera_center = self.world_view_transform.inverse()[3, :3]  # cuda


class PhysCamera(Camera):
    def __init__(self, colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask,
                 image_name, uid, view: str, step: int,
                 trans=np.array([0.0, 0.0, 0.0]), scale=1.0, data_device = "cuda"
                 ):
        
        super(PhysCamera, self).__init__(colmap_id, R, T, FoVx, FoVy, image, gt_alpha_mask,
                                         image_name, uid, trans, scale, data_device)
        # view id
        self.view = view
        # step id
        self.step = step


class MiniCam:
    def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform):
        self.image_width = width
        self.image_height = height    
        self.FoVy = fovy
        self.FoVx = fovx
        self.znear = znear
        self.zfar = zfar
        self.world_view_transform = world_view_transform
        self.full_proj_transform = full_proj_transform
        view_inv = torch.inverse(self.world_view_transform)
        self.camera_center = view_inv[3][:3]
