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:]) * params.loss_alpha2

            # transformation sensitivity loss (TSL)
            if i < 2:
                all_R_feats = endpoints["all_R_feats"][i]
                all_t_feats = endpoints["all_t_feats"][i]
                # R feats triplet loss
                R_feats_pos = F.nn.square_loss(all_t_feats[0], all_t_feats[1])
                R_feats_neg = F.nn.square_loss(all_R_feats[0], all_R_feats[1])
                triplet_loss["R_feats_triplet_pos_{}".format(i)] = R_feats_pos
                triplet_loss["R_feats_triplet_neg_{}".format(i)] = R_feats_neg
                loss["R_feats_triplet_{}".format(i)] = (F.clip(-R_feats_neg + params.margin[i], lower=0.0) +
                                                        R_feats_pos) * params.loss_alpha3
                # t feats triplet loss
                t_feats_pos = F.nn.square_loss(all_R_feats[0], all_R_feats[2])
                t_feats_neg = F.nn.square_loss(all_t_feats[0], all_t_feats[2])
                triplet_loss["t_feats_triplet_pos_{}".format(i)] = t_feats_pos
                triplet_loss["t_feats_triplet_neg_{}".format(i)] = t_feats_neg
                loss["t_feats_triplet_{}".format(i)] = (F.clip(-t_feats_neg + params.margin[i], lower=0.0) +
                                                        t_feats_pos) * params.loss_alpha3

            # point-wise feature dropout loss (PFDL)
            all_dropout_R_feats = endpoints["all_dropout_R_feats"][i]
            all_dropout_t_feats = endpoints["all_dropout_t_feats"][i]
            loss["src_R_feats_dropout_{}".format(i)] = F.nn.square_loss(all_dropout_R_feats[0], all_dropout_R_feats[1]) * params.loss_alpha4
            loss["ref_R_feats_dropout_{}".format(i)] = F.nn.square_loss(all_dropout_R_feats[2], all_dropout_R_feats[3]) * params.loss_alpha4
            loss["src_t_feats_dropout_{}".format(i)] = F.nn.square_loss(all_dropout_t_feats[0], all_dropout_t_feats[1]) * params.loss_alpha4
            loss["ref_t_feats_dropout_{}".format(i)] = F.nn.square_loss(all_dropout_t_feats[2], all_dropout_t_feats[3]) * params.loss_alpha4
        # total loss
        total_losses = []
        for k in loss:
            total_losses.append(loss[k])
        loss["total"] = F.sum(
F.concat(total_losses)