import torch
import numpy as np


def batch_episym(x1, x2, F):
    batch_size, num_pts = x1.shape[0], x1.shape[1]
    x1 = torch.cat([x1, x1.new_ones(batch_size, num_pts, 1)], dim=-1).reshape(
        batch_size, num_pts, 3, 1
    )
    x2 = torch.cat([x2, x2.new_ones(batch_size, num_pts, 1)], dim=-1).reshape(
        batch_size, num_pts, 3, 1
    )
    F = F.reshape(-1, 1, 3, 3).repeat(1, num_pts, 1, 1)
    x2Fx1 = torch.matmul(x2.transpose(2, 3), torch.matmul(F, x1)).reshape(
        batch_size, num_pts
    )
    Fx1 = torch.matmul(F, x1).reshape(batch_size, num_pts, 3)
    Ftx2 = torch.matmul(F.transpose(2, 3), x2).reshape(batch_size, num_pts, 3)
    ys = (
        x2Fx1**2
        * (
            1.0 / (Fx1[:, :, 0] ** 2 + Fx1[:, :, 1] ** 2 + 1e-15)
            + 1.0 / (Ftx2[:, :, 0] ** 2 + Ftx2[:, :, 1] ** 2 + 1e-15)
        )
    ).sqrt()
    return ys


def CELoss(seed_x1, seed_x2, e, confidence, inlier_th, batch_mask=1):
    # seed_x: b*k*2
    ys = batch_episym(seed_x1, seed_x2, e)
    mask_pos, mask_neg = (ys <= inlier_th).float(), (ys > inlier_th).float()
    num_pos, num_neg = (
        torch.relu(torch.sum(mask_pos, dim=1) - 1.0) + 1.0,
        torch.relu(torch.sum(mask_neg, dim=1) - 1.0) + 1.0,
    )
    loss_pos, loss_neg = (
        -torch.log(abs(confidence) + 1e-8) * mask_pos,
        -torch.log(abs(1 - confidence) + 1e-8) * mask_neg,
    )
    classif_loss = torch.mean(
        loss_pos * 0.5 / num_pos.unsqueeze(-1) + loss_neg * 0.5 / num_neg.unsqueeze(-1),
        dim=-1,
    )
    classif_loss = classif_loss * batch_mask
    classif_loss = classif_loss.mean()
    precision = torch.mean(
        torch.sum((confidence > 0.5).type(confidence.type()) * mask_pos, dim=1)
        / (torch.sum((confidence > 0.5).type(confidence.type()), dim=1) + 1e-8)
    )
    recall = torch.mean(
        torch.sum((confidence > 0.5).type(confidence.type()) * mask_pos, dim=1)
        / num_pos
    )
    return classif_loss, precision, recall


def CorrLoss(desc_mat, batch_num_corr, batch_num_incorr1, batch_num_incorr2):
    total_loss_corr, total_loss_incorr = 0, 0
    total_acc_corr, total_acc_incorr = 0, 0
    batch_size = desc_mat.shape[0]
    log_p = torch.log(abs(desc_mat) + 1e-8)

    for i in range(batch_size):
        cur_log_p = log_p[i]
        num_corr = batch_num_corr[i]
        num_incorr1, num_incorr2 = batch_num_incorr1[i], batch_num_incorr2[i]

        # loss and acc
        loss_corr = -torch.diag(cur_log_p)[:num_corr].mean()
        loss_incorr = (
            -cur_log_p[num_corr : num_corr + num_incorr1, -1].mean()
            - cur_log_p[-1, num_corr : num_corr + num_incorr2].mean()
        ) / 2

        value_row, row_index = torch.max(desc_mat[i, :-1, :-1], dim=-1)
        value_col, col_index = torch.max(desc_mat[i, :-1, :-1], dim=-2)
        acc_incorr = (
            (value_row[num_corr : num_corr + num_incorr1] < 0.2).float().mean()
            + (value_col[num_corr : num_corr + num_incorr2] < 0.2).float().mean()
        ) / 2

        acc_row_mask = row_index[:num_corr] == torch.arange(num_corr).cuda()
        acc_col_mask = col_index[:num_corr] == torch.arange(num_corr).cuda()
        acc = (acc_col_mask & acc_row_mask).float().mean()

        total_loss_corr += loss_corr
        total_loss_incorr += loss_incorr
        total_acc_corr += acc
        total_acc_incorr += acc_incorr

    total_acc_corr /= batch_size
    total_acc_incorr /= batch_size
    total_loss_corr /= batch_size
    total_loss_incorr /= batch_size
    return total_loss_corr, total_loss_incorr, total_acc_corr, total_acc_incorr


