import math
import torch
from dust3r.utils.geometry import inv
from src.utils.cuda_splatting import DummyCamera

def get_scaled_camera(ref_camera_extrinsics, target_camera_extrinsics, target_camera_intrinsics, scale, image_shape):
    """
    get a scaled camera from a reference camera to a target camera
    
    """
    
    # get extrinsics(target_camera to ref_camera)
    target_camera_extrinsics = inv(ref_camera_extrinsics) @ target_camera_extrinsics
    # scale translation
    target_camera_extrinsics[:3, 3] = target_camera_extrinsics[:3, 3] * scale
    # invert extrinsics(ref_camera to target_camera)
    target_camera_extrinsics_inv = inv(target_camera_extrinsics)
    # calculate fov
    fovx = 2 * math.atan(image_shape[1] / (2 * target_camera_intrinsics[0, 0]))
    fovy = 2 * math.atan(image_shape[0] / (2 * target_camera_intrinsics[1, 1]))
    # return camera(numpy)
    R = target_camera_extrinsics_inv[:3, :3].cpu().numpy().transpose() # R.transpose() : ref_camera_2_target_camera
    T = target_camera_extrinsics_inv[:3, 3].cpu().numpy() # T : ref_camera_2_target_camera
    image_shape = image_shape.cpu().numpy()
    return DummyCamera(R, T, fovx, fovy, image_shape[1], image_shape[0])

def move_c2w_along_z(extrinsics: torch.Tensor, distance: float) -> torch.Tensor:
    """
    向后移动多个 Camera-to-World (C2W) 矩阵，使相机沿各自 Z 轴方向远离原点。

    参数:
        extrinsics (torch.Tensor): 形状为 [N, 4, 4] 的张量，包含 N 个 C2W 矩阵。
        distance (float): 向后移动的距离。

    返回:
        torch.Tensor: 更新后的 C2W 矩阵，形状与输入相同。
    """
    # 确保输入是一个四维矩阵，且最后一维是 4x4
    assert extrinsics.dim() == 3 and extrinsics.shape[1:] == (4, 4), \
        "输入的 extrinsics 必须是形状为 [N, 4, 4] 的张量"

    # 创建一个拷贝以免修改原矩阵
    updated_extrinsics = extrinsics.clone()

    # 遍历每个 C2W 矩阵
    for i in range(updated_extrinsics.shape[0]):
        # 提取旋转矩阵 R 和平移向量 t
        R = updated_extrinsics[i, :3, :3]  # 形状为 [3, 3]
        t = updated_extrinsics[i, :3, 3]   # 形状为 [3]

        # 获取相机的 Z 轴方向（第三列）
        z_axis = R[:, 2]  # 形状为 [3]

        # 计算新的平移向量，沿 Z 轴方向向后移动
        t_new = t - distance * z_axis

        # 更新 C2W 矩阵的平移部分
        updated_extrinsics[i, :3, 3] = t_new

    return updated_extrinsics