import torch
import torch.nn as nn


net_architecture = [
    (32, 3, 1),
    (64, 3, 2),
    ["B", 1],
    (128, 3, 2),
    ["B", 2],
    (256, 3, 2),
    ["B", 8],
    (512, 3, 2),
    ["B", 8],
    (1024, 3, 2),
    ["B", 4],
    (512, 1, 1),
    (1024, 3, 1),
    "Scale",
    (256, 1, 1),
    "Upsample",
    (256, 1, 1),
    (512, 3, 1),
    "Scale",
    (128, 1, 1),
    "Upsample",
    (128, 1, 1),
    (256, 3, 1),
    "Scale",
]


class CNN(nn.Module):
    def __init__(self, in_channels, out_channels, batch_norm=True, **kwargs):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, bias=not batch_norm, **kwargs)
        self.bn = nn.BatchNorm2d(out_channels)
        self.leaky = nn.LeakyReLU(0.1)
        self.batch_norm = batch_norm

    def forward(self, x):
        if self.batch_norm:
            return self.leaky(self.bn(self.conv(x)))
        else:
            return self.conv(x)


class ResBlock(nn.Module):
    def __init__(self, channels, res=True, num_repeats=1):
        super().__init__()
        self.layers = nn.ModuleList()
        for repeat in range(num_repeats):
            self.layers += [
                nn.Sequential(
                    CNN(channels, channels // 2, kernel_size=1),
                    CNN(channels // 2, channels, kernel_size=3, padding=1),
                )
            ]

        self.use_residual = res
        self.num_repeats = num_repeats

    def forward(self, x):
        for layer in self.layers:
            if self.use_residual:
                x = x + layer(x)
            else:
                x = layer(x)

        return x


class ScalePrediction(nn.Module):
    def __init__(self, in_channels, no_classes):
        super().__init__()
        self.pred = nn.Sequential(
            CNN(in_channels, 2 * in_channels, kernel_size=3, padding=1),
            CNN(
                2 * in_channels, (no_classes + 5) * 3, batch_norm=False, kernel_size=1
            ),
        )
        self.num_classes = no_classes

    def forward(self, x):
        return (
            self.pred(x)
            .reshape(x.shape[0], 3, self.num_classes + 5, x.shape[2], x.shape[3])
            .permute(0, 1, 3, 4, 2)
        )


class YoloBlock(nn.Module):
    def __init__(self, in_channels=3, no_classes=9):
        super().__init__()
        self.in_channels = in_channels
        self.no_classes = no_classes
        self.layers = self.craft_layers()

    def forward(self, x):
        outputs = []  # for each scale
        route_connections = []
        for layer in self.layers:
            if isinstance(layer, ScalePrediction):
                outputs.append(layer(x))
                continue

            x = layer(x)

            if isinstance(layer, ResBlock) and layer.num_repeats == 8:
                route_connections.append(x)

            elif isinstance(layer, nn.Upsample):
                x = torch.cat([x, route_connections[-1]], dim=1)
                route_connections.pop()

        return outputs

    def craft_layers(self):
        layers = nn.ModuleList()
        in_channels = self.in_channels

        for module in net_architecture:
            if isinstance(module, tuple):
                out_channels, kernel_size, stride = module
                layers.append(
                    CNN(
                        in_channels,
                        out_channels,
                        kernel_size=kernel_size,
                        stride=stride,
                        padding=1 if kernel_size == 3 else 0,
                    )
                )
                in_channels = out_channels
            elif isinstance(module, str):
                if module == "Scale":
                    layers += [
                        ResBlock(in_channels, res=False, num_repeats=1),
                        CNN(in_channels, in_channels // 2, kernel_size=1),
                        ScalePrediction(in_channels // 2, no_classes=self.no_classes),
                    ]
                    in_channels = in_channels // 2

                elif module == "Upsample":
                    layers.append(nn.Upsample(scale_factor=2), )
                    in_channels = in_channels * 3

            elif isinstance(module, list):
                num_repeats = module[1]
                layers.append(ResBlock(in_channels, num_repeats=num_repeats,))


        return layers
