text: # Copyright 2019 Xilinx Inc.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  
  # encoding: utf-8
  """
  @author:  liaoxingyu
  @contact: xyliao1993@qq.com 
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from __future__ import unicode_literals
  
  import torch
  from torch import nn
  import torch.nn.functional as F
  from ipdb import set_trace
  
  def normalize(x, axis=-1):
      """Normalizing to unit length along the specified dimension.
      Args:
        x: pytorch Variable
      Returns:
        x: pytorch Variable, same shape as input
      """
      x = 1. * x / (torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12)
      return x
  
  
  def euclidean_dist(x, y):
      """
      Args:
        x: pytorch Variable, with shape [m, d]
        y: pytorch Variable, with shape [n, d]
      Returns:
        dist: pytorch Variable, with shape [m, n]
      """
      m, n = x.size(0), y.size(0)
      xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)
      yy = torch.pow(y, 2).sum(1, keepdim=True).expand(n, m).t()
      dist = xx + yy
      dist.addmm_(1, -2, x, y.t())
      dist = dist.clamp(min=1e-12).sqrt()  # for numerical stability
      return dist
  
  
  def hard_example_mining(dist_mat, labels, return_inds=False):
      """For each anchor, find the hardest positive and negative sample.
      Args:
        dist_mat: pytorch Variable, pair wise distance between samples, shape [N, N]
        labels: pytorch LongTensor, with shape [N]
        return_inds: whether to return the indices. Save time if `False`(?)
      Returns:
        dist_ap: pytorch Variable, distance(anchor, positive); shape [N]
        dist_an: pytorch Variable, distance(anchor, negative); shape [N]
        p_inds: pytorch LongTensor, with shape [N];
          indices of selected hard positive samples; 0 <= p_inds[i] <= N - 1
        n_inds: pytorch LongTensor, with shape [N];
          indices of selected hard negative samples; 0 <= n_inds[i] <= N - 1
      NOTE: Only consider the case in which all labels have same num of samples,
        thus we can cope with all anchors in parallel.
      """
  
      assert len(dist_mat.size()) == 2
      assert dist_mat.size(0) == dist_mat.size(1)
      N = dist_mat.size(0)
  
      # shape [N, N]
      is_pos = labels.expand(N, N).eq(labels.expand(N, N).t())
      is_neg = labels.expand(N, N).ne(labels.expand(N, N).t())
  
      temp1 = dist_mat.clone()
      temp1[is_neg]=0
      temp2 = dist_mat.clone()
      temp2[is_pos]=0
      dist_ap, relative_p_inds = torch.max(temp1, 1, keepdim=True)
      dist_an, relative_n_inds = torch.min(temp2, 1, keepdim=True)
      dist_ap = dist_ap.squeeze(1)
      dist_an = dist_an.squeeze(1)
  
      if return_inds:
          # shape [N, N]
          ind = (labels.new().resize_as_(labels)
                 .copy_(torch.arange(0, N).long())
                 .unsqueeze(0).expand(N, N))
          # shape [N, 1]
          p_inds = torch.gather(
              ind[is_pos].contiguous().view(N, -1), 1, relative_p_inds.data)
          n_inds = torch.gather(
              ind[is_neg].contiguous().view(N, -1), 1, relative_n_inds.data)
          # shape [N]
          p_inds = p_inds.squeeze(1)
          n_inds = n_inds.squeeze(1)
          return dist_ap, dist_an, p_inds, n_inds
  
      return dist_ap, dist_an
  
  
  class TripletLoss(object):
      """Modified from Tong Xiao's open-reid (https://github.com/Cysu/open-reid).
      Related Triplet Loss theory can be found in paper 'In Defense of the Triplet
      Loss for Person Re-Identification'."""
  
      def __init__(self, margin=None):
          self.margin = margin
          if margin is not None:
              self.ranking_loss = nn.MarginRankingLoss(margin=margin)
          else:
              self.ranking_loss = nn.SoftMarginLoss()
  
      def __call__(self, global_feat, labels, normalize_feature=True):
          if normalize_feature:
              global_feat = normalize(global_feat, axis=-1)
          dist_mat = euclidean_dist(global_feat, global_feat)
          dist_ap, dist_an = hard_example_mining(
              dist_mat, labels)
          y = dist_an.new().resize_as_(dist_an).fill_(1)
          if self.margin is not None:
              loss = self.ranking_loss(dist_an, dist_ap, y)
          else:
              loss = self.ranking_loss(dist_an - dist_ap, y)
          return loss, dist_ap, dist_an
  
  class CirclePairwiseLoss(object):
      def __init__(self):
          self.m = 0.45
          self.s = 32
  
      def __call__(self, global_features, targets):
          global_features = normalize(global_features, axis=-1)
  
          sim_mat = torch.matmul(global_features, global_features.t())
  
          N = sim_mat.size(0)
          is_pos = targets.expand(N, N).eq(targets.expand(N, N).t()).float() - torch.eye(N).to(sim_mat.device)
  #        is_pos = is_pos.bool()
          is_pos = is_pos.type(torch.uint8)
          is_neg = targets.expand(N, N).ne(targets.expand(N, N).t())
  
          total_loss=0
          for i in range(N):
              s_p = sim_mat[i][is_pos[i]]
              s_n = sim_mat[i][is_neg[i]]
      
              alpha_p = F.relu(-s_p.detach() + 1 + self.m)
              alpha_n = F.relu(s_n.detach() + self.m)
              delta_p = 1 - self.m
              delta_n = self.m
      
              logit_p = - self.s * alpha_p * (s_p - delta_p)
              logit_n = self.s * alpha_n * (s_n - delta_n)
      
              loss = F.softplus(torch.logsumexp(logit_p, dim=0) + torch.logsumexp(logit_n, dim=0))
              total_loss += loss
  
          return  total_loss/N
  
  
  '''
  add by wjk. Need to complete.
  '''
  class PairwiseLoss(object):
      def __call__(self, global_feat, labels, normalize_feature=True):
          if normalize_feature:
              global_feat = normalize(global_feat, axis=-1)
          
   
      
  class CrossEntropyLabelSmooth(nn.Module):
      """Cross entropy loss with label smoothing regularizer.
      Reference:
      Szegedy et al. Rethinking the Inception Architecture for Computer Vision. CVPR 2016.
      Equation: y = (1 - epsilon) * y + epsilon / K.
      Args:
          num_classes (int): number of classes.
          epsilon (float): weight.
      """
  
      def __init__(self, num_classes, epsilon=0.1, use_gpu=True):
          super(CrossEntropyLabelSmooth, self).__init__()
          self.num_classes = num_classes
          self.epsilon = epsilon
          self.use_gpu = use_gpu
          self.logsoftmax = nn.LogSoftmax(dim=1)
  
      def forward(self, inputs, targets):
          """
          Args:
              inputs: prediction matrix (before softmax) with shape (batch_size, num_classes)
              targets: ground truth labels with shape (num_classes)
          """
          log_probs = self.logsoftmax(inputs)
          targets = torch.zeros(log_probs.size()).scatter_(1, targets.unsqueeze(1).cpu(), 1)
          if self.use_gpu: targets = targets.cuda()
          targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes
          loss = (- targets * log_probs).mean(0).sum()
          return loss
  
  
  