# ------------------------------------------------------------------------------
# Portions of this code are from
# CornerNet (https://github.com/princeton-vl/CornerNet)
# Copyright (c) 2018, University of Michigan
# Licensed under the BSD 3-Clause License
# ------------------------------------------------------------------------------
"""
当前模块是对各种损失函数的一个封装 即根据掩码计算有效位置的损失
实际上大多数损失函数在训练过程中都没有用到也没有来得及用到
"""

import torch, math, random
import torch.nn as nn
import torch.nn.functional as F
from fvcore.nn import sigmoid_focal_loss_jit

from .base import _neg_loss, _reg_loss
from lib.models.utils.models import _tranpose_and_gather_feat, _sigmoid

class FocalLoss(nn.Module):
  '''nn.Module warpper for focal loss'''
  def __init__(self):
    super(FocalLoss, self).__init__()
    self.neg_loss = _neg_loss

  def forward(self, out, target):
    return self.neg_loss(out, target)

class RegSL1Loss(nn.Module):
  '''Regression loss for an output tensor
    Arguments:
      output (batch x dim x h x w)
      mask (batch x max_objects)
      ind (batch x max_objects)
      target (batch x max_objects x dim)
    最后得到的结果取平均值
  '''
  def __init__(self):
    super(RegSL1Loss, self).__init__()
  
  def forward(self, output, mask, ind, target):
    pred = _tranpose_and_gather_feat(output, ind)
    loss = _reg_loss(pred, target, mask)
    return loss

class RegL1Loss(nn.Module):
  """
  L1损失 最后得到的结构取平均值
  """
  def __init__(self):
    super(RegL1Loss, self).__init__()
  
  def forward(self, output, mask, ind, target):
    pred = _tranpose_and_gather_feat(output, ind)
    mask = mask.unsqueeze(2).expand_as(pred).float()
    # loss = F.l1_loss(pred * mask, target * mask, reduction='elementwise_mean')
    loss = F.l1_loss(pred * mask, target * mask, size_average=False)
    loss = loss / (mask.sum() + 1e-4)
    return loss

class NormRegL1Loss(nn.Module):
  """
  进行简单放缩处理的L1回归损失 主要求的是和原始值的相对比例
  """
  def __init__(self):
    super(NormRegL1Loss, self).__init__()
  
  def forward(self, output, mask, ind, target):
    pred = _tranpose_and_gather_feat(output, ind)
    mask = mask.unsqueeze(2).expand_as(pred).float()
    # loss = F.l1_loss(pred * mask, target * mask, reduction='elementwise_mean')
    pred = pred / (target + 1e-4)
    target = target * 0 + 1
    loss = F.l1_loss(pred * mask, target * mask, size_average=False)
    loss = loss / (mask.sum() + 1e-4)
    return loss

class RegWeightedL1Loss(nn.Module):
  def __init__(self):
    super(RegWeightedL1Loss, self).__init__()
  
  def forward(self, output, mask, ind, target):
    pred = _tranpose_and_gather_feat(output, ind)
    mask = mask.float()
    # loss = F.l1_loss(pred * mask, target * mask, reduction='elementwise_mean')
    loss = F.l1_loss(pred * mask, target * mask, size_average=False)
    loss = loss / (mask.sum() + 1e-4)
    return loss

class L1Loss(nn.Module):
  def __init__(self):
    super(L1Loss, self).__init__()
  
  def forward(self, output, mask, ind, target):
    pred = _tranpose_and_gather_feat(output, ind)
    mask = mask.unsqueeze(2).expand_as(pred).float()
    loss = F.l1_loss(pred * mask, target * mask, reduction='elementwise_mean')
    return loss

class BinRotLoss(nn.Module):
  def __init__(self):
    super(BinRotLoss, self).__init__()
  
  def forward(self, output, mask, ind, rotbin, rotres):
    pred = _tranpose_and_gather_feat(output, ind)
    loss = compute_rot_loss(pred, rotbin, rotres, mask)
    return loss

def compute_res_loss(output, target):
    return F.smooth_l1_loss(output, target, reduction='elementwise_mean')

def compute_bin_loss(output, target, mask):
    mask = mask.expand_as(output)
    output = output * mask.float()
    return F.cross_entropy(output, target, reduction='elementwise_mean')

