import numpy as np
import megengine.functional as F
from common import se3, so3


def compute_losses(endpoints, params):
    loss = {}
    # compute losses
    if params.loss_type == "finet":
        num_iter = len(endpoints["all_pose_pair"])
        triplet_loss = {}
        for i in range(num_iter):
            # reg loss
            pose_pair = endpoints["all_pose_pair"][i]
            loss["quat_{}".format(i)] = F.nn.l1_loss(pose_pair[0][:, :4], pose_pair[1][:, :4]) * params.loss_alpha1
            loss["translate_{}".format(i)] = 
F.nn.square_loss(pose_pair[0][:, 4:], pose_pair[1][:, 4:])