import torch
import math
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
from scene.gaussian_model import GaussianModel
from utils.sh_utils import eval_sh
from utils.rigid_utils import from_homogenous, to_homogenous, R_from_quaternions, quaternion_to_axis_angle, rigid_transform
import torch.nn.functional as F

def quaternion_to_rotation_matrix(q, get_R4 = False):
    """
    将四元数转换为 4x4 旋转矩阵
    :param q: 四元数 (w, x, y, z) -> Tensor of shape (N, 4)
    :return: 4x4 旋转矩阵 Tensor of shape (N, 4, 4)
    """
    # 提取四元数分量
    w = q[:, 0]
    x = q[:, 1]
    y = q[:, 2]
    z = q[:, 3]

    # 计算旋转矩阵的分量 (3x3)
    R = torch.stack([
        1 - 2 * (y ** 2 + z ** 2), 2 * (x * y - w * z), 2 * (x * z + w * y),
        2 * (x * y + w * z), 1 - 2 * (x ** 2 + z ** 2), 2 * (y * z - w * x),
        2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x ** 2 + y ** 2)
    ], dim=-1).reshape(-1, 3, 3)




    # 扩展为 4x4 旋转矩阵
    R_4 = torch.cat([
        torch.cat([R, torch.zeros(R.shape[0], 3, 1).cuda()], dim=-1),  # 添加第三列（0, 0, 0）
        torch.tensor([0.0, 0.0, 0.0, 1.0]).reshape(1, 4).repeat(R.shape[0], 1).unsqueeze(1).cuda()  # 最后一行 (0, 0, 0, 1)
    ], dim=-2)
    if get_R4:
        return R_4
    else:
        return R


def update_rotation_matrix(current_rotation_matrix, delta_quaternion):
    """
    使用四元数增量更新当前的旋转矩阵
    :param current_rotation_matrix: 当前旋转矩阵 (3, 3)
    :param delta_quaternion: 四元数增量 (w, x, y, z) -> Tensor of shape (4,)
    :return: 更新后的旋转矩阵 (3, 3)
    """
    # 将四元数增量转换为旋转矩阵
    delta_rotation_matrix = quaternion_to_rotation_matrix(delta_quaternion)

    # 更新旋转矩阵
    updated_rotation_matrix = torch.mm(delta_rotation_matrix, current_rotation_matrix)
    return updated_rotation_matrix

def quaternion_multiply(q1, q2):
    w1, x1, y1, z1 = q1[..., 0], q1[..., 1], q1[..., 2], q1[..., 3]
    w2, x2, y2, z2 = q2[..., 0], q2[..., 1], q2[..., 2], q2[..., 3]

    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2

    return torch.stack((w, x, y, z), dim=-1)

