import torch
import torch.nn as nn
from yolox_microbt.model.backbone.efficientnet import MBNetV2Backbone
from yolox_microbt.model.head.yolox_head import YOLOXHead
from yolox_microbt.model.loss.yolox_loss import YOLOXLoss
from yolox.utils import meshgrid


class YOLOXDetector(nn.Module):
    def __init__(self, 
                 backbone_params,
                 head_params):
        super(YOLOXDetector, self).__init__()
        self.add_module('backbone0', MBNetV2Backbone(**backbone_params))
        self.add_module('head0', YOLOXHead(**head_params))
    
    def forward(self, x):
        features = self.backbone0(x)
        outs = self.head0(features)
        return outs


class YOLOXTrainer(nn.Module):
    def __init__(self, 
                 backbone_params,
                 head_params,
                 loss_params):
        super(YOLOXTrainer, self).__init__()
        self.yolox = YOLOXDetector(backbone_params, head_params)
        self.loss = YOLOXLoss(**loss_params)
        self.decode_in_inference = True


    def forward(self, x, targets=None):
        outs = self.yolox(x)
        if self.training:
            assert targets is not None
            loss, iou_loss, conf_loss, cls_loss, l1_loss, num_fg = self.loss(
                outs, targets, x
            )
            outputs = {
                "total_loss": loss,
                "iou_loss": iou_loss,
                "l1_loss": l1_loss,
                "conf_loss": conf_loss,
                "cls_loss": cls_loss,
                "num_fg": num_fg,
            }
        else:
            if self.decode_in_inference:
                outputs = []
                for stride_i in range(len(outs[0])):
                    obj_output = outs[0][stride_i]
                    cls_output = outs[1][stride_i]
                    reg_output = outs[2][stride_i]
                    output = torch.cat([reg_output, obj_output.sigmoid(), cls_output.sigmoid()], 1)
                    outputs.append(output)
                self.hw = [x_out.shape[-2:] for x_out in outputs]
                # [batch, n_anchors_all, 85]
                outputs = torch.cat(
                    [x_out.flatten(start_dim=2) for x_out in outputs], dim=2
                ).permute(0, 2, 1)
                self.strides = self.loss.strides
                outputs = self.decode_outputs(outputs, reg_output.type())
            else:
                outputs = outs
        return outputs
    
    def decode_outputs(self, outputs, dtype):
        grids = []
        strides = []
        for (hsize, wsize), stride in zip(self.hw, self.strides):
            yv, xv = meshgrid([torch.arange(hsize), torch.arange(wsize)])
            grid = torch.stack((xv, yv), 2).view(1, -1, 2)
            grids.append(grid)
            shape = grid.shape[:2]
            strides.append(torch.full((*shape, 1), stride))

        grids = torch.cat(grids, dim=1).type(dtype)
        strides = torch.cat(strides, dim=1).type(dtype)

        outputs = torch.cat([
            (outputs[..., 0:2] + grids) * strides,
            torch.exp(outputs[..., 2:4]) * strides,
            outputs[..., 4:]
        ], dim=-1)
        return outputs