import torch.nn as nn
import torch.nn.functional as F
from darknet2any.tool.torch_utils import *


def build_targets(pred_boxes, target, anchors, num_anchors, num_classes, nH, nW, noobject_scale, object_scale,
          sil_thresh, seen):
  nB = target.size(0)
  nA = num_anchors
  nC = num_classes
  anchor_step = len(anchors) / num_anchors
  conf_mask = torch.ones(nB, nA, nH, nW) * noobject_scale
  coord_mask = torch.zeros(nB, nA, nH, nW)
  cls_mask = torch.zeros(nB, nA, nH, nW)
  tx = torch.zeros(nB, nA, nH, nW)
  ty = torch.zeros(nB, nA, nH, nW)
  tw = torch.zeros(nB, nA, nH, nW)
  th = torch.zeros(nB, nA, nH, nW)
  tconf = torch.zeros(nB, nA, nH, nW)
  tcls = torch.zeros(nB, nA, nH, nW)

  nAnchors = nA * nH * nW
  nPixels = nH * nW
  for b in range(nB):
    cur_pred_boxes = pred_boxes[b * nAnchors:(b + 1) * nAnchors].t()
    cur_ious = torch.zeros(nAnchors)
    for t in range(50):
      if target[b][t * 5 + 1] == 0:
        break
      gx = target[b][t * 5 + 1] * nW
      gy = target[b][t * 5 + 2] * nH
      gw = target[b][t * 5 + 3] * nW
      gh = target[b][t * 5 + 4] * nH
      cur_gt_boxes = torch.FloatTensor([gx, gy, gw, gh]).repeat(nAnchors, 1).t()
      cur_ious = torch.max(cur_ious, bbox_ious(cur_pred_boxes, cur_gt_boxes, x1y1x2y2=False))
    conf_mask[b][cur_ious > sil_thresh] = 0
  if seen < 12800:
    if anchor_step == 4:
      tx = torch.FloatTensor(anchors).view(nA, anchor_step).index_select(1, torch.LongTensor([2])).view(1, nA, 1,
                                                        1).repeat(
        nB, 1, nH, nW)
      ty = torch.FloatTensor(anchors).view(num_anchors, anchor_step).index_select(1, torch.LongTensor([2])).view(
        1, nA, 1, 1).repeat(nB, 1, nH, nW)
    else:
      tx.fill_(0.5)
      ty.fill_(0.5)
    tw.zero_()
    th.zero_()
    coord_mask.fill_(1)

  nGT = 0
  nCorrect = 0
  for b in range(nB):
    for t in range(50):
      if target[b][t * 5 + 1] == 0:
        break
      nGT = nGT + 1
      best_iou = 0.0
      best_n = -1
      min_dist = 10000
      gx = target[b][t * 5 + 1] * nW
      gy = target[b][t * 5 + 2] * nH
      gi = int(gx)
      gj = int(gy)
      gw = target[b][t * 5 + 3] * nW
      gh = target[b][t * 5 + 4] * nH
      gt_box = [0, 0, gw, gh]
      for n in range(nA):
        aw = anchors[anchor_step * n]
        ah = anchors[anchor_step * n + 1]
        anchor_box = [0, 0, aw, ah]
        iou = bbox_iou(anchor_box, gt_box, x1y1x2y2=False)
        if anchor_step == 4:
          ax = anchors[anchor_step * n + 2]
          ay = anchors[anchor_step * n + 3]
          dist = pow(((gi + ax) - gx), 2) + pow(((gj + ay) - gy), 2)
        if iou > best_iou:
          best_iou = iou
          best_n = n
        elif anchor_step == 4 and iou == best_iou and dist < min_dist:
          best_iou = iou
          best_n = n
          min_dist = dist

      gt_box = [gx, gy, gw, gh]
      pred_box = pred_boxes[b * nAnchors + best_n * nPixels + gj * nW + gi]

      coord_mask[b][best_n][gj][gi] = 1
      cls_mask[b][best_n][gj][gi] = 1
      conf_mask[b][best_n][gj][gi] = object_scale
      tx[b][best_n][gj][gi] = target[b][t * 5 + 1] * nW - gi
      ty[b][best_n][gj][gi] = target[b][t * 5 + 2] * nH - gj
      tw[b][best_n][gj][gi] = math.log(gw / anchors[anchor_step * best_n])
      th[b][best_n][gj][gi] = math.log(gh / anchors[anchor_step * best_n + 1])
      iou = bbox_iou(gt_box, pred_box, x1y1x2y2=False)  # best_iou
      tconf[b][best_n][gj][gi] = iou
      tcls[b][best_n][gj][gi] = target[b][t * 5]
      if iou > 0.5:
        nCorrect = nCorrect + 1

  return nGT, nCorrect, coord_mask, conf_mask, cls_mask, tx, ty, tw, th, tconf, tcls