def new_render(viewpoint_camera, pc: GaussianModel, pipe, bg_color: torch.Tensor, fid= None, scaling_modifier=1.0, override_color=None):
    """
    Render the scene.

    Background tensor (bg_color) must be on GPU!
    """

    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    screenspace_points_densify = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()
        screenspace_points_densify.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug,
    )

    cov3D_precomp = None
    if pipe.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)
    else:
        scales = pc.get_scaling
        rotations = pc.get_rotation


    rasterizer = GaussianRasterizer(raster_settings=raster_settings)
    if fid is not None:
        if pc.revolute:
            #means3D = pc.get_xyz
            #dynamic_means3D = rigid_transform(pc, fid)
            #means3D[pc.dynamic_part_mask] = dynamic_means3D
            means3D = rigid_transform(pc, fid)
            if pc.use_canonical:
                scaling = (pc.canonical - fid) / pc.canonical
            else:
                scaling = fid
            axis, angle = quaternion_to_axis_angle(pc.quaternions)
            #scaling = (pc.canonical - fid) / pc.canonical
            angle = scaling * angle
            q_rot = torch.tensor([[1.0, 0., 0., 0.]]).repeat(pc.get_rotation.shape[0], 1).cuda()
            qw = torch.cos(angle / 2)
            qx = axis[0] * torch.sin(angle / 2)
            qy = axis[1] * torch.sin(angle / 2)
            qz = axis[2] * torch.sin(angle / 2)
            q_rot[pc.dynamic_part_mask] = torch.tensor([qw, qx, qy, qz]).cuda()  # 四元数顺序为 (w, x, y, z)
            rotations = quaternion_multiply(q_rot, pc.get_rotation)

        elif pc.prismatic:
            # d_xyz = torch.zeros_like(pc.get_xyz)
            # d_xyz[~pc.dynamic_part_mask] = pc.dir * fid
            # means3D = pc.get_xyz + d_xyz
            if pc.use_canonical:
                scaling = (fid - pc.canonical) / pc.canonical
            else:
                scaling = fid
            #scaling = (pc.canonical - fid) / pc.canonical
            d_xyz = torch.zeros_like(pc.get_xyz)
            #d_xyz[pc.dynamic_part_mask] = scaling * pc.dir
            unit_d = F.normalize(pc.dir, p=2, dim=0)
            d_xyz[pc.dynamic_part_mask] = scaling * unit_d * pc.dist
            #d_xyz[pc.dynamic_part_mask] = unit_d * pc.dist * fid
            #d_xyz[pc.dynamic_part_mask] = unit_d * pc.dist.unsqueeze(dim=-1) * fid
            means3D = pc.get_xyz + d_xyz
        else:
            means3D = pc.get_xyz
    else:
        means3D = pc.get_xyz


    opacity = pc.get_opacity

    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
    # scaling / rotation by the rasterizer.





    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
    shs = None
    colors_precomp = None
    if colors_precomp is None:
        if pipe.convert_SHs_python:
            shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree + 1) ** 2)
            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))
            dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
        else:
            shs = pc.get_features
    else:
        colors_precomp = override_color

    # Rasterize visible Gaussians to image, obtain their radii (on screen).
    rendered_image, radii, depth = rasterizer(
        means3D=means3D,
        means2D=screenspace_points,
        means2D_densify=screenspace_points_densify,
        shs=shs,
        colors_precomp=colors_precomp,
        opacities=opacity,
        scales=scales,
        rotations=rotations,
        cov3D_precomp=cov3D_precomp)

    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    # They will be excluded from value updates used in the splitting criteria.
    return {"render": rendered_image,
            "viewspace_points": screenspace_points,
            "viewspace_points_densify": screenspace_points_densify,
            #"visibility_filter": radii > 0,
            "radii": radii,
            "depth": depth}


def vanilla_render(viewpoint_camera, pc: GaussianModel, pipe, bg_color: torch.Tensor, d_xyz, d_rotation, d_scaling, is_6dof=False,
           scaling_modifier=1.0, override_color=None):
    """
    Render the scene.

    Background tensor (bg_color) must be on GPU!
    """

    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    screenspace_points_densify = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()
        screenspace_points_densify.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug,
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    if is_6dof:
        if torch.is_tensor(d_xyz) is False:
            means3D = pc.get_xyz
        else:
            means3D = from_homogenous(
                torch.bmm(d_xyz, to_homogenous(pc.get_xyz).unsqueeze(-1)).squeeze(-1))
    else:
        means3D = pc.get_xyz + d_xyz
    opacity = pc.get_opacity

    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
    # scaling / rotation by the rasterizer.
    scales = None
    rotations = None
    cov3D_precomp = None
    if pipe.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)
    else:
        scales = pc.get_scaling + d_scaling
        rotations = pc.get_rotation + d_rotation

    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
    shs = None
    colors_precomp = None
    if colors_precomp is None:
        if pipe.convert_SHs_python:
            shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree + 1) ** 2)
            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))
            dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
        else:
            shs = pc.get_features
    else:
        colors_precomp = override_color

    # Rasterize visible Gaussians to image, obtain their radii (on screen).
    rendered_image, radii, depth = rasterizer(
        means3D=means3D,
        means2D=screenspace_points,
        means2D_densify=screenspace_points_densify,
        shs=shs,
        colors_precomp=colors_precomp,
        opacities=opacity,
        scales=scales,
        rotations=rotations,
        cov3D_precomp=cov3D_precomp)

    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    # They will be excluded from value updates used in the splitting criteria.
    return {"render": rendered_image,
            "viewspace_points": screenspace_points,
            "viewspace_points_densify": screenspace_points_densify,
            "visibility_filter": radii > 0,
            "radii": radii,
            "depth": depth}



