import torch
import torch.nn as nn

from .yolo_layer import *
from .yolov3_base import *

class Yolov3(Yolov3Base):
    def __init__(self, num_classes=80):
        super().__init__()
        self.backbone = Darknet([1,2,8,8,4])
        
        anchors_per_region = 3
        self.yolo_0_pre = Yolov3UpsamplePrep([512, 1024], 1024, anchors_per_region*(5+num_classes))
        self.yolo_0 = YoloLayer(anchors=[(116.,  90.), (156., 198.), (373., 326.)], stride=32, num_classes=num_classes)

        self.yolo_1_c = ConvBN(512, 256, 1)
        self.yolo_1_prep = Yolov3UpsamplePrep([256, 512], 512+256, anchors_per_region*(5+num_classes))
        self.yolo_1 = YoloLayer(anchors=[(30., 61.), (62., 45.), (59., 119.)], stride=16, num_classes=num_classes)

        self.yolo_2_c = ConvBN(256, 128, 1)
        self.yolo_2_prep = Yolov3UpsamplePrep([128, 256], 256+128, anchors_per_region*(5+num_classes))
        self.yolo_2 = YoloLayer(anchors=[(10., 13.), (16., 30.), (33., 23.)], stride=8, num_classes=num_classes)

    def get_loss_layers(self):
        return [self.yolo_0, self.yolo_1, self.yolo_2]

    def forward_yolo(self, xb):
        x, y0 = self.yolo_0_pre(xb[-1])

        x = self.yolo_1_c(x)
        x = nn.Upsample(scale_factor=2, mode='nearest')(x)
        x = torch.cat([x, xb[-2]], 1)
        x, y1 = self.yolo_1_prep(x)

        x = self.yolo_2_c(x)
        x = nn.Upsample(scale_factor=2, mode='nearest')(x)
        x = torch.cat([x, xb[-3]], 1)
        x, y2 = self.yolo_2_prep(x)
        
        return [y0, y1, y2]


###################################################################
## Backbone and helper modules

class DarknetBlock(nn.Module):
    def __init__(self, ch_in):
        super().__init__()
        ch_hid = ch_in//2
        self.conv1 = ConvBN(ch_in, ch_hid, kernel_size=1, stride=1, padding=0)
        self.conv2 = ConvBN(ch_hid, ch_in, kernel_size=3, stride=1, padding=1)

    def forward(self, x): return self.conv2(self.conv1(x)) + x


class Darknet(nn.Module):
    def __init__(self, num_blocks, start_nf=32):
        super().__init__()
        nf = start_nf
        self.base = ConvBN(3, nf, kernel_size=3, stride=1) #, padding=1)
        self.layers = []
        for i, nb in enumerate(num_blocks):
            # dn_layer = make_group_layer(nf, nb, stride=(1 if i==-1 else 2))
            dn_layer = self.make_group_layer(nf, nb, stride=2)
            self.add_module(f"darknet_{i}", dn_layer)
            self.layers.append(dn_layer)
            nf *= 2

    def make_group_layer(self, ch_in, num_blocks, stride=2):
        layers = [ConvBN(ch_in, ch_in*2, stride=stride)]
        for i in range(num_blocks): layers.append(DarknetBlock(ch_in*2))
        return nn.Sequential(*layers)

    def forward(self, x):
        y = [self.base(x)]
        for l in self.layers:
            y.append(l(y[-1]))
        return y


class Yolov3UpsamplePrep(nn.Module):
    def __init__(self, filters_list, in_filters, out_filters):
        super().__init__()
        self.branch = nn.ModuleList([
                        ConvBN(in_filters, filters_list[0], 1),
                        ConvBN(filters_list[0], filters_list[1], kernel_size=3),
                        ConvBN(filters_list[1], filters_list[0], kernel_size=1),
                        ConvBN(filters_list[0], filters_list[1], kernel_size=3),
                        ConvBN(filters_list[1], filters_list[0], kernel_size=1),])
        self.for_yolo = nn.ModuleList([
                        ConvBN(filters_list[0], filters_list[1], kernel_size=3),
                        nn.Conv2d(filters_list[1], out_filters, kernel_size=1, stride=1,
                                   padding=0, bias=True)])

    def forward(self, x):
        for m in self.branch: x = m(x)
        branch_out = x
        for m in self.for_yolo: x = m(x)
        return branch_out, x


[((32, 12, 32, 37), 'torch.cat', 'torch.cat', ({(32, 22, 32, 33): '[x, xb[-2]]', (32, 35, 32, 36): '1'}, {}), '([x, xb[-2]], 1)', False, 'import torch\n'), ((37, 12, 37, 37), 'torch.cat', 'torch.cat', ({(37, 22, 37, 33): '[x, xb[-3]]', (37, 35, 37, 36): '1'}, {}), '([x, xb[-3]], 1)', False, 'import torch\n'), ((72, 15, 72, 37), 'torch.nn.Sequential', 'nn.Sequential', ({(72, 29, 72, 36): '*layers'}, {}), '(*layers)', True, 'import torch.nn as nn\n'), ((31, 12, 31, 55), 'torch.nn.Upsample', 'nn.Upsample', (), '', True, 'import torch.nn as nn\n'), ((36, 12, 36, 55), 'torch.nn.Upsample', 'nn.Upsample', (), '', True, 'import torch.nn as nn\n'), ((92, 24, 93, 56), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n')]