# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by <NAME> (<EMAIL>)
# ------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import torch
import torch.nn as nn
import torch.nn.functional as F


class KeypointIOULoss(nn.Module):
    def __init__(self, use_target_weight):
        super().__init__()
        self.use_target_weight = use_target_weight

    # def forward(self, output, target, target_weight):
    #     # output [bs, 14, 2]
    #     # target same
    #     # target_weight [bs, 14, 1]
    #     target = target.float()
    #     inter = torch.min(output, target)
    #     union = output + target - inter
    #     losses = -torch.log((inter + 1.0) / (union + 1.0))
    #     # print('angle loss : ', losses[:, :, 0].sum() * 0.5, " length loss : ", losses[:, :, 1].sum() * 0.5)
    #     losses = losses[:, :, 0] * 0.5 + losses[:, :, 1] * 0.5
    #     losses.unsqueeze_(2)
    #     # print('pred : ', output[0])
    #     # print('target : ', target[0])

    #     if self.use_target_weight and target_weight.sum() > 0:
    #         return (losses * target_weight).sum() / (target_weight.sum())
    #     else:
    #         losses.mean()

    # def forward(self, output, target, target_weight):
    #     # output [bs, 16, 361]
    #     index = target_weight > 0
    #     # output [bs * 16]
    #     index = index.view(-1)
    #     degree = F.cross_entropy(output[:, :, :361].view(-1, 361)[index, :], target[:, :, 0].view(-1).long()[index])
        
    #     length_target = target[:, :, 1].float().view(-1)[index]
    #     length_pred = output[:, :, 361].view(-1)[index]
    #     inter = torch.min(length_pred, length_target)
    #     union = length_pred + length_target - inter
    #     length = -torch.log((inter + 1.) / (union + 1.))
    #     length = length.sum() / index.sum()
    #     return degree + 0.5 + length * 0.5
    
    def forward(self, output, target, target_weight):
        # index (bs * joints)
        # print('output: ', output[0])
        # print('target: ', target[0])
        index = target_weight > 0
        index = index.view(-1)

        output = output.view(-1, 3)[index, :]
        target = target.view(-1, 3)[index, :].float()

        angle_loss = F.mse_loss(output[:, :2], target[:, :2])

        inter = torch.min(output[:, 2], target[:, 2])
        union = output[:, 2] + target[:, 2] - inter
        length_loss = -torch.log((inter + 1.) / (union + 1.))
        length_loss = length_loss.sum() / index.sum()
        # length_loss = F.mse_loss(output[:, 9], target[:, 9])
        # print('angle : ', angle_loss, " length : ", length_loss)
        
        return 0.5 * angle_loss + 0.5 * length_loss

    # def forward(self, output, target, target_weight):
    #     # index (bs * joints)
    #     index = target_weight > 0
    #     index = index.view(-1)

    #     output = output.view(-1, 10)[index, :]
    #     target = target.view(-1, 10)[index, :].float()

    #     angle_loss = F.mse_loss(output[:, :9].view(-1, 3, 3)[:, :2, :2], target[:, :9].view(-1, 3, 3)[:, :2, :2])
    #     length_loss = F.mse_loss(output[:, 9], target[:, 9])
    #     return 0.5 * angle_loss + 0.5 * length_loss
    #     # length_loss = F.mse_loss(output[:, 9], target[:, 9])
    #     # print('angle : ', angle_loss, " length : ", length_loss)
        
    #     # return 0.5 * angle_loss + 0.5 * length_loss

    #     # length = -torch.log((inter + 1.0) / ( + 1.0))-torch.log(output[:, :, 361], target[:, :, 1].float())

class JointsMSELoss(nn.Module):
    def __init__(self, use_target_weight):
        super(JointsMSELoss, self).__init__()
        self.criterion = nn.MSELoss(reduction='mean')
        self.use_target_weight = use_target_weight

    def forward(self, output, target, target_weight):
        target = target.float()
        batch_size = output.size(0)
        num_joints = output.size(1)
        heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
        heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
        loss = 0

        for idx in range(num_joints):
            heatmap_pred = heatmaps_pred[idx].squeeze()
            heatmap_gt = heatmaps_gt[idx].squeeze()
            if self.use_target_weight:
                loss += 0.5 * self.criterion(
                    heatmap_pred.mul(target_weight[:, idx]),
                    heatmap_gt.mul(target_weight[:, idx])
                )
            else:
                loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)

        return loss / num_joints