def mini_render(viewpoint_camera, pc: GaussianModel, pipe, bg_color: torch.Tensor, d_xyz, d_rotation, d_scaling, is_6dof=False,
           scaling_modifier=1.0, override_color=None):
    """
    Render the scene.

    Background tensor (bg_color) must be on GPU!
    """

    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    screenspace_points_densify = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()
        screenspace_points_densify.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug,
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    rotations = pc.get_rotation
    if is_6dof or pc.revolute:
        if torch.is_tensor(d_xyz) is False:
            if torch.is_tensor(d_rotation):
                #means3D = pc.get_xyz
                transformation1 = quaternion_to_rotation_matrix(d_rotation, get_R4= True)
                means3D = from_homogenous(
                torch.bmm(transformation1, to_homogenous(pc.get_xyz).unsqueeze(-1)).squeeze(-1))

                # transformation2 = quaternion_to_rotation_matrix(d_rotation)
                # stds = pc.get_scaling
                # means = torch.zeros((stds.size(0), 3), device="cuda")
                # samples = torch.normal(mean=means, std=stds)
                # means3D = torch.bmm(transformation2, samples.unsqueeze(-1)).squeeze(-1) + pc.get_xyz

                rotations = quaternion_multiply(d_rotation, rotations)
            else:
                means3D = pc.get_xyz


        else:
            means3D = from_homogenous(
                torch.bmm(d_xyz, to_homogenous(pc.get_xyz).unsqueeze(-1)).squeeze(-1))
    else:
        means3D = pc.get_xyz + d_xyz
    opacity = pc.get_opacity

    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
    # scaling / rotation by the rasterizer.
    scales = None

    cov3D_precomp = None
    if pipe.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)
    else:
        scales = pc.get_scaling + d_scaling
       #rotations = pc.get_rotation + d_rotation



    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
    shs = None
    colors_precomp = None
    if colors_precomp is None:
        if pipe.convert_SHs_python:
            shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree + 1) ** 2)
            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))
            dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
        else:
            shs = pc.get_features
    else:
        colors_precomp = override_color

    # Rasterize visible Gaussians to image, obtain their radii (on screen).
    rendered_image, radii, depth = rasterizer(
        means3D=means3D,
        means2D=screenspace_points,
        means2D_densify=screenspace_points_densify,
        shs=shs,
        colors_precomp=colors_precomp,
        opacities=opacity,
        scales=scales,
        rotations=rotations,
        cov3D_precomp=cov3D_precomp)

    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    # They will be excluded from value updates used in the splitting criteria.
    return {"render": rendered_image,
            "viewspace_points": screenspace_points,
            "viewspace_points_densify": screenspace_points_densify,
            "visibility_filter": radii > 0,
            "radii": radii,
            "depth": depth}