def compute_rot_loss(output, target_bin, target_res, mask):
    # output: (B, 128, 8) [bin1_cls[0], bin1_cls[1], bin1_sin, bin1_cos, 
    #                 bin2_cls[0], bin2_cls[1], bin2_sin, bin2_cos]
    # target_bin: (B, 128, 2) [bin1_cls, bin2_cls]
    # target_res: (B, 128, 2) [bin1_res, bin2_res]
    # mask: (B, 128, 1)
    # import pdb; pdb.set_trace()
    output = output.view(-1, 8)
    target_bin = target_bin.view(-1, 2)
    target_res = target_res.view(-1, 2)
    mask = mask.view(-1, 1)
    loss_bin1 = compute_bin_loss(output[:, 0:2], target_bin[:, 0], mask)
    loss_bin2 = compute_bin_loss(output[:, 4:6], target_bin[:, 1], mask)
    loss_res = torch.zeros_like(loss_bin1)
    if target_bin[:, 0].nonzero().shape[0] > 0:
        idx1 = target_bin[:, 0].nonzero()[:, 0]
        valid_output1 = torch.index_select(output, 0, idx1.long())
        valid_target_res1 = torch.index_select(target_res, 0, idx1.long())
        loss_sin1 = compute_res_loss(
          valid_output1[:, 2], torch.sin(valid_target_res1[:, 0]))
        loss_cos1 = compute_res_loss(
          valid_output1[:, 3], torch.cos(valid_target_res1[:, 0]))
        loss_res += loss_sin1 + loss_cos1
    if target_bin[:, 1].nonzero().shape[0] > 0:
        idx2 = target_bin[:, 1].nonzero()[:, 0]
        valid_output2 = torch.index_select(output, 0, idx2.long())
        valid_target_res2 = torch.index_select(target_res, 0, idx2.long())
        loss_sin2 = compute_res_loss(
          valid_output2[:, 6], torch.sin(valid_target_res2[:, 1]))
        loss_cos2 = compute_res_loss(
          valid_output2[:, 7], torch.cos(valid_target_res2[:, 1]))
        loss_res += loss_sin2 + loss_cos2
    return loss_bin1 + loss_bin2 + loss_res

class TripletLoss(nn.Module):
    """Triplet loss with hard positive/negative mining.
    Reference:
    Hermans et al. In Defense of the Triplet Loss for Person Re-Identification. arXiv:1703.07737.
    Code imported from https://github.com/Cysu/open-reid/blob/master/reid/loss/triplet.py.
    Args:
        margin (float): margin for triplet.
    """

    def __init__(self, margin=0.3, mutual_flag=False):
        super(TripletLoss, self).__init__()
        self.margin = margin
        self.ranking_loss = nn.MarginRankingLoss(margin=margin)
        self.mutual = mutual_flag

    def forward(self, inputs, targets):
        """
        Args:
            inputs: feature matrix with shape (batch_size, feat_dim)
            targets: ground truth labels with shape (num_classes)
        """
        n = inputs.size(0)
        # inputs = 1. * inputs / (torch.norm(inputs, 2, dim=-1, keepdim=True).expand_as(inputs) + 1e-12)
        # Compute pairwise distance, replace by the official when merged
        dist = torch.pow(inputs, 2).sum(dim=1, keepdim=True).expand(n, n)
        dist = dist + dist.t()
        dist.addmm_(1, -2, inputs, inputs.t())
        dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability
        # For each anchor, find the hardest positive and negative
        mask = targets.expand(n, n).eq(targets.expand(n, n).t())
        dist_ap, dist_an = [], []
        for i in range(n):
            dist_ap.append(dist[i][mask[i]].max().unsqueeze(0))
            dist_an.append(dist[i][mask[i] == 0].min().unsqueeze(0))
        dist_ap = torch.cat(dist_ap)
        dist_an = torch.cat(dist_an)
        # Compute ranking hinge loss
        y = torch.ones_like(dist_an)
        loss = self.ranking_loss(dist_an, dist_ap, y)
        if self.mutual:
            return loss, dist
        return loss

