import torch
import torch.nn as nn
import torch.nn.functional as F
from . import backbone_picie as backbone


class PanopticFPN(nn.Module):
    def __init__(self, args):
        super(PanopticFPN, self).__init__()
        self.backbone = backbone.__dict__[args.arch](pretrained=args.pretrain)
        if args.arch == 'vit_small':
            self.decoder = FPNDecoderViT(args)
        else:
            self.decoder = FPNDecoder(args)

    def forward(self, x, encoder_features=False, decoder_features=False):
        feats = self.backbone(x)
        dec_outs = self.decoder(feats)

        if encoder_features:
            return feats['res5'], dec_outs
        else:
            return dec_outs


class FPNDecoder(nn.Module):
    def __init__(self, args):
        super(FPNDecoder, self).__init__()
        if args.arch == 'resnet18':
            mfactor = 1
            out_dim = 128
        else:
            mfactor = 4
            out_dim = 256

        self.layer4 = nn.Conv2d(512 * mfactor // 8, out_dim, kernel_size=1, stride=1, padding=0)
        self.layer3 = nn.Conv2d(512 * mfactor // 4, out_dim, kernel_size=1, stride=1, padding=0)
        self.layer2 = nn.Conv2d(512 * mfactor // 2, out_dim, kernel_size=1, stride=1, padding=0)
        self.layer1 = nn.Conv2d(512 * mfactor, out_dim, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        o1 = self.layer1(x['res5'])
        o2 = self.upsample_add(o1, self.layer2(x['res4']))
        o3 = self.upsample_add(o2, self.layer3(x['res3']))
        o4 = self.upsample_add(o3, self.layer4(x['res2']))

        return o4

    def upsample_add(self, x, y):
        _, _, H, W = y.size()

        return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=False) + y


class FPNDecoderViT(nn.Module):
    def __init__(self, args):
        super(FPNDecoderViT, self).__init__()
        if args.arch == 'resnet18' or args.arch == 'vit_small':
            mfactor = 1
            out_dim = 128
        else:
            mfactor = 4
            out_dim = 256

        self.upsample_rate = 4

        self.layer4 = nn.Conv2d(384, out_dim, kernel_size=1, stride=1, padding=0)
        self.layer3 = nn.Conv2d(384, out_dim, kernel_size=1, stride=1, padding=0)
        self.layer2 = nn.Conv2d(384, out_dim, kernel_size=1, stride=1, padding=0)
        self.layer1 = nn.Conv2d(384, out_dim, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        o1 = self.layer1(x[3])
        o1 = F.interpolate(o1, scale_factor=4, mode='bilinear', align_corners=False)
        o2 = self.upsample_add(o1, self.layer2(x[2]))
        o3 = self.upsample_add(o2, self.layer3(x[1]))
        o4 = self.upsample_add(o3, self.layer4(x[0]))

        return o4

    def upsample_add(self, x, y):
        return F.interpolate(y, scale_factor=self.upsample_rate, mode='bilinear', align_corners=False) + x