def part_render(viewpoint_camera, pc: GaussianModel, pipe, bg_color: torch.Tensor, fid= None, scaling_modifier=1.0, override_color=None):
    """
    Render the scene.

    Background tensor (bg_color) must be on GPU!
    """

    # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
    screenspace_points = torch.zeros_like(pc.get_xyz[pc.dynamic_part_mask], dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    screenspace_points_densify = torch.zeros_like(pc.get_xyz[pc.dynamic_part_mask], dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 0
    try:
        screenspace_points.retain_grad()
        screenspace_points_densify.retain_grad()
    except:
        pass

    # Set up rasterization configuration
    tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
    tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)

    raster_settings = GaussianRasterizationSettings(
        image_height=int(viewpoint_camera.image_height),
        image_width=int(viewpoint_camera.image_width),
        tanfovx=tanfovx,
        tanfovy=tanfovy,
        bg=bg_color,
        scale_modifier=scaling_modifier,
        viewmatrix=viewpoint_camera.world_view_transform,
        projmatrix=viewpoint_camera.full_proj_transform,
        sh_degree=pc.active_sh_degree,
        campos=viewpoint_camera.camera_center,
        prefiltered=False,
        debug=pipe.debug,
    )

    rasterizer = GaussianRasterizer(raster_settings=raster_settings)

    cov3D_precomp = None
    if pipe.compute_cov3D_python:
        cov3D_precomp = pc.get_covariance(scaling_modifier)[pc.dynamic_part_mask]
    else:
        scales = pc.get_scaling[pc.dynamic_part_mask]
        rotations = pc.get_rotation[pc.dynamic_part_mask]

    if fid is not None:
        if pc.revolute:

            #means3D = pc.get_xyz
            #dynamic_means3D = rigid_transform(pc, fid)
            #means3D[pc.dynamic_part_mask] = dynamic_means3D
            means3D = rigid_transform(pc, fid)[pc.dynamic_part_mask]
            if pc.use_canonical:
                scaling = (pc.canonical - fid) / pc.canonical
            else:
                scaling = fid
            axis, angle = quaternion_to_axis_angle(pc.quaternions)

            angle = scaling * angle
            q_rot = torch.tensor([[1.0, 0., 0., 0.]]).repeat(pc.get_rotation.shape[0], 1).cuda()
            qw = torch.cos(angle / 2)
            qx = axis[0] * torch.sin(angle / 2)
            qy = axis[1] * torch.sin(angle / 2)
            qz = axis[2] * torch.sin(angle / 2)
            q_rot[pc.dynamic_part_mask] = torch.tensor([qw, qx, qy, qz]).cuda()  # 四元数顺序为 (w, x, y, z)
            rotations = quaternion_multiply(q_rot, pc.get_rotation)[pc.dynamic_part_mask]

        elif pc.prismatic:
            if pc.use_canonical:
                scaling = (fid - pc.canonical) / pc.canonical
            else:
                scaling = fid
            #scaling = (pc.canonical - fid) / pc.canonical

            d_xyz = torch.zeros_like(pc.get_xyz)

            #d_xyz[pc.dynamic_part_mask] = scaling * pc.dir
            unit_d = F.normalize(pc.dir, p=2, dim=0)
            d_xyz[pc.dynamic_part_mask] = scaling * unit_d * pc.dist
            #d_xyz[pc.dynamic_part_mask] = unit_d * pc.dist * fid
            means3D = (pc.get_xyz + d_xyz)[pc.dynamic_part_mask]
        else:
            means3D = pc.get_xyz[pc.dynamic_part_mask]
    else:
        means3D = pc.get_xyz[pc.dynamic_part_mask]


    opacity = pc.get_opacity[pc.dynamic_part_mask]

    # If precomputed 3d covariance is provided, use it. If not, then it will be computed from
    # scaling / rotation by the rasterizer.




    # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
    # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
    shs = None
    colors_precomp = None
    if colors_precomp is None:
        if pipe.convert_SHs_python:
            shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree + 1) ** 2)
            dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))
            dir_pp_normalized = dir_pp / dir_pp.norm(dim=1, keepdim=True)
            sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)
            colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
        else:
            shs = pc.get_features[pc.dynamic_part_mask]
    else:
        colors_precomp = override_color

    # Rasterize visible Gaussians to image, obtain their radii (on screen).
    rendered_image, radii, depth = rasterizer(
        means3D=means3D,
        means2D=screenspace_points,
        means2D_densify=screenspace_points_densify,
        shs=shs,
        colors_precomp=colors_precomp,
        opacities=opacity,
        scales=scales,
        rotations=rotations,
        cov3D_precomp=cov3D_precomp)

    # Those Gaussians that were frustum culled or had a radius of 0 were not visible.
    # They will be excluded from value updates used in the splitting criteria.
    return {"render": rendered_image,
            "viewspace_points": screenspace_points,
            "viewspace_points_densify": screenspace_points_densify,
            "visibility_filter": radii > 0,
            "radii": radii,
            "depth": depth}