import torch
import math
import numpy as np
import torch.nn as nn
from .conv_block import Compose


class YOLOXHead(nn.Module):
    def __init__(self, num_classes,
                       feature_strides=[8, 16, 32],
                       feature_in_channels=[8, 24, 80],
                       shared_channels=[64, 64, 64],
                       shared_op_list = ('conv', 'act')
                       ):
        super(YOLOXHead, self).__init__()
        self.feature_strides = feature_strides
        for f_idx, stride in enumerate(feature_strides):
            shared_layer_name = 'shared_layer_%s'%(stride)
            
            shared_op_params_list = [
                    {
                        'type': 'ConvBlock',
                        'out_channels': shared_channels[0],
                        'kernel_size': 1,
                        'padding': 0,
                        'op_list': shared_op_list,
                    },
                    {
                        'type': 'ConvBlock',
                        'out_channels': shared_channels[1],
                        'kernel_size': 3,
                        'padding': 1,
                        'groups':  shared_channels[1],
                        'op_list': shared_op_list,
                    },
                    {
                        'type': 'ConvBlock',
                        'out_channels': shared_channels[2],
                        'kernel_size': 1,
                        'padding': 0,
                        'op_list': shared_op_list,
                    },
                ]

            self.add_module(shared_layer_name, Compose(op_params_list=shared_op_params_list, in_channels=feature_in_channels[f_idx]))
            obj_layer_name = 'layer_%s_obj'%(stride)
            obj_op_params_list = [
                    {
                        'type': 'ConvBlock',
                        'out_channels': 1,
                        'kernel_size': 1,
                        'padding': 0,
                    }
                ]
            self.add_module(obj_layer_name, Compose(op_params_list=obj_op_params_list, in_channels=shared_channels[-1]))
            cls_layer_name = 'layer_%s_cls'%(stride)
            cls_op_params_list = [
                    {
                        'type': 'ConvBlock',
                        'out_channels': num_classes,
                        'kernel_size': 1,
                        'padding': 0,
                    }
                ]
            self.add_module(cls_layer_name, Compose(op_params_list=cls_op_params_list, in_channels=shared_channels[-1]))
            box_layer_name = 'layer_%s_box'%(stride)
            box_op_params_list = [
                    {
                        'type': 'ConvBlock',
                        'out_channels': 4,
                        'kernel_size': 1,
                        'padding': 0,
                    }
                ]
            self.add_module(box_layer_name, Compose(op_params_list=box_op_params_list, in_channels=shared_channels[-1]))

    def initialized_params(self, prior_prob):
        for name, param in self.named_parameters():
            nn.init.normal_(param, mean=0, std=0.01)
        for f_idx, stride in enumerate(self.feature_strides):
            obj_layer_name = 'layer_%s_obj'%(stride)
            cls_layer_name = 'layer_%s_cls'%(stride)
            for conv in getattr(self, cls_layer_name).modules():
                if isinstance(conv, nn.Conv2d):
                    b = conv.bias.view(1, -1)
                    b.data.fill_(-math.log((1 - prior_prob) / prior_prob))
                    conv.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)

            for conv in getattr(self, obj_layer_name).modules():
                if isinstance(conv, nn.Conv2d):
                    b = conv.bias.view(1, -1)
                    b.data.fill_(-math.log((1 - prior_prob) / prior_prob))
                    conv.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)

    def forward(self, x):
        outs_obj = [] 
        outs_cls = [] 
        outs_box = [] 
        feat_8 = self.shared_layer_8 (x[0])
        outs_obj.append(self.layer_8_obj(feat_8))
        outs_cls.append(self.layer_8_cls(feat_8))
        outs_box.append(self.layer_8_box(feat_8))
        feat_16 = self.shared_layer_16 (x[1])
        outs_obj.append(self.layer_16_obj(feat_16))
        outs_cls.append(self.layer_16_cls(feat_16))
        outs_box.append(self.layer_16_box(feat_16))
        feat32 = self.shared_layer_32 (x[2])
        outs_obj.append(self.layer_32_obj(feat32))
        outs_cls.append(self.layer_32_cls(feat32))
        outs_box.append(self.layer_32_box(feat32))
        return [outs_obj, outs_cls, outs_box]