import copy
import random

import numpy as np
import torch
import torch.nn.functional as F
from torch.autograd import Variable

from ...util import box_transform_inv_numpy, clip_boxes_numpy, box_transform_numpy, convert_xyxyzz
from ....layers.overlap_3d import overlap_3d


class BoxModule(object):
    def __init__(self, cfg):
        self.cfg = cfg
        self.pre_nms_score_threshold = cfg.MODEL.ROI_BOX_HEAD.TEST_PRE_NMS_SCORE_THRESH
        self.num_class = cfg.MODEL.ROI_BOX_HEAD.NUM_CLASS
        self.box_reg_weight = cfg.MODEL.BBOX_REG_WEIGHT
        self.batch_size = cfg.MODEL.ROI_BOX_HEAD.ROI_BATCH_SIZE
        self.box_train_bg_thresh_high = cfg.MODEL.ROI_BOX_HEAD.BG_THRESH_HIGH
        self.box_train_fg_thresh_low = cfg.MODEL.ROI_BOX_HEAD.FG_THRESH_LOW
        self.fg_fraction = cfg.MODEL.ROI_BOX_HEAD.FG_FRACTION

    def loss(self, logits, deltas, labels, targets, deltas_sigma=1.0):
        batch_size, num_class = logits.size(0), logits.size(1)

        # Weighted cross entropy for imbalance class distribution
        weight = torch.ones(num_class).cuda()
        total = len(labels)
        for i in range(num_class):
            num_pos = float((labels == i).sum())
            num_pos = max(num_pos, 1)
            weight[i] = total / num_pos

        weight = weight / weight.sum()
        rcnn_cls_loss = F.cross_entropy(logits, labels, weight=weight, size_average=True)

        # If multi-class classification, compute the confusion metric to understand the mistakes
        confusion_matrix = np.zeros((num_class, num_class))
        probs = F.softmax(logits, dim=1)
        v, cat = torch.max(probs, dim=1)
        for i in labels.nonzero():
            i = i.item()
            confusion_matrix[labels.long().detach()[i].item()][cat[i].detach().item()] += 1

        # one hot encode
        select = Variable(torch.zeros((batch_size, num_class))).cuda()
        select.scatter_(1, labels.view(-1, 1), 1)
        select[:, 0] = 0
        select = select.view(batch_size, num_class, 1).expand((batch_size, num_class, 6)).contiguous().bool()

        deltas = deltas.view(batch_size, num_class, 6)
        deltas = deltas[select].view(-1, 6)

        rcnn_reg_loss = 0
        reg_losses = []
        for i in range(6):
            l = F.smooth_l1_loss(deltas[:, i], targets[:, i])
            rcnn_reg_loss += l
            reg_losses.append(l.data.item())

        if torch.isnan(rcnn_reg_loss.detach().cpu()):
            rcnn_reg_loss.data = torch.tensor(0.).cuda()
            reg_losses = [Variable(torch.cuda.FloatTensor(1).zero_()) for i in range(6)]

        return rcnn_cls_loss, rcnn_reg_loss, [reg_losses[0], reg_losses[1], reg_losses[2],
                                              reg_losses[3], reg_losses[4], reg_losses[5], confusion_matrix]

    def get_probability(self, inputs, proposals, logits, deltas):
        if logits is not None:
            probs = F.softmax(logits).cpu().data.numpy()
            deltas = deltas.cpu().data.numpy().reshape(-1, self.num_class, 6)
            proposals = proposals.cpu().data.numpy()
        else:
            raise NotImplementedError
        for j in range(1, self.num_class):  # skip background
            idx = np.where(probs[:, j] > self.pre_nms_score_threshold)[0]
            if len(idx) > 0:
                p = probs[idx, j].reshape(-1, 1)
                d = deltas[idx, j]
                box = box_transform_inv_numpy(proposals[idx, 2:8], d, self.box_reg_weight)
                box = clip_boxes_numpy(box, inputs.shape[2:])
                js = np.expand_dims(np.array([j] * len(p)), axis=-1)
                output = np.concatenate((p, box, js), 1)

        return torch.from_numpy(output).cuda().float()

    def make_one_box_target(self, proposal, truth_bbox):
        sampled_proposal = torch.zeros((0, 8)).float().cuda()
        sampled_label = torch.zeros((0)).long().cuda()
        sampled_assign = np.zeros((0, 1), dtype=np.int32) - 1
        sampled_target = torch.zeros((0, 6)).float().cuda()

        bboxes_label = truth_bbox.get_field("bboxes_label")
        truth_bbox = truth_bbox.data
        truth_bbox = truth_bbox[bboxes_label > -1]

        # Even if there is no ground truth box in this batch
        if len(proposal) == 0:
            # print('[RCNN] No proposal')
            return sampled_proposal, sampled_label, sampled_assign, sampled_target

        if len(truth_bbox) == 0:
            # print('[RCNN] No truth_bbox')
            num_bg = min(len(proposal), self.batch_size)
            bg_length = len(proposal)
            bg_index = np.arange(len(proposal))
            bg_index = bg_index[
                np.random.choice(bg_length, size=num_bg, replace=bg_length < num_bg)
            ]
            sampled_proposal = proposal[bg_index]
            sampled_proposal = torch.from_numpy(sampled_proposal).cuda()
            sampled_label = torch.zeros((num_bg)).long().cuda()

            return sampled_proposal, sampled_label, sampled_assign, sampled_target

        num_proposal = len(proposal)
        box = proposal[:, 2:8]

        # Determine positive or negative purely based on threshold
        # Since the GT box has been added to proposal, it is gauranteed that
        # each ground truth would have one proposal
        overlap = overlap_3d(convert_xyxyzz(torch.from_numpy(box).float().cuda()),
                             convert_xyxyzz(torch.from_numpy(truth_bbox).float().cuda())).cpu().data.numpy()
        argmax_overlap = np.argmax(overlap, 1)
        max_overlap = overlap[np.arange(num_proposal), argmax_overlap]

        fg_index = np.where(max_overlap >= self.box_train_fg_thresh_low)[0]
        bg_index = np.where(max_overlap < self.box_train_bg_thresh_high)[0]

        # sampling for class balance
        num = self.batch_size
        num_fg = int(np.round(self.fg_fraction * num))

        fg_length = len(fg_index)
        bg_length = len(bg_index)
        # print(fg_inds_length)

        # sampled_assign = argmax_overlap[fg_index]

        # Need to consider four cases, corner cases
        if fg_length > 0 and bg_length > 0:
            idx = []
            idx = random.sample(range(len(fg_index)), min(num_fg, len(fg_index)))

            fg_index = fg_index[idx]
            num_fg = len(fg_index)
            # print('[RCNN] has fgs and bgs, num_fg, num_bg, num_proposal',fg_length,bg_length,num_proposal)
            num_bg = num - num_fg
            bg_index = bg_index[
                np.random.choice(bg_length, size=num_bg, replace=bg_length < num_bg)
            ]
        elif fg_length > 0:  # no bgs
            idx = []
            idx = random.sample(range(len(fg_index)), min(num_fg, len(fg_index)))

            fg_index = fg_index[idx]
            num_fg = len(fg_index)
            num = num_fg
            num_bg = 0
            # print('[RCNN] No bgs, num_fg, num_bg, num_proposal',fg_length, bg_length,num_proposal)
            # print('truth_bbox', truth_bbox)
            # print('box',box)

        elif bg_length > 0:  # no fgs
            # print('[RCNN] No fgs, num_fg, num_bg, num_proposal',fg_length, bg_length,num_proposal)
            # print(truth_bbox)
            # print('---------------------------')
            # print(proposal)

            num_fg = 0
            num_bg = num
            bg_index = bg_index[
                np.random.choice(bg_length, size=num_bg, replace=bg_length < num_bg)
            ]
            num_fg_proposal = 0
        else:
            # no bgs and no fgs?
            # raise NotImplementedError
            # print('[RCNN] No bgs or fgs, num_fg, num_bg, num_proposal',fg_length,bg_length,num_proposal)
            # print(truth_bbox)
            # print('---------------------------')
            # print(proposal)
            num_fg = 0
            num_bg = num
            sampled_assign = np.zeros((0, 1), dtype=np.int32) - 1
            return sampled_proposal, sampled_label, sampled_assign, sampled_target
            # bg_index = np.random.choice(num_proposal, size=num_bg, replace=num_proposal<num_bg)

        assert ((num_fg + num_bg) == num)

        # selecting both fg and bg
        index = np.concatenate([fg_index, bg_index], 0)
        sampled_proposal = np.take(proposal, index, axis=0)

        # label
        sampled_assign = np.take(argmax_overlap, index)

        truth_label = np.ones(len(truth_bbox))
        sampled_label = np.take(truth_label, sampled_assign)

        sampled_label[num_fg:] = 0  # Clamp labels for the background to 0
        sampled_assign[num_fg:] = -1  # Clamp label assignments for the background to -1

        # bounding box regression terms
        if num_fg > 0:
            target_truth_bbox = truth_bbox[sampled_assign[:num_fg], :]
            # print("target_truth_bbox",target_truth_bbox)
            if len(target_truth_bbox.shape) < 2:  # one dimension lost after slicing
                target_truth_bbox = target_truth_bbox[np.newaxis, ...]
            target_box = sampled_proposal[:num_fg, :][:, 2:8]
            # print("target_box",target_box)
            sampled_target = box_transform_numpy(target_box, target_truth_bbox, self.box_reg_weight)
            sampled_target = Variable(torch.from_numpy(sampled_target)).float().cuda()

        sampled_label = Variable(torch.from_numpy(sampled_label)).long().cuda()
        sampled_proposal = Variable(torch.from_numpy(sampled_proposal)).cuda()

        # print("len sampled_proposal ",len(sampled_proposal), "len sampled_label ",len(sampled_label))

        return sampled_proposal, sampled_label, sampled_assign, sampled_target

    def make_box_target(self, proposals, truth_bboxes):
        truth_bboxes = copy.deepcopy(truth_bboxes)
        batch_size = len(truth_bboxes)
        for b in range(batch_size):
            truth_bboxes[b] = truth_bboxes[b]

        proposals = proposals.cpu().data.numpy()
        sampled_proposals = []
        sampled_labels = []
        sampled_assigns = []
        sampled_targets = []

        batch_size = len(truth_bboxes)
        for b in range(batch_size):
            truth_bbox = truth_bboxes[b]

            if len(proposals) == 0:
                proposal = np.zeros((0, 8), np.float32)
            else:
                proposal = proposals[proposals[:, 0] == b]

            # Add ground truth box to proposal, so that even if the RPN branch fails to find something,
            # we can still get classification branch to work
            proposal = select_proposal_by_diameter(proposal, 23)

            proposal = add_truth_bbox_to_proposal(proposal, b, truth_bbox)

            sampled_proposal, sampled_label, _, sampled_target = \
                self.make_one_box_target(proposal, truth_bbox)

            sampled_proposals.append(sampled_proposal)
            sampled_labels.append(sampled_label)
            sampled_targets.append(sampled_target)

        sampled_proposals = torch.cat(sampled_proposals, 0)
        sampled_labels = torch.cat(sampled_labels, 0)
        sampled_targets = torch.cat(sampled_targets, 0)

        return sampled_proposals, sampled_labels, sampled_targets


def select_proposal_by_diameter(proposal, diameter_thre):
    if len(proposal) > 0:
        proposal_idx = np.amax(proposal[:, 5:], 1) >= diameter_thre  # max(diameter_z，diameter_y,diameter_x)
        # print('proposal_idx',proposal_idx)
        proposal = proposal[proposal_idx]

    return proposal


def select_truth_bbox(truth_bbox):
    pass


def add_truth_bbox_to_proposal(proposal, b, truth_bbox, score=1):
    bboxes_label = truth_bbox.get_field("bboxes_label")
    truth_bbox = truth_bbox.data
    truth_bbox = truth_bbox[bboxes_label > 0]
    if len(truth_bbox) != 0:
        truth = np.zeros((len(truth_bbox), 8), np.float32)
        truth[:, 0] = b
        truth[:, 1] = score  # 1  #
        truth[:, 2:8] = truth_bbox
    else:
        truth = np.zeros((0, 8), np.float32)

    sampled_proposal = np.vstack([proposal, truth])
    return sampled_proposal
