import torch

import torch.nn as nn

import math

def cart2polar(coord):

    """ coord: (N, 2, ...)

    """

    x = coord[:, 0, ...]

    y = coord[:, 1, ...]

    theta = torch.atan(y / (x + 1e-12))

    theta = theta + (x < 0).to(coord.dtype) * math.pi

    theta = theta + ((x > 0).to(coord.dtype) * (y < 0).to(coord.dtype)) * 2 * math.pi

    return theta / (2 * math.pi)

class EuclideanAngleLossWithOHEM(nn.Module):

    def __init__(self, npRatio=3):

        super(EuclideanAngleLossWithOHEM, self).__init__()

        self.npRatio = npRatio

    def __cal_weight(self, gt):

        _, H, W = gt.shape  # N=1

        labels = torch.unique(gt, sorted=True)[1:]

        weight = torch.zeros((H, W), dtype=torch.float, device=gt.device)

        posRegion = gt[0, ...] > 0

        posCount = torch.sum(posRegion)

        if posCount != 0:

            segRemain = 0

            for segi in labels:

                overlap_segi = gt[0, ...] == segi

                overlapCount_segi = torch.sum(overlap_segi)

                if overlapCount_segi == 0: continue

                segRemain = segRemain + 1

            segAve = float(posCount) / segRemain

            for segi in labels:

                overlap_segi = gt[0, ...] == segi

                overlapCount_segi = torch.sum(overlap_segi, dtype=torch.float)

                if overlapCount_segi == 0: continue

                pixAve = segAve / overlapCount_segi

                weight = weight * (~overlap_segi).to(torch.float) + pixAve * overlap_segi.to(torch.float)

        # weight = weight[None]

        return weight

    def forward(self, pred, gt_df, gt, weight=None):

        """ pred: (N, C, H, W)

            gt_df: (N, C, H, W)

            gt: (N, 1, H, W)

        """

        # L1 and L2 distance

        N, _, H, W = pred.shape

        distL1 = pred - gt_df

        distL2 = distL1 ** 2

        theta_p = cart2polar(pred)

        theta_g = cart2polar(gt_df)

        angleDistL1 = theta_g - theta_p

        if weight is None:

            weight = torch.zeros((N, H, W), device=pred.device)

            for i in range(N):

                weight[i] = self.__cal_weight(gt[i])

        # the amount of positive and negtive pixels

        regionPos = (weight > 0).to(torch.float)

        regionNeg = (weight == 0).to(torch.float)

        sumPos = torch.sum(regionPos, dim=(1,2))  # (N,)

        sumNeg = torch.sum(regionNeg, dim=(1,2))

        # the amount of hard negative pixels

        sumhardNeg = torch.min(self.npRatio * sumPos, sumNeg).to(torch.int)  # (N,)

        # angle loss on ~(top - sumhardNeg) negative pixels to 0

        angleLossNeg = (angleDistL1 ** 2) * regionNeg

        angleLossNegFlat = torch.flatten(angleLossNeg, start_dim=1)  # (N, ...)

        # set loss on ~(top - sumhardNeg) negative pixels to 0

        lossNeg = (distL2[:,0,...] + distL2[:, 1, ...]) * regionNeg

        lossFlat = torch.flatten(lossNeg, start_dim=1)  # (N, ...)

        # l2-norm distance and angle distance

        lossFlat = lossFlat + angleLossNegFlat

        arg = torch.argsort(lossFlat, dim=1)

        for i in range(N):

            lossFlat[i, arg[i, :-sumhardNeg[i]]] = 0

        lossHard = lossFlat.view(lossNeg.shape)

        # weight for positive and negative pixels

        weightPos = torch.zeros_like(gt, dtype=pred.dtype)

        weightNeg = torch.zeros_like(gt, dtype=pred.dtype)

        weightPos = weight.clone()

        weightNeg[:,0,...] = (lossHard != 0).to(torch.float32)

        # total loss

        total_loss = torch.sum(((distL2[:,0,...] + distL2[:, 1, ...]) + angleDistL1 ** 2) *

                               (weightPos + weightNeg)) / N / 2. / torch.sum(weightPos + weightNeg)

        return total_loss

if __name__ == "__main__":

    import os

    import torch.nn as nn

    import torch.optim as optim

    os.environ["CUDA_VISIBLE_DEVICES"] = "4"

    criterion = EuclideanAngleLossWithOHEM()

    # models = nn.Sequential(nn.Conv2d(2, 2, 1),

    #                        nn.ReLU())

    # models.to(device="cuda")

    # epoch_n = 200

    # learning_rate = 1e-4

    # optimizer = optim.Adam(params=models.parameters(), lr=learning_rate)

    # for i in range(100):

    #     pred = torch.randn((32, 2, 224, 224)).cuda()

    #     gt_df = torch.randn((32, 2, 224, 224)).cuda()

    #     gt = torch.randint(0, 4, (32, 1, 224, 224)).cuda()

    #     pred = models(gt_df)

    #     loss = criterion(pred, gt_df, gt)

    #     optimizer.zero_grad()

    #     loss.backward()

    #     optimizer.step()

    #     print("{:6} loss:{}".format(i, loss))

    for i in range(100):

        pred = torch.randn((32, 2, 224, 224)).cuda()

        gt_df = torch.randn((32, 2, 224, 224)).cuda()

        gt = torch.randint(0, 4, (32, 1, 224, 224)).cuda()

        loss = criterion(-gt_df, gt_df, gt)

        print("{:6} loss:{}".format(i, loss))