# =====================================================
# 多目标跟踪的各个损失部分
# =====================================================
class DetLoss(nn.Module):
  """
  针对检测部分的损失
  """
  def __init__(self, opt):
    super(DetLoss, self).__init__()

    self.opt = opt
    # 对于中心点的计算是选择回归还是MSE回归还是FocalLoss分类
    self.hm_crit = torch.nn.MSELoss() if opt.mse_loss else FocalLoss()
    # 中心点相对偏移的计算
    self.offset_reg_crit = RegL1Loss() if opt.reg_loss == 'l1' else RegSL1Loss() if opt.reg_loss == 'sl1' else None
    # 对于尺寸的损失定义 一般采用ltrb的形式 dense_wh定义了sum的损失规约方式 其它默认的都是求平均值
    self.size_crit = torch.nn.L1Loss(reduction='sum') if opt.dense_wh else NormRegL1Loss() if opt.norm_wh else RegWeightedL1Loss() if opt.cat_spec_wh else self.offset_reg_crit

  def hm_loss(self, output_hm, target_hm):
    """
    计算热力图的损失
    Args:
        output_hm: 网络输出的hm部分
        target_hm: 需要对比的target部分
    Return:
        loss: hm热力图部分的损失
    """
    if not self.opt.mse_loss:
      # for focal loss
      output_hm = _sigmoid(output_hm)
    return self.hm_crit(output_hm, target_hm)
      
  def center_offset_loss(self, output_center_offset, reg_masks, reg_position_ind, target_center_offset):
      """
      计算中心点偏移量的损失
      Args:
          output_center_offset: 网络输出的中心点偏移量
          reg_masks: 需要回归的掩码 即1代表需要进行回归计算
          reg_position_ind:需要进行回归计算的位置的1d索引
          target_center_offset: 目标偏移量
      Return:
          loss: 偏移量的回归损失
      """
      if self.opt.reg_offset and self.opt.off_weight > 0:
        return self.offset_reg_crit(output_center_offset, reg_masks, reg_position_ind, target_center_offset)
      return 0.0

  def size_loss(self, output_size, reg_masks, reg_position_ind, target_size):
      """
      计算bbox大小的损失 这里转化为中心点距离边框的距离
      Args:
          output_size: 网络输出的size部分
          reg_masks: 需要回归的掩码 即1代表需要进行回归计算
          reg_position_ind:需要进行回归计算的位置的1d索引
          target_size: 目标大小
      Return:
          loss: size的回归损失
      """
      if self.opt.wh_weight > 0:
        return self.size_crit(output_size, reg_masks, reg_position_ind, target_size)
      return 0.0

  def forward(self, output, batch):
    """
    针对检测部分的损失
    """
    hm_loss, wh_loss, off_loss = 0, 0, 0
    for s in range(self.opt.num_stacks):
      output = output[s]
      # ========================================================
      # 计算hm损失
      # ========================================================
      hm_loss += self.hm_loss(output['hm'], batch['hm']) / self.opt.num_stacks
      # ========================================================
      # 计算wh损失
      # ========================================================
      wh_loss += self.size_loss(output['wh'], batch['reg_mask'], batch['ind'], batch['wh']) / self.opt.num_stacks
      # ========================================================
      # 计算offset损失
      # ========================================================
      if 'reg' in output.keys():
        off_loss += self.center_offset_loss(output['reg'], batch['reg_mask'], batch['ind'], batch['reg']) / self.opt.num_stacks
      else:
        off_loss = torch.tensor(0.0, requires_grad=False)

      # # ====================================================
      # output = output[s]
      # batch_size = batch['hm'].size()[0]
      # # 计算hm损失
      # # ========================================================
      # hm_group = torch.split(output['hm'], [batch_size, batch_size, batch_size, batch_size], dim=0)
      # hm_loss += (0.1 * self.hm_loss(hm_group[0], batch['hm']) + 
      #             0.1 * self.hm_loss(hm_group[1], batch['hm']) + 
      #             0.1 * self.hm_loss(hm_group[2], batch['hm']) + 
      #             self.hm_loss(hm_group[3], batch['hm'])) / self.opt.num_stacks
      # # ========================================================
      # # 计算wh损失
      # # ========================================================
      # wh_group = torch.split(output['wh'], [batch_size, batch_size, batch_size, batch_size], dim=0)
      # wh_loss += (0.1 * self.size_loss(wh_group[0], batch['reg_mask'], batch['ind'], batch['wh']) + 
      #             0.1 * self.size_loss(wh_group[1], batch['reg_mask'], batch['ind'], batch['wh']) + 
      #             0.1 * self.size_loss(wh_group[2], batch['reg_mask'], batch['ind'], batch['wh']) +
      #             self.size_loss(wh_group[3], batch['reg_mask'], batch['ind'], batch['wh']) ) / self.opt.num_stacks
      # # ========================================================
      # # 计算offset损失
      # # ========================================================
      # if 'reg' in output.keys():
      #   off_group = torch.split(output['reg'], [batch_size, batch_size, batch_size, batch_size], dim=0)
      #   off_loss += (0.1 * self.center_offset_loss(off_group[0], batch['reg_mask'], batch['ind'], batch['reg']) + 
      #                0.1 * self.center_offset_loss(off_group[1], batch['reg_mask'], batch['ind'], batch['reg']) + 
      #                0.1 * self.center_offset_loss(off_group[2], batch['reg_mask'], batch['ind'], batch['reg']) + 
      #                self.center_offset_loss(off_group[3], batch['reg_mask'], batch['ind'], batch['reg'])) / self.opt.num_stacks
      # else:
      #  off_loss = torch.tensor(0.0, requires_grad=False)

    det_loss = self.opt.hm_weight * hm_loss + self.opt.wh_weight * wh_loss + self.opt.off_weight * off_loss

    return det_loss, {
      "det_loss": det_loss,
      'hm_loss': hm_loss,
      'wh_loss': wh_loss,
      'off_loss': off_loss
    }

