from PIL import Image
import numpy as np
import torch
import torch.nn.functional as F
from scipy.spatial.transform import Rotation as R

def backproject(depth_maps, normal_maps, Ks, Es, rgb=None, masks=None):
    points3d = {}
    normals3d = {}
    rgb3d = {}
    for cam_id in depth_maps.keys():
        depth_map = depth_maps[cam_id]
        normal_map = normal_maps[cam_id]


        ys = np.arange(depth_map.shape[0])
        xs = np.arange(depth_map.shape[1])
        p_screen = np.dstack(np.meshgrid(xs, ys, [1])).reshape((-1, 3))
        depth_mask = (depth_map > 0)  & (depth_map < 1.4)
        if masks is not None:
            # upsample mask
            I = Image.fromarray(masks[cam_id])
            I = I.resize((I.size[0]*2, I.size[1]*2))
            depth_mask = np.logical_and(depth_mask, np.array(I).astype(np.bool))
        depths = depth_map[depth_mask]
        p_screen = p_screen[depth_mask.reshape(-1)]
        p_screen_canonical = p_screen @ Ks[cam_id].invert().T
        p_cam = p_screen_canonical * np.expand_dims(depths, 1)
        p_cam_hom = np.hstack([p_cam, np.ones((p_cam.shape[0], 1))])
        p_world = p_cam_hom @ Es[cam_id].T
        ns = np.ones_like(p_world)
        ns[:, :3] = normal_map[depth_mask]
        n_world = ns @ Es[cam_id].T

        points3d[cam_id] = p_world[:, :3]
        normals3d[cam_id] = n_world[:, :3]
        if rgb is not None:
            rgb_lin = rgb[cam_id].reshape((-1, 3))
            rgb_valid = rgb_lin[depth_mask.reshape(-1)]
            rgb3d[cam_id] = rgb_valid

    if rgb is None:
        return points3d, normals3d
    else:
        return points3d, normals3d, rgb3d


def get_view_dirs(Ks, Es, image_shape, rgb=None, masks=None):
    points3d = {}
    view_dirs = {}
    for cam_id in Ks.keys():
        ys = np.arange(image_shape[0])
        xs = np.arange(image_shape[1])
        p_screen = np.dstack(np.meshgrid(xs, ys, [1])).reshape((-1, 3))
        if masks is not None:
            # upsample mask
            I = Image.fromarray(masks[cam_id])
            I = I.resize((I.size[0]*2, I.size[1]*2))
            depth_mask = np.logical_and(depth_mask, np.array(I).astype(np.bool))
        p_screen = np.reshape(p_screen, [-1, 3])
        p_screen_canonical = p_screen @ Ks[cam_id].invert().T
        p_cam = p_screen_canonical * 1
        p_cam_hom = np.hstack([p_cam, np.ones((p_cam.shape[0], 1))])
        p_world = p_cam_hom @ Es[cam_id].T

        points3d[cam_id] = p_world[:, :3]

        origin = Es[cam_id][:3, 3]
        view_dirs[cam_id] = p_world[:, :3] - origin
        view_dirs[cam_id] /= np.linalg.norm(view_dirs[cam_id], axis=-1, keepdims=True)

        #if rgb is not None:
        #    rgb_lin = rgb[cam_id].reshape((-1, 3))
        #    rgb_valid = rgb_lin[depth_mask.reshape(-1)]
        #    rgb3d[cam_id] = rgb_valid

    #if rgb is None:
    #    return points3d, normals3d
    #else:

    return view_dirs





    return


def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
    """
    Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
    using Gram--Schmidt orthogonalization per Section B of [1].
    Args:
        d6: 6D rotation representation, of size (*, 6)

    Returns:
        batch of rotation matrices of size (*, 3, 3)

    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
    On the Continuity of Rotation Representations in Neural Networks.
    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
    Retrieved from http://arxiv.org/abs/1812.07035
    """

    a1, a2 = d6[..., :3], d6[..., 3:]
    b1 = F.normalize(a1, dim=-1)
    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
    b2 = F.normalize(b2, dim=-1)
    b3 = torch.cross(b1, b2, dim=-1)
    return torch.stack((b1, b2, b3), dim=-2)