class JointsOHKMMSELoss(nn.Module):
    def __init__(self, use_target_weight, topk=8):
        super(JointsOHKMMSELoss, self).__init__()
        self.criterion = nn.MSELoss(reduction='none')
        self.use_target_weight = use_target_weight
        self.topk = topk

    def ohkm(self, loss):
        ohkm_loss = 0.
        for i in range(loss.size()[0]):
            sub_loss = loss[i]
            topk_val, topk_idx = torch.topk(
                sub_loss, k=self.topk, dim=0, sorted=False
            )
            tmp_loss = torch.gather(sub_loss, 0, topk_idx)
            ohkm_loss += torch.sum(tmp_loss) / self.topk
        ohkm_loss /= loss.size()[0]
        return ohkm_loss

    def forward(self, output, target, target_weight):
        batch_size = output.size(0)
        num_joints = output.size(1)
        heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
        heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)

        loss = []
        for idx in range(num_joints):
            heatmap_pred = heatmaps_pred[idx].squeeze()
            heatmap_gt = heatmaps_gt[idx].squeeze()
            if self.use_target_weight:
                loss.append(0.5 * self.criterion(
                    heatmap_pred.mul(target_weight[:, idx]),
                    heatmap_gt.mul(target_weight[:, idx])
                ))
            else:
                loss.append(
                    0.5 * self.criterion(heatmap_pred, heatmap_gt)
                )

        loss = [l.mean(dim=1).unsqueeze(dim=1) for l in loss]
        loss = torch.cat(loss, dim=1)

        return self.ohkm(loss)

# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by <NAME> (<EMAIL>)
# ------------------------------------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import torch
import torch.nn as nn
import torch.nn.functional as F


class KeypointIOULoss(nn.Module):
    def __init__(self, use_target_weight):
        super().__init__()
        self.use_target_weight = use_target_weight

    # def forward(self, output, target, target_weight):
    #     # output [bs, 14, 2]
    #     # target same
    #     # target_weight [bs, 14, 1]
    #     target = target.float()
    #     inter = torch.min(output, target)
    #     union = output + target - inter
    #     losses = -torch.log((inter + 1.0) / (union + 1.0))
    #     # print('angle loss : ', losses[:, :, 0].sum() * 0.5, " length loss : ", losses[:, :, 1].sum() * 0.5)
    #     losses = losses[:, :, 0] * 0.5 + losses[:, :, 1] * 0.5
    #     losses.unsqueeze_(2)
    #     # print('pred : ', output[0])
    #     # print('target : ', target[0])

    #     if self.use_target_weight and target_weight.sum() > 0:
    #         return (losses * target_weight).sum() / (target_weight.sum())
    #     else:
    #         losses.mean()

    # def forward(self, output, target, target_weight):
    #     # output [bs, 16, 361]
    #     index = target_weight > 0
    #     # output [bs * 16]
    #     index = index.view(-1)
    #     degree = F.cross_entropy(output[:, :, :361].view(-1, 361)[index, :], target[:, :, 0].view(-1).long()[index])
        
    #     length_target = target[:, :, 1].float().view(-1)[index]
    #     length_pred = output[:, :, 361].view(-1)[index]
    #     inter = torch.min(length_pred, length_target)
    #     union = length_pred + length_target - inter
    #     length = -torch.log((inter + 1.) / (union + 1.))
    #     length = length.sum() / index.sum()
    #     return degree + 0.5 + length * 0.5
    
    def forward(self, output, target, target_weight):
        # index (bs * joints)
        # print('output: ', output[0])
        # print('target: ', target[0])
        index = target_weight > 0
        index = index.view(-1)

        output = output.view(-1, 3)[index, :]
        target = target.view(-1, 3)[index, :].float()

        angle_loss = F.mse_loss(output[:, :2], target[:, :2])

        inter = torch.min(output[:, 2], target[:, 2])
        union = output[:, 2] + target[:, 2] - inter
        length_loss = -torch.log((inter + 1.) / (union + 1.))
        length_loss = length_loss.sum() / index.sum()
        # length_loss = F.mse_loss(output[:, 9], target[:, 9])
        # print('angle : ', angle_loss, " length : ", length_loss)
        
        return 0.5 * angle_loss + 0.5 * length_loss

    # def forward(self, output, target, target_weight):
    #     # index (bs * joints)
    #     index = target_weight > 0
    #     index = index.view(-1)

    #     output = output.view(-1, 10)[index, :]
    #     target = target.view(-1, 10)[index, :].float()

    #     angle_loss = F.mse_loss(output[:, :9].view(-1, 3, 3)[:, :2, :2], target[:, :9].view(-1, 3, 3)[:, :2, :2])
    #     length_loss = F.mse_loss(output[:, 9], target[:, 9])
    #     return 0.5 * angle_loss + 0.5 * length_loss
    #     # length_loss = F.mse_loss(output[:, 9], target[:, 9])
    #     # print('angle : ', angle_loss, " length : ", length_loss)
        
    #     # return 0.5 * angle_loss + 0.5 * length_loss

    #     # length = -torch.log((inter + 1.0) / ( + 1.0))-torch.log(output[:, :, 361], target[:, :, 1].float())