class SGMLoss:
    def __init__(self, config, model_config):
        self.config = config
        self.model_config = model_config

    def run(self, data, result):
        loss_corr, loss_incorr, acc_corr, acc_incorr = CorrLoss(
            result["p"], data["num_corr"], data["num_incorr1"], data["num_incorr2"]
        )
        loss_mid_corr_tower, loss_mid_incorr_tower, acc_mid_tower = [], [], []

        # mid loss
        for i in range(len(result["mid_p"])):
            mid_p = result["mid_p"][i]
            loss_mid_corr, loss_mid_incorr, mid_acc_corr, mid_acc_incorr = CorrLoss(
                mid_p, data["num_corr"], data["num_incorr1"], data["num_incorr2"]
            )
            loss_mid_corr_tower.append(loss_mid_corr), loss_mid_incorr_tower.append(
                loss_mid_incorr
            ), acc_mid_tower.append(mid_acc_corr)
        if len(result["mid_p"]) != 0:
            loss_mid_corr_tower, loss_mid_incorr_tower, acc_mid_tower = (
                torch.stack(loss_mid_corr_tower),
                torch.stack(loss_mid_incorr_tower),
                torch.stack(acc_mid_tower),
            )
        else:
            loss_mid_corr_tower, loss_mid_incorr_tower, acc_mid_tower = (
                torch.zeros(1).cuda(),
                torch.zeros(1).cuda(),
                torch.zeros(1).cuda(),
            )

        # seed confidence loss
        classif_loss_tower, classif_precision_tower, classif_recall_tower = [], [], []
        for layer in range(len(result["seed_conf"])):
            confidence = result["seed_conf"][layer]
            seed_index = result["seed_index"][
                (np.asarray(self.model_config.seedlayer) <= layer).nonzero()[0][-1]
            ]
            seed_x1, seed_x2 = data["x1"].gather(
                dim=1, index=seed_index[:, :, 0, None].expand(-1, -1, 2)
            ), data["x2"].gather(
                dim=1, index=seed_index[:, :, 1, None].expand(-1, -1, 2)
            )
            classif_loss, classif_precision, classif_recall = CELoss(
                seed_x1, seed_x2, data["e_gt"], confidence, self.config.inlier_th
            )
            classif_loss_tower.append(classif_loss), classif_precision_tower.append(
                classif_precision
            ), classif_recall_tower.append(classif_recall)
        classif_loss, classif_precision_tower, classif_recall_tower = (
            torch.stack(classif_loss_tower).mean(),
            torch.stack(classif_precision_tower),
            torch.stack(classif_recall_tower),
        )

        classif_loss *= self.config.seed_loss_weight
        loss_mid_corr_tower *= self.config.mid_loss_weight
        loss_mid_incorr_tower *= self.config.mid_loss_weight
        total_loss = (
            loss_corr
            + loss_incorr
            + classif_loss
            + loss_mid_corr_tower.sum()
            + loss_mid_incorr_tower.sum()
        )

        return {
            "loss_corr": loss_corr,
            "loss_incorr": loss_incorr,
            "acc_corr": acc_corr,
            "acc_incorr": acc_incorr,
            "loss_seed_conf": classif_loss,
            "pre_seed_conf": classif_precision_tower,
            "recall_seed_conf": classif_recall_tower,
            "loss_corr_mid": loss_mid_corr_tower,
            "loss_incorr_mid": loss_mid_incorr_tower,
            "mid_acc_corr": acc_mid_tower,
            "total_loss": total_loss,
        }


class SGLoss:
    def __init__(self, config, model_config):
        self.config = config
        self.model_config = model_config

    def run(self, data, result):
        loss_corr, loss_incorr, acc_corr, acc_incorr = CorrLoss(
            result["p"], data["num_corr"], data["num_incorr1"], data["num_incorr2"]
        )
        total_loss = loss_corr + loss_incorr
        return {
            "loss_corr": loss_corr,
            "loss_incorr": loss_incorr,
            "acc_corr": acc_corr,
            "acc_incorr": acc_incorr,
            "total_loss": total_loss,
        }