class RegionLoss(nn.Module):
  def __init__(self, num_classes=0, anchors=[], num_anchors=1):
    super(RegionLoss, self).__init__()
    self.num_classes = num_classes
    self.anchors = anchors
    self.num_anchors = num_anchors
    self.anchor_step = len(anchors) / num_anchors
    self.coord_scale = 1
    self.noobject_scale = 1
    self.object_scale = 5
    self.class_scale = 1
    self.thresh = 0.6
    self.seen = 0

  def forward(self, output, target):
    # output : BxAs*(4+1+num_classes)*H*W
    t0 = time.time()
    nB = output.data.size(0)
    nA = self.num_anchors
    nC = self.num_classes
    nH = output.data.size(2)
    nW = output.data.size(3)

    output = output.view(nB, nA, (5 + nC), nH, nW)
    x = F.sigmoid(output.index_select(2, Variable(torch.cuda.LongTensor([0]))).view(nB, nA, nH, nW))
    y = F.sigmoid(output.index_select(2, Variable(torch.cuda.LongTensor([1]))).view(nB, nA, nH, nW))
    w = output.index_select(2, Variable(torch.cuda.LongTensor([2]))).view(nB, nA, nH, nW)
    h = output.index_select(2, Variable(torch.cuda.LongTensor([3]))).view(nB, nA, nH, nW)
    conf = F.sigmoid(output.index_select(2, Variable(torch.cuda.LongTensor([4]))).view(nB, nA, nH, nW))
    cls = output.index_select(2, Variable(torch.linspace(5, 5 + nC - 1, nC).long().cuda()))
    cls = cls.view(nB * nA, nC, nH * nW).transpose(1, 2).contiguous().view(nB * nA * nH * nW, nC)
    t1 = time.time()

    pred_boxes = torch.cuda.FloatTensor(4, nB * nA * nH * nW)
    grid_x = torch.linspace(0, nW - 1, nW).repeat(nH, 1).repeat(nB * nA, 1, 1).view(nB * nA * nH * nW).cuda()
    grid_y = torch.linspace(0, nH - 1, nH).repeat(nW, 1).t().repeat(nB * nA, 1, 1).view(nB * nA * nH * nW).cuda()
    anchor_w = torch.Tensor(self.anchors).view(nA, self.anchor_step).index_select(1, torch.LongTensor([0])).cuda()
    anchor_h = torch.Tensor(self.anchors).view(nA, self.anchor_step).index_select(1, torch.LongTensor([1])).cuda()
    anchor_w = anchor_w.repeat(nB, 1).repeat(1, 1, nH * nW).view(nB * nA * nH * nW)
    anchor_h = anchor_h.repeat(nB, 1).repeat(1, 1, nH * nW).view(nB * nA * nH * nW)
    pred_boxes[0] = x.data + grid_x
    pred_boxes[1] = y.data + grid_y
    pred_boxes[2] = torch.exp(w.data) * anchor_w
    pred_boxes[3] = torch.exp(h.data) * anchor_h
    pred_boxes = convert2cpu(pred_boxes.transpose(0, 1).contiguous().view(-1, 4))
    t2 = time.time()

    nGT, nCorrect, coord_mask, conf_mask, cls_mask, tx, ty, tw, th, tconf, tcls = build_targets(pred_boxes,
                                                  target.data,
                                                  self.anchors, nA,
                                                  nC, \
                                                  nH, nW,
                                                  self.noobject_scale,
                                                  self.object_scale,
                                                  self.thresh,
                                                  self.seen)
    cls_mask = (cls_mask == 1)
    nProposals = int((conf > 0.25).sum().data[0])

    tx = Variable(tx.cuda())
    ty = Variable(ty.cuda())
    tw = Variable(tw.cuda())
    th = Variable(th.cuda())
    tconf = Variable(tconf.cuda())
    tcls = Variable(tcls.view(-1)[cls_mask].long().cuda())

    coord_mask = Variable(coord_mask.cuda())
    conf_mask = Variable(conf_mask.cuda().sqrt())
    cls_mask = Variable(cls_mask.view(-1, 1).repeat(1, nC).cuda())
    cls = cls[cls_mask].view(-1, nC)

    t3 = time.time()

    loss_x = self.coord_scale * nn.MSELoss(reduction='sum')(x * coord_mask, tx * coord_mask) / 2.0
    loss_y = self.coord_scale * nn.MSELoss(reduction='sum')(y * coord_mask, ty * coord_mask) / 2.0
    loss_w = self.coord_scale * nn.MSELoss(reduction='sum')(w * coord_mask, tw * coord_mask) / 2.0
    loss_h = self.coord_scale * nn.MSELoss(reduction='sum')(h * coord_mask, th * coord_mask) / 2.0
    loss_conf = nn.MSELoss(reduction='sum')(conf * conf_mask, tconf * conf_mask) / 2.0
    loss_cls = self.class_scale * nn.CrossEntropyLoss(reduction='sum')(cls, tcls)
    loss = loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls
    t4 = time.time()
    if False:
      print('-----------------------------------')
      print('    activation : %f' % (t1 - t0))
      print(' create pred_boxes : %f' % (t2 - t1))
      print('   build targets : %f' % (t3 - t2))
      print('     create loss : %f' % (t4 - t3))
      print('       total : %f' % (t4 - t0))
    print('%d: nGT %d, recall %d, proposals %d, loss: x %f, y %f, w %f, h %f, conf %f, cls %f, total %f' % (
    self.seen, nGT, nCorrect, nProposals, loss_x.data[0], loss_y.data[0], loss_w.data[0], loss_h.data[0],
    loss_conf.data[0], loss_cls.data[0], loss.data[0]))
    return loss