def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
    """
    Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
    by dropping the last row. Note that 6D representation is not unique.
    Args:
        matrix: batch of rotation matrices of size (*, 3, 3)

    Returns:
        6D rotation representation, of size (*, 6)

    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
    On the Continuity of Rotation Representations in Neural Networks.
    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
    Retrieved from http://arxiv.org/abs/1812.07035
    """
    batch_dim = matrix.size()[:-2]
    return matrix[..., :2, :].clone().reshape(batch_dim + (6,))


def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
    """
    Return the rotation matrices for one of the rotations about an axis
    of which Euler angles describe, for each value of the angle given.

    Args:
        axis: Axis label "X" or "Y or "Z".
        angle: any shape tensor of Euler angles in radians

    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """

    cos = torch.cos(angle)
    sin = torch.sin(angle)
    one = torch.ones_like(angle)
    zero = torch.zeros_like(angle)

    if axis == "X":
        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
    elif axis == "Y":
        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
    elif axis == "Z":
        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
    else:
        raise ValueError("letter must be either X, Y or Z.")

    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))


def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor:
    """
    Convert rotations given as Euler angles in radians to rotation matrices.

    Args:
        euler_angles: Euler angles in radians as tensor of shape (..., 3).
        convention: Convention string of three uppercase letters from
            {"X", "Y", and "Z"}.

    Returns:
        Rotation matrices as tensor of shape (..., 3, 3).
    """
    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
        raise ValueError("Invalid input euler angles.")
    if len(convention) != 3:
        raise ValueError("Convention must have 3 letters.")
    if convention[1] in (convention[0], convention[2]):
        raise ValueError(f"Invalid convention {convention}.")
    for letter in convention:
        if letter not in ("X", "Y", "Z"):
            raise ValueError(f"Invalid letter {letter} in convention string.")
    matrices = [
        _axis_angle_rotation(c, e)
        for c, e in zip(convention, torch.unbind(euler_angles, -1))
    ]
    # return functools.reduce(torch.matmul, matrices)
    return torch.matmul(torch.matmul(matrices[0], matrices[1]), matrices[2])

def rodrigues_to_rot_vec(
    rot_mats: torch.Tensor,
    epsilon: float = 1e-8,
) -> torch.Tensor:
    ''' Calculates the rotation vectors for a batch of rotation matrices
        Parameters
        ----------
        rot_mats: torch.tensor Nx3x3
            The rotation matrices
        Returns
        -------
        rot_vecs: torch.tensor Nx3
            The axis-angle representations of the rotation matrices
    '''
    batch_size = rot_mats.shape[0]
    device, dtype = rot_mats.device, rot_mats.dtype

    # Extract the trace from the rotation matrix: tr(R) = 1 + 2*cos(angle)
    trace = torch.einsum('bii->b', rot_mats)
    angle = torch.acos(torch.clamp((trace - 1) / 2, -1 + epsilon, 1 - epsilon))

    # The axis is encoded in the skew-symmetric part of R
    # (R - R^T) / (2 * sin(angle))
    axis = torch.stack([
        rot_mats[:, 2, 1] - rot_mats[:, 1, 2],
        rot_mats[:, 0, 2] - rot_mats[:, 2, 0],
        rot_mats[:, 1, 0] - rot_mats[:, 0, 1]
    ], dim=1)

    # Handle the case where the angle is close to 0 or pi
    sin_angle = torch.sin(angle)
    mask = sin_angle.abs() < epsilon

    # If sin(angle) is close to 0, the angle is 0 or pi
    # If angle is 0, axis is undefined but rot_vec is the zero vector
    # If angle is pi, we need to find the axis from R + I
    safe_sin_angle = torch.where(mask, torch.ones_like(sin_angle), sin_angle)
    axis = axis / (2 * safe_sin_angle.unsqueeze(-1))

    # For angles close to pi, the axis needs to be recomputed
    pi_mask = mask & (torch.cos(angle) < 0)
    if torch.any(pi_mask):
        # In this case, R = I + 2 * axis * axis^T
        # (R + I) / 2 = axis * axis^T
        diag = torch.einsum('bii->bi', rot_mats[pi_mask])
        axis_pi = torch.sqrt(torch.clamp((diag + 1) / 2, min=0))
        # Find the largest component to determine the sign
        _, max_idx = torch.max(axis_pi, dim=1)
        signs = torch.sign(torch.stack([
            rot_mats[pi_mask, 0, 1] + rot_mats[pi_mask, 1, 0],
            rot_mats[pi_mask, 0, 2] + rot_mats[pi_mask, 2, 0],
            rot_mats[pi_mask, 1, 2] + rot_mats[pi_mask, 2, 1]
        ], dim=1)[torch.arange(axis_pi.shape[0]), max_idx])
        
        # Correct the signs of the axes
        axis_pi_corrected = axis_pi * signs.unsqueeze(-1)
        axis[pi_mask] = axis_pi_corrected / torch.norm(axis_pi_corrected, dim=1, keepdim=True)
    
    rot_vecs = axis * angle.unsqueeze(-1)
    return rot_vecs

