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

from cvpods.layers import ShapeSpec
from cvpods.utils import get_event_storage, comm


def weighted_cross_entropy(pred, label, weight, avg_factor=None,
                           reduce=True):
    if avg_factor is None:
        avg_factor = max(torch.sum(weight > 0).float().item(), 1.)
    raw = F.cross_entropy(pred, label, reduction='none')
    if reduce:
        return torch.sum(raw * weight)[None] / avg_factor
    else:
        return raw * weight / avg_factor

def accuracy(pred, target, topk=1):
    if isinstance(topk, int):
        topk = (topk, )
        return_single = True
    else:
        return_single = False
    maxk = max(topk)
    _, pred_label = pred.topk(maxk, 1, True, True)
    pred_label = pred_label.t()
    correct = pred_label.eq(target.view(1, -1).expand_as(pred_label))

    res = []
    for k in topk:
        correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
        res.append(correct_k.mul_(100.0 / pred.size(0)))
    return res[0] if return_single else res

class Trackhead(nn.Module):
    """
        Tracking head, predict tracking features and match with reference objects
        Using dynamic option to deal with different number of objects in different images.
        A non-match entry is add to the reference objects with all-zero features.
        Object matched with the non-match entry is condisidered as a new object.
    """

    def __init__(self, cfg, input_shape: ShapeSpec):

        super(Trackhead, self).__init__()

        self.avg_pool_on = cfg.MODEL.ROI_TRACK_HEAD.AVG_POOL_ON
        self.num_fcs = cfg.MODEL.ROI_TRACK_HEAD.NUM_FC
        self.fc_channels = cfg.MODEL.ROI_TRACK_HEAD.FC_DIM
        self.roi_feature_size = cfg.MODEL.ROI_TRACK_HEAD.POOLER_RESOLUTION
        self.match_coeff = cfg.MODEL.ROI_TRACK_HEAD.MATCH_COEFF
        self.bbox_dummy_iou = cfg.MODEL.ROI_TRACK_HEAD.BBOX_DUMMY_IOU
        self.dynamic = cfg.MODEL.ROI_TRACK_HEAD.DYNAMIC

        if self.avg_pool_on:
            self.avg_pool = nn.AvgPool2d(self.roi_feature_size)
            in_channels = input_shape.channels
        else:
            in_channels = input_shape.channels * self.roi_feature_size ** 2

        self.fcs = nn.ModuleList()
        for i in range(self.num_fcs):
            in_channels = in_channels if i == 0 else self.fc_channels
            fc = nn.Linear(in_channels, self.fc_channels)
            self.fcs.append(fc)
        self.relu = nn.ReLU()

        for fc in self.fcs:
            nn.init.normal_(fc.weight, 0, 0.01)
            nn.init.constant_(fc.bias, 0)

    def compute_comp_scores(self, match_ll, bbox_scores, bbox_ious, label_delta, add_bbox_dummy=False):
        """
        Compute comprehensize matching score based on match likelihood, bbox_confidence and ious

        Args:
            match_ll: Match scores generated by self.forward, shape=[m, n+1], m: number of bbox in current frame
                    n: number of bbox in previous frame.
            bbox_scores: Detection scores, shape=[m, ]
            bbox_ious: Ious between current frame bbox and previous frame bbox, shape = [m, n]
            label_delta: Kronecker delta function shape=[m, n]

        Returns:
            The same meaning of match_ll, shape=[m, n+1]
        """
        if self.match_coeff is None:
            return match_ll

        if add_bbox_dummy:
            bbox_iou_dummy = torch.ones(bbox_ious.size(0), 1, device=torch.cuda.current_device()) * self.bbox_dummy_iou
            bbox_ious = torch.cat([bbox_iou_dummy, bbox_ious], dim=1)
            label_dummy = torch.ones(bbox_ious.size(0), 1, device=torch.cuda.current_device())
            label_delta = torch.cat((label_dummy, label_delta), dim=1)

        assert len(self.match_coeff) == 3
        alpha, beta, gamma = self.match_coeff
        return match_ll + alpha*torch.log(bbox_scores) + beta * bbox_ious + gamma * label_delta

    def forward(self, x, ref_x, x_n, ref_x_n):
        """
        Here we compute a correlation matrix of x and ref_x. We also add a all 0 column
        denote no matching.
        Args:
            x(Tensor): Grouped bbox features generated by roi align op of current images
                        shape = [M, C, Roi_size, Roi_size]
            ref_x(Tensor): Grouped bbox features generated by roi align op of reference
                        images. shape=[N, C, Roi_size, Roi_size]
            x_n(list(int)): Numbers of proposals in the current(query) images in the
                        mini-batch. len=Number of images
            ref_x_n(list(int)): Number of ground truth bbox in the reference images in
                        the mini-batch. len=Number of images

        Returns:
            list[Tensor]:
                Each tensor is the match score in one image. Element's shape=[m, n+1],
                m is the number of bbox features in query images, n is the number of
                bbox features in reference images.
        """

        assert len(x_n) == len(ref_x_n)
        if self.avg_pool_on:
            x = self.avg_pool(x)
            ref_x = self.avg_pool(ref_x)
        x = x.view(x.size(0), -1)               # [M, C*Roi_size*Roi_size]
        ref_x = ref_x.view(ref_x.size(0), -1)   # [N, C*Roi_size*Roi_size]
        for idx, fc in enumerate(self.fcs):
            x = fc(x)
            ref_x = fc(ref_x)
            if idx < len(self.fcs) - 1:
                x = self.relu(x)
                ref_x = self.relu(ref_x)

        x_split = torch.split(x, x_n, dim=0)            # list of tensor (shape=[m, 1024])
        ref_x_split = torch.split(ref_x, ref_x_n, dim=0)        # list of tensor (shape=[n, 1024])
        prods = []
        for x_feat, ref_x_feat in zip(x_split, ref_x_split):
            prod = torch.mm(x_feat, torch.transpose(ref_x_feat, 0, 1))          # [m, n]
            prods.append(prod)                          # list of tensor (shape=[m,n])

        if self.dynamic:
            match_score = []
            for prod in prods:
                m = prod.size(0)          # number of query bboxes in a current image
                dummy = torch.zeros(m, 1, device=torch.cuda.current_device())

                prod_ext = torch.cat([dummy, prod], dim=1)
                match_score.append(prod_ext)
        else:
            pass
        for match_scores in match_score:
            if not torch.isfinite(match_scores).all():
                import logging
                logger = logging.getLogger(__name__)
                logger.info("match_score: ", match_scores,
                            "x_split: ", x_split,
                            "ref_x_split: ", ref_x_split)
        return match_score

    def loss(self, match_score, ids, id_weights, reduce=True):
        """
            Args:
                match_score (list[torch.Tensor]) shape=[m,n+1], len=number_images, m is the number of features in query
                image while n is the number of features in reference image.
                ids: (list[torch.Tensor]), len=N pids of all proposals per image
                id_weights: (list[torch.Tensor]), len=N, weights of pids per image
        """
        losses = dict()
        if self.dynamic:
            n = len(match_score)                # number of images
            loss_match = 0.
            match_acc = 0.
            n_total = 0
            for score, cur_ids, cur_weights in zip(match_score, ids, id_weights):
                valid_idx = torch.nonzero(cur_weights).squeeze(0)
                if len(valid_idx.size()) == 0:
                    loss_match += score.sum() * 0
                    continue
                n_valid = valid_idx.size(0)
                n_total += n_valid
                loss_match += weighted_cross_entropy(score, cur_ids, cur_weights, reduce=reduce)
                match_acc += accuracy(torch.index_select(score, 0, valid_idx.squeeze()),
                                      torch.index_select(cur_ids, 0, valid_idx.squeeze())) * n_valid
            if isinstance(match_acc, float):
                import logging
                logger = logging.getLogger("trackdebug")
                logger.info("match_score: ", match_score,
                            "loss_match: ", loss_match,
                            "ids: ", ids,
                            "id_weights", id_weights)
            losses['loss_match'] = loss_match / n
            if n_total > 0:
                storage = get_event_storage()
                storage.put_scalar("track_head/fg_cls_accuracy", match_acc/n_total)

        else:
            pass

        return losses