class JointsMSELoss(nn.Module):
    def __init__(self, use_target_weight):
        super(JointsMSELoss, self).__init__()
        self.criterion = nn.MSELoss(reduction='mean')
        self.use_target_weight = use_target_weight

    def forward(self, output, target, target_weight):
        target = target.float()
        batch_size = output.size(0)
        num_joints = output.size(1)
        heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
        heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
        loss = 0

        for idx in range(num_joints):
            heatmap_pred = heatmaps_pred[idx].squeeze()
            heatmap_gt = heatmaps_gt[idx].squeeze()
            if self.use_target_weight:
                loss += 0.5 * self.criterion(
                    heatmap_pred.mul(target_weight[:, idx]),
                    heatmap_gt.mul(target_weight[:, idx])
                )
            else:
                loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)

        return loss / num_joints


class JointsOHKMMSELoss(nn.Module):
    def __init__(self, use_target_weight, topk=8):
        super(JointsOHKMMSELoss, self).__init__()
        self.criterion = nn.MSELoss(reduction='none')
        self.use_target_weight = use_target_weight
        self.topk = topk

    def ohkm(self, loss):
        ohkm_loss = 0.
        for i in range(loss.size()[0]):
            sub_loss = loss[i]
            topk_val, topk_idx = torch.topk(
                sub_loss, k=self.topk, dim=0, sorted=False
            )
            tmp_loss = torch.gather(sub_loss, 0, topk_idx)
            ohkm_loss += torch.sum(tmp_loss) / self.topk
        ohkm_loss /= loss.size()[0]
        return ohkm_loss

    def forward(self, output, target, target_weight):
        batch_size = output.size(0)
        num_joints = output.size(1)
        heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
        heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)

        loss = []
        for idx in range(num_joints):
            heatmap_pred = heatmaps_pred[idx].squeeze()
            heatmap_gt = heatmaps_gt[idx].squeeze()
            if self.use_target_weight:
                loss.append(0.5 * self.criterion(
                    heatmap_pred.mul(target_weight[:, idx]),
                    heatmap_gt.mul(target_weight[:, idx])
                ))
            else:
                loss.append(
                    0.5 * self.criterion(heatmap_pred, heatmap_gt)
                )

        loss = [l.mean(dim=1).unsqueeze(dim=1) for l in loss]
        loss = torch.cat(loss, dim=1)

        return self.ohkm(loss)

# ------------------------------------------------------------------------------ # Copyright (c) Microsoft # Licensed under the MIT License. # Written by <NAME> (<EMAIL>) # ------------------------------------------------------------------------------ # def forward(self, output, target, target_weight): #     # output [bs, 14, 2] #     # target same #     # target_weight [bs, 14, 1] #     target = target.float() #     inter = torch.min(output, target) #     union = output + target - inter #     losses = -torch.log((inter + 1.0) / (union + 1.0)) #     # print('angle loss : ', losses[:, :, 0].sum() * 0.5, " length loss : ", losses[:, :, 1].sum() * 0.5) #     losses = losses[:, :, 0] * 0.5 + losses[:, :, 1] * 0.5 #     losses.unsqueeze_(2) #     # print('pred : ', output[0]) #     # print('target : ', target[0]) #     if self.use_target_weight and target_weight.sum() > 0: #         return (losses * target_weight).sum() / (target_weight.sum()) #     else: #         losses.mean() # def forward(self, output, target, target_weight): #     # output [bs, 16, 361] #     index = target_weight > 0 #     # output [bs * 16] #     index = index.view(-1) #     degree = F.cross_entropy(output[:, :, :361].view(-1, 361)[index, :], target[:, :, 0].view(-1).long()[index]) #     length_target = target[:, :, 1].float().view(-1)[index] #     length_pred = output[:, :, 361].view(-1)[index] #     inter = torch.min(length_pred, length_target) #     union = length_pred + length_target - inter #     length = -torch.log((inter + 1.) / (union + 1.)) #     length = length.sum() / index.sum() #     return degree + 0.5 + length * 0.5 # index (bs * joints) # print('output: ', output[0]) # print('target: ', target[0]) # length_loss = F.mse_loss(output[:, 9], target[:, 9]) # print('angle : ', angle_loss, " length : ", length_loss) # def forward(self, output, target, target_weight): #     # index (bs * joints) #     index = target_weight > 0 #     index = index.view(-1) #     output = output.view(-1, 10)[index, :] #     target = target.view(-1, 10)[index, :].float() #     angle_loss = F.mse_loss(output[:, :9].view(-1, 3, 3)[:, :2, :2], target[:, :9].view(-1, 3, 3)[:, :2, :2]) #     length_loss = F.mse_loss(output[:, 9], target[:, 9]) #     return 0.5 * angle_loss + 0.5 * length_loss #     # length_loss = F.mse_loss(output[:, 9], target[:, 9]) #     # print('angle : ', angle_loss, " length : ", length_loss) #     # return 0.5 * angle_loss + 0.5 * length_loss #     # length = -torch.log((inter + 1.0) / ( + 1.0))-torch.log(output[:, :, 361], target[:, :, 1].float())