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


def compute_losses(data_batch, endpoints, params):
    loss = {}
    # compute losses
    if params.loss_type == "omnet":
        num_iter = len(endpoints["all_pose_pair"])
        for i in range(num_iter):
            # mask loss
            src_cls_pair, ref_cls_pair = endpoints["all_src_cls_pair"][i], endpoints["all_ref_cls_pair"][i]
            src_cls = F.nn.frequency_weighted_cross_entropy(src_cls_pair[1], src_cls_pair[0], weight=mge.tensor([0.7, 0.3]))
            ref_cls = F.nn.frequency_weighted_cross_entropy(ref_cls_pair[1], ref_cls_pair[0], weight=mge.tensor([0.7, 0.3]))
            loss["cls_{}".format(i)] = (src_cls + ref_cls) / 2.0
            # reg loss
            pose_pair = endpoints["all_pose_pair"][i]
            loss["quat_{}".format(i)] = F.nn.l1_loss(pose_pair[1][:, :4], pose_pair[0][:, :4]) * params.loss_alpha1
            loss["translate_{}".format(i)] = F.nn.square_loss(pose_pair[1][:, 4:], pose_pair[0][:, 4:]) * params.loss_alpha2
        # total loss
        total_losses = []
        for k in loss:
            total_losses.append(loss[k])
        loss["total"] = F.sum(
F.concat(total_losses)