import torch

#####  Utils ###
def angle_bw_vectors_batch(v1, v2, eps=1e-12):
    ## Input shapes Tensors: Batch X 3 range of arccos ins [0, pi)
    return torch.acos(
        torch.clamp(
            torch.mul(v1[:, :3], v2[:, :3]).sum(dim=-1)
            / (torch.norm(v1[:, :3], dim=-1) * torch.norm(v2[:, :3], dim=-1) + eps),
            min=-1,
            max=1,
        )
    )


def distance_bw_plucker_lines(target, prediction, eps=1e-10):
    """Input shapes Tensors: Batch X #Images X 8
    # Based on formula from Plücker Coordinates for Lines in the Space by Prof. Yan-bin Jia
    # Verified by https://keisan.casio.com/exec/system/1223531414
    """
    norm_cross_prod = torch.norm(
        torch.cross(target[:, :, :3], prediction[:, :, :3], dim=-1), dim=-1
    )
    dist = torch.zeros_like(norm_cross_prod)

    # Checking for Parallel Lines
    if torch.any(norm_cross_prod <= eps):
        zero_idxs = (norm_cross_prod <= eps).nonzero(as_tuple=True)
        scales = (
            torch.norm(prediction[zero_idxs][:, :3], dim=-1)
            / torch.norm(target[zero_idxs][:, :3], dim=-1)
            + eps
        )
        dist[zero_idxs] = torch.norm(
            torch.cross(
                target[zero_idxs][:, :3],
                (
                    target[zero_idxs][:, 3:6]
                    - prediction[zero_idxs][:, 3:6] / scales.unsqueeze(-1)
                ),
            ),
            dim=-1,
        ) / (
            torch.mul(target[zero_idxs][:, :3], target[zero_idxs][:, :3]).sum(dim=-1)
            + eps
        )

    # Skew Lines: Non zero cross product
    nonzero_idxs = (norm_cross_prod > eps).nonzero(as_tuple=True)
    dist[nonzero_idxs] = torch.abs(
        torch.mul(target[nonzero_idxs][:, :3], prediction[nonzero_idxs][:, 3:6]).sum(
            dim=-1
        )
        + torch.mul(target[nonzero_idxs][:, 3:6], prediction[nonzero_idxs][:, :3]).sum(
            dim=-1
        )
    ) / (norm_cross_prod[nonzero_idxs] + eps)
    return dist


def orientation_difference_bw_plucker_lines(target, prediction, eps=1e-6):
    """Input shapes Tensors: Batch X #Images X 8
    range of arccos ins [0, pi)"""
    return torch.acos(
        torch.clamp(
            torch.mul(target[:, :, :3], prediction[:, :, :3]).sum(dim=-1)
            / (
                torch.norm(target[:, :, :3], dim=-1)
                * torch.norm(prediction[:, :, :3], dim=-1)
                + eps
            ),
            min=-1,
            max=1,
        )
    )


def theta_config_error(target, prediction):
    rot_tar = angle_axis_to_rotation_matrix(target[:, :, :3], target[:, :, 6]).view(
        -1, 3, 3
    )
    rot_pred = angle_axis_to_rotation_matrix(
        prediction[:, :, :3], prediction[:, :, 6]
    ).view(-1, 3, 3)
    I_ = torch.eye(3).reshape((1, 3, 3))
    I_ = I_.repeat(rot_tar.size(0), 1, 1).to(target.device)
    return torch.norm(
        I_ - torch.bmm(rot_pred, rot_tar.transpose(1, 2)), dim=(1, 2), p=2
    ).view(target.shape[:2])


def angle_axis_to_rotation_matrix(angle_axis, theta):
    # Stolen from PyTorch geometry library. Modified for our code
    angle_axis_shape = angle_axis.shape
    angle_axis_ = angle_axis.contiguous().view(-1, 3)
    theta_ = theta.contiguous().view(-1, 1)

    k_one = 1.0
    normed_axes = angle_axis_ / angle_axis_.norm(dim=-1, keepdim=True)
    wx, wy, wz = torch.chunk(normed_axes, 3, dim=1)
    cos_theta = torch.cos(theta_)
    sin_theta = torch.sin(theta_)

    r00 = cos_theta + wx * wx * (k_one - cos_theta)
    r10 = wz * sin_theta + wx * wy * (k_one - cos_theta)
    r20 = -wy * sin_theta + wx * wz * (k_one - cos_theta)
    r01 = wx * wy * (k_one - cos_theta) - wz * sin_theta
    r11 = cos_theta + wy * wy * (k_one - cos_theta)
    r21 = wx * sin_theta + wy * wz * (k_one - cos_theta)
    r02 = wy * sin_theta + wx * wz * (k_one - cos_theta)
    r12 = -wx * sin_theta + wy * wz * (k_one - cos_theta)
    r22 = cos_theta + wz * wz * (k_one - cos_theta)
    rotation_matrix = torch.cat([r00, r01, r02, r10, r11, r12, r20, r21, r22], dim=1)
    return rotation_matrix.view(list(angle_axis_shape[:-1]) + [3, 3])


def d_config_error(target, prediction):
    tar_d = target[:, :, 7].unsqueeze(-1)
    pred_d = prediction[:, :, 7].unsqueeze(-1)
    tar_d = target[:, :, :3] * tar_d
    pred_d = prediction[:, :, :3] * pred_d
    return (tar_d - pred_d).norm(dim=-1)


def quaternion_inner_product(q, r):
    assert q.shape[-1] == 4
    assert r.shape[-1] == 4
    original_shape = q.shape
    return torch.bmm(q.view(-1, 1, 4), r.view(-1, 4, 1)).view(original_shape[:-1])


def quaternion_inner_product(q, r):
    assert q.shape[-1] == 4
    assert r.shape[-1] == 4
    original_shape = q.shape
    return torch.bmm(q.view(-1, 1, 4), r.view(-1, 4, 1)).view(original_shape[:-1])


def difference_between_quaternions_tensors(q1, q2, eps=1e-6):
    return torch.acos(
        torch.clamp(2 * quaternion_inner_product(q1, q2) ** 2 - 1, -1 + eps, 1 - eps)
    )


###############


def maad(target_, prediction_):
    return (torch.abs(target_ - prediction_)).mean()


def angular_maad(target_, prediction_):
    maad_loss = target_.new_zeros((1,))
    for i in range(target_.size(0)):
        maad_loss += torch.abs(
            angle_bw_vectors_batch(target_[i, :, :], prediction_[i, :, :])
        ).mean()
    return maad_loss / target_.size(0)


## Screw Loss
def screw_loss(target_, pred_, detailed=False):
    """Based on Spatial distance
    Input shapes: Batch X Objects X images+
    """
    # Spatial Distance loss
    ori_error = orientation_difference_bw_plucker_lines(target_, pred_)
    dist_error = distance_bw_plucker_lines(target_, pred_)

    # Configuration Loss
    theta_error = theta_config_error(target_, pred_)
    d_error = d_config_error(target_, pred_)

    # orthogonal
    ortho_error = torch.mul(pred_[:, :, :3], pred_[:, :, 3:6]).sum(dim=-1)

    if detailed:
        return ori_error, dist_error, theta_error, d_error, ortho_error
    else:
        return (
            ori_error.abs().mean(),
            dist_error.abs().mean(),
            theta_error.abs().mean(),
            d_error.abs().mean(),
            ortho_error.abs().mean(),
        )