def rotation_6d_to_rot_vec(
    rot_6d: torch.Tensor,
    epsilon: float = 1e-8,
) -> torch.Tensor:
    """
    Converts 6D rotation representation to rotation vector.
    
    Args:
        rot_6d: 6D rotation representation, of size (*, 6)
        epsilon: small value to avoid division by zero

    Returns:
        Rotation vectors of size (*, 3)
    """
    rot_mats = rotation_6d_to_matrix(rot_6d)
    return rodrigues_to_rot_vec(rot_mats, epsilon)


def rigid_alignment_npy(source: np.ndarray, target: np.ndarray):
    """
    Aligns source (N x 3) to target (N x 3) via least-squares rigid transform (no scaling).
    Returns: (R, t) such that aligned = source @ R.T + t
    """
    assert source.shape == target.shape

    # Subtract centroids
    src_center = source.mean(axis=0)
    tgt_center = target.mean(axis=0)
    src_centered = source - src_center
    tgt_centered = target - tgt_center

    # Compute rotation via SVD
    H = src_centered.T @ tgt_centered
    U, _, Vt = np.linalg.svd(H)
    R = Vt.T @ U.T

    # Handle reflection
    if np.linalg.det(R) < 0:
        Vt[-1, :] *= -1
        R = Vt.T @ U.T

    t = tgt_center - src_center @ R.T
    return R, t

def rigid_alignment(src: torch.Tensor, tgt: torch.Tensor):
    """
    Batched rigid alignment using SVD (Kabsch algorithm).
    
    Args:
        src: (B, N, 3) source point sets
        tgt: (B, N, 3) target point sets
    Returns:
        R: (B, 3, 3) rotation matrices
        t: (B, 3) translation vectors
    """
    if (src.ndim == 2) and (tgt.ndim == 2):
        src = src[None] 
        tgt = tgt[None]
    assert src.shape == tgt.shape and src.ndim == 3 and src.shape[-1] == 3

    src_mean = src.mean(dim=1, keepdim=True)   # (B, 1, 3)
    tgt_mean = tgt.mean(dim=1, keepdim=True)

    src_centered = src - src_mean              # (B, N, 3)
    tgt_centered = tgt - tgt_mean

    H = torch.matmul(src_centered.transpose(1, 2), tgt_centered)  # (B, 3, 3)
    U, S, Vt = torch.linalg.svd(H)  # Vt: (B, 3, 3)

    # Compute R = V U^T
    R = torch.matmul(Vt.transpose(-2, -1), U.transpose(-2, -1))  # (B, 3, 3)

    # Handle reflection case
    det = torch.det(R)
    mask = det < 0
    if mask.any():
        Vt[mask, -1, :] *= -1
        R[mask] = torch.matmul(Vt[mask].transpose(-2, -1), U[mask].transpose(-2, -1))

    t = tgt_mean.squeeze(1) - torch.matmul(R, src_mean.squeeze(1).unsqueeze(-1)).squeeze(-1)  # (B, 3)

    return R, t