class ReIDLoss(nn.Module):
  """
  针对ReID部分的损失实现
  """
  def __init__(self, opt):
      super(ReIDLoss, self).__init__()

      self.opt = opt
      self.emb_dim = opt.reid_dim
      self.nID = opt.nID
      
      self.once_classifier = nn.Linear(self.emb_dim, self.nID)
      # 根据选择的损失类型来对分类器进行初始化
      if opt.id_loss == 'focal':
        torch.nn.init.normal_(self.classifier.weight, std=0.01)
        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        torch.nn.init.constant_(self.classifier.bias, bias_value)

      # 分类依然采用交叉熵来计算
      # 标签为-1的哪一项忽略计算 不产生reid损失 仅仅用于det部分 
      self.reid_crit = nn.CrossEntropyLoss(ignore_index=-1)
      
      # 缩放尺寸
      self.emb_scale = math.sqrt(2) * math.log(self.nID)
  
  def reid_loss(self, output, batch):
      """
      ReID损失
      """
      id_loss = 0.0
      # 根据ind信息找到对应的那个位置的特征
      if self.opt.id_weight > 0:
        # !!! 普通模式训练
        if "hd_patches" not in batch.keys():     
          # 从feat中抽取指定空间位置处的通道特征
          # batch * k * reid_dim
          id_head = _tranpose_and_gather_feat(output['id'], batch['ind'])
          # id特征 (batch * total_k) * 128
          id_head = id_head[batch['reg_mask'] > 0].contiguous()
        # !!! 高分辨率图像块增强模式训练
        else:                                     
          # id特征[(batch * total_k) * 2] * 128 包括了增强前的和增强后的 
          id_head = output['fine_id']
          # 以50%的概率进行融合
          partition_total_num = int(id_head.size()[0] * 0.5)
          low_id_head, hr_id_head = torch.split(id_head, [partition_total_num, partition_total_num], dim=0)
          # for i in range(partition_total_num):
          #     if random.random() <= 0.5:
          #       low_id_head[i] = hr_id_head[i]
          # id_head = low_id_head 
          merged_id_head = torch.stack([
            low_id_head[i] if random.random() <= 0.5 else hr_id_head[i] for i in range(partition_total_num)
          ],dim=0)
          id_head = merged_id_head
        
        # 规范化
        id_head = self.emb_scale * F.normalize(id_head)

        id_target = batch['ids'][batch['reg_mask'] > 0]
        id_output = self.once_classifier(id_head).contiguous()
        if self.opt.id_loss == 'focal':
          id_target_one_hot = id_output.new_zeros((id_head.size(0), self.nID)).scatter_(1, id_target.long().view(-1, 1), 1)
          id_loss += sigmoid_focal_loss_jit(id_output, id_target_one_hot, alpha=0.25, gamma=2.0, reduction="sum") / id_output.size(0)
        else:
          id_loss += self.reid_crit(id_output, id_target)

      return id_loss

  def forward(self, outputs, batch):

    reid_loss = 0.0
    for i in range(self.opt.num_stacks):
      output = outputs[i]
      reid_loss += self.reid_loss(output, batch) / self.opt.num_stacks

    return reid_loss, {"reid_loss": reid_loss}

class MotLoss(nn.Module):
    """
    定义的MotLoss代表了整体的MOT损失
    """
    def __init__(self, opt):
      super(MotLoss, self).__init__()

      self.opt = opt

      self.det_loss_model = DetLoss(opt)
      
      if self.opt.id_weight > 0:
        self.reid_loss_model = ReIDLoss(opt)
        # uncertainty weights
        self.s_det = nn.Parameter(-1.85 * torch.ones(1))
        self.s_id = nn.Parameter(-1.05 * torch.ones(1))

    def forward(
      self,
      output,
      batch,
    ):
      """
      计算各类损失值
      四项损失: hm_loss 回归或者分类 wh_loss bbox尺寸大小预测 off_loss 中心点的偏移量 id_loss reid分类损失
      """
      opt = self.opt
      
      det_loss, det_loss_details = self.det_loss_model(output, batch)

      if opt.id_weight > 0:
        id_loss, id_loss_details = self.reid_loss_model(output, batch)
        if opt.multi_loss == 'uncertainty':
            loss = torch.exp(-self.s_det) * det_loss + torch.exp(-self.s_id) * id_loss + (self.s_det + self.s_id)
            loss *= 0.5
        else:
            loss = det_loss + 0.1 * id_loss
        loss_stats = {"loss": loss}
        loss_stats.update(det_loss_details)
        loss_stats.update(id_loss_details)
      else:
        loss = det_loss
        loss_stats = {"loss": loss}
        loss_stats.update(det_loss_details)
      return loss, loss_stats

