import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np
from .model_helpers import *

import torchvision

import torch.fft
from torch import nn


class Net2DFast(nn.Module):
    def __init__(self, num_filts, num_classes=0, emb_dim=0, ip_height=128, resize_factor=0.5):
        super(Net2DFast, self).__init__()
        self.num_classes = num_classes
        self.emb_dim = emb_dim
        self.num_filts = num_filts
        self.resize_factor = resize_factor
        self.ip_height_rs = ip_height
        self.bneck_height = self.ip_height_rs//32

        # encoder
        self.conv_dn_0 = ConvBlockDownCoordF(1, num_filts//4, self.ip_height_rs, k_size=3, pad_size=1, stride=1)
        self.conv_dn_1 = ConvBlockDownCoordF(num_filts//4, num_filts//2, self.ip_height_rs//2, k_size=3, pad_size=1, stride=1)
        self.conv_dn_2 = ConvBlockDownCoordF(num_filts//2, num_filts, self.ip_height_rs//4, k_size=3, pad_size=1, stride=1)
        self.conv_dn_3 = nn.Conv2d(num_filts, num_filts*2, 3, padding=1)
        self.conv_dn_3_bn = nn.BatchNorm2d(num_filts*2)

        # bottleneck
        self.conv_1d = nn.Conv2d(num_filts*2, num_filts*2, (self.ip_height_rs//8,1), padding=0)
        self.conv_1d_bn = nn.BatchNorm2d(num_filts*2)
        self.att = SelfAttention(num_filts*2, num_filts*2)

        # decoder
        self.conv_up_2 = ConvBlockUpF(num_filts*2, num_filts//2, self.ip_height_rs//8)
        self.conv_up_3 = ConvBlockUpF(num_filts//2, num_filts//4, self.ip_height_rs//4)
        self.conv_up_4 = ConvBlockUpF(num_filts//4, num_filts//4, self.ip_height_rs//2)

        # output
        # +1 to include background class for class output
        self.conv_op  = nn.Conv2d(num_filts//4, num_filts//4, kernel_size=3, padding=1)
        self.conv_op_bn = nn.BatchNorm2d(num_filts//4)
        self.conv_size_op  = nn.Conv2d(num_filts//4, 2, kernel_size=1, padding=0)
        self.conv_classes_op = nn.Conv2d(num_filts//4, self.num_classes+1, kernel_size=1, padding=0)

        if self.emb_dim > 0:
            self.conv_emb = nn.Conv2d(num_filts, self.emb_dim, kernel_size=1, padding=0)


    def forward(self, ip, return_feats=False):

        # encoder
        x1 = self.conv_dn_0(ip)
        x2 = self.conv_dn_1(x1)
        x3 = self.conv_dn_2(x2)
        x3 = F.relu(self.conv_dn_3_bn(self.conv_dn_3(x3)), inplace=True)

        # bottleneck
        x = F.relu(self.conv_1d_bn(self.conv_1d(x3)), inplace=True)
        x = self.att(x)
        x = x.repeat([1,1,self.bneck_height*4,1])

        # decoder
        x = self.conv_up_2(x+x3)
        x = self.conv_up_3(x+x2)
        x = self.conv_up_4(x+x1)

        # output
        x = F.relu(self.conv_op_bn(self.conv_op(x)), inplace=True)
        cls  = self.conv_classes_op(x)
        comb = torch.softmax(cls, 1)

        op = {}
        op['pred_det']   = comb[:,:-1, :, :].sum(1).unsqueeze(1)
        op['pred_size']  = F.relu(self.conv_size_op(x), inplace=True)
        op['pred_class'] = comb
        op['pred_class_un_norm'] = cls
        if self.emb_dim > 0:
            op['pred_emb'] = self.conv_emb(x)
        if return_feats:
            op['features'] = x

        return op


class Net2DFastNoAttn(nn.Module):
    def __init__(self, num_filts, num_classes=0, emb_dim=0, ip_height=128, resize_factor=0.5):
        super(Net2DFastNoAttn, self).__init__()

        self.num_classes = num_classes
        self.emb_dim = emb_dim
        self.num_filts = num_filts
        self.resize_factor = resize_factor
        self.ip_height_rs = ip_height
        self.bneck_height = self.ip_height_rs//32

        self.conv_dn_0 = ConvBlockDownCoordF(1, num_filts//4, self.ip_height_rs, k_size=3, pad_size=1, stride=1)
        self.conv_dn_1 = ConvBlockDownCoordF(num_filts//4, num_filts//2, self.ip_height_rs//2, k_size=3, pad_size=1, stride=1)
        self.conv_dn_2 = ConvBlockDownCoordF(num_filts//2, num_filts, self.ip_height_rs//4, k_size=3, pad_size=1, stride=1)
        self.conv_dn_3 = nn.Conv2d(num_filts, num_filts*2, 3, padding=1)
        self.conv_dn_3_bn = nn.BatchNorm2d(num_filts*2)

        self.conv_1d = nn.Conv2d(num_filts*2, num_filts*2, (self.ip_height_rs//8,1), padding=0)
        self.conv_1d_bn = nn.BatchNorm2d(num_filts*2)


        self.conv_up_2 = ConvBlockUpF(num_filts*2, num_filts//2, self.ip_height_rs//8)
        self.conv_up_3 = ConvBlockUpF(num_filts//2, num_filts//4, self.ip_height_rs//4)
        self.conv_up_4 = ConvBlockUpF(num_filts//4, num_filts//4, self.ip_height_rs//2)

        # output
        # +1 to include background class for class output
        self.conv_op  = nn.Conv2d(num_filts//4, num_filts//4, kernel_size=3, padding=1)
        self.conv_op_bn = nn.BatchNorm2d(num_filts//4)
        self.conv_size_op  = nn.Conv2d(num_filts//4, 2, kernel_size=1, padding=0)
        self.conv_classes_op = nn.Conv2d(num_filts//4, self.num_classes+1, kernel_size=1, padding=0)

        if self.emb_dim > 0:
            self.conv_emb = nn.Conv2d(num_filts, self.emb_dim, kernel_size=1, padding=0)

    def forward(self, ip, return_feats=False):

        x1 = self.conv_dn_0(ip)
        x2 = self.conv_dn_1(x1)
        x3 = self.conv_dn_2(x2)
        x3 = F.relu(self.conv_dn_3_bn(self.conv_dn_3(x3)), inplace=True)

        x = F.relu(self.conv_1d_bn(self.conv_1d(x3)), inplace=True)
        x = x.repeat([1,1,self.bneck_height*4,1])

        x = self.conv_up_2(x+x3)
        x = self.conv_up_3(x+x2)
        x = self.conv_up_4(x+x1)

        x = F.relu(self.conv_op_bn(self.conv_op(x)), inplace=True)
        cls  = self.conv_classes_op(x)
        comb = torch.softmax(cls, 1)

        op = {}
        op['pred_det']   = comb[:,:-1, :, :].sum(1).unsqueeze(1)
        op['pred_size']  = F.relu(self.conv_size_op(x), inplace=True)
        op['pred_class'] = comb
        op['pred_class_un_norm'] = cls
        if self.emb_dim > 0:
            op['pred_emb'] = self.conv_emb(x)
        if return_feats:
            op['features'] = x

        return op


class Net2DFastNoCoordConv(nn.Module):
    def __init__(self, num_filts, num_classes=0, emb_dim=0, ip_height=128, resize_factor=0.5):
        super(Net2DFastNoCoordConv, self).__init__()

        self.num_classes = num_classes
        self.emb_dim = emb_dim
        self.num_filts = num_filts
        self.resize_factor = resize_factor
        self.ip_height_rs = ip_height
        self.bneck_height = self.ip_height_rs//32

        self.conv_dn_0 = ConvBlockDownStandard(1, num_filts//4, self.ip_height_rs, k_size=3, pad_size=1, stride=1)
        self.conv_dn_1 = ConvBlockDownStandard(num_filts//4, num_filts//2, self.ip_height_rs//2, k_size=3, pad_size=1, stride=1)
        self.conv_dn_2 = ConvBlockDownStandard(num_filts//2, num_filts, self.ip_height_rs//4, k_size=3, pad_size=1, stride=1)
        self.conv_dn_3 = nn.Conv2d(num_filts, num_filts*2, 3, padding=1)
        self.conv_dn_3_bn = nn.BatchNorm2d(num_filts*2)

        self.conv_1d = nn.Conv2d(num_filts*2, num_filts*2, (self.ip_height_rs//8,1), padding=0)
        self.conv_1d_bn = nn.BatchNorm2d(num_filts*2)

        self.att = SelfAttention(num_filts*2, num_filts*2)

        self.conv_up_2 = ConvBlockUpStandard(num_filts*2, num_filts//2, self.ip_height_rs//8)
        self.conv_up_3 = ConvBlockUpStandard(num_filts//2, num_filts//4, self.ip_height_rs//4)
        self.conv_up_4 = ConvBlockUpStandard(num_filts//4, num_filts//4, self.ip_height_rs//2)

        # output
        # +1 to include background class for class output
        self.conv_op  = nn.Conv2d(num_filts//4, num_filts//4, kernel_size=3, padding=1)
        self.conv_op_bn = nn.BatchNorm2d(num_filts//4)
        self.conv_size_op  = nn.Conv2d(num_filts//4, 2, kernel_size=1, padding=0)
        self.conv_classes_op = nn.Conv2d(num_filts//4, self.num_classes+1, kernel_size=1, padding=0)

        if self.emb_dim > 0:
            self.conv_emb = nn.Conv2d(num_filts, self.emb_dim, kernel_size=1, padding=0)

    def forward(self, ip, return_feats=False):

        x1 = self.conv_dn_0(ip)
        x2 = self.conv_dn_1(x1)
        x3 = self.conv_dn_2(x2)
        x3 = F.relu(self.conv_dn_3_bn(self.conv_dn_3(x3)), inplace=True)

        x = F.relu(self.conv_1d_bn(self.conv_1d(x3)), inplace=True)
        x = self.att(x)
        x = x.repeat([1,1,self.bneck_height*4,1])

        x = self.conv_up_2(x+x3)
        x = self.conv_up_3(x+x2)
        x = self.conv_up_4(x+x1)

        x = F.relu(self.conv_op_bn(self.conv_op(x)), inplace=True)
        cls  = self.conv_classes_op(x)
        comb = torch.softmax(cls, 1)

        op = {}
        op['pred_det']   = comb[:,:-1, :, :].sum(1).unsqueeze(1)
        op['pred_size']  = F.relu(self.conv_size_op(x), inplace=True)
        op['pred_class'] = comb
        op['pred_class_un_norm'] = cls
        if self.emb_dim > 0:
            op['pred_emb'] = self.conv_emb(x)
        if return_feats:
            op['features'] = x

        return op