import torch
import torch.nn as nn
import torch.nn.functional as F


class h_sigmoid(nn.Module):
    def __init__(self, inplace=True):
        super(h_sigmoid, self).__init__()
        self.relu = nn.ReLU6(inplace=inplace)

    def forward(self, x):
        return self.relu(x + 3) / 6


class h_swish(nn.Module):
    def __init__(self, inplace=True):
        super(h_swish, self).__init__()
        self.sigmoid = h_sigmoid(inplace=inplace)

    def forward(self, x):
        return x * self.sigmoid(x)


class CoordAtt(nn.Module):
    def __init__(self, inp, oup, reduction=32):
        super(CoordAtt, self).__init__()
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))

        mip = max(8, inp // reduction)

        self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
        self.bn1 = nn.BatchNorm2d(mip)

        self.bn2 = nn.BatchNorm2d(1)
        self.bn3 = nn.BatchNorm2d(1)
        self.act = h_swish()

        self.bn4 = nn.BatchNorm2d(mip)
        self.bn5 = nn.BatchNorm2d(mip)

        self.bn6 = nn.BatchNorm2d(1)
        self.bn7 = nn.BatchNorm2d(1)

        self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
        self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        x = torch.unsqueeze(x, 1) #2 1 2304 196
        identity = x

        n, c, h, w = x.size()#2 1 2304 196
        x_h = self.bn2(self.pool_h(x))#2 1 2304 1
        x_w = self.bn3(self.pool_w(x).permute(0, 1, 3, 2)) #2 1 196 1
        identity_x_w = x_w
        identity_x_h = x_h
        y = torch.cat([x_h, x_w], dim=2)
        y = self.conv1(y) #2 8 2500 1
        y = self.bn1(y)
        y = self.act(y)

        x_h, x_w = torch.split(y, [h, w], dim=2) #2 8 2304 1 | 2 8 196 1
        x_h = self.bn4(x_h)+identity_x_h
        x_w = self.bn5(x_w)+identity_x_w
        x_w = x_w.permute(0, 1, 3, 2)

        a_h = self.bn6(self.conv_h(x_h)).sigmoid() #2 1 2304 1
        a_w = self.bn7(self.conv_w(x_w)).sigmoid() #24 1 1 196

        out = identity * a_w * a_h #点×
        out = torch.squeeze(out, 1)
        return out

class CoordAtt_ori(nn.Module):
    def __init__(self, inp, oup, reduction=32):
        super(CoordAtt_ori, self).__init__()
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))

        mip = max(8, inp // reduction)

        self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
        self.bn1 = nn.BatchNorm2d(mip)
        self.act = h_swish()

        self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
        self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        x = torch.unsqueeze(x, 1)
        identity = x

        n, c, h, w = x.size()
        x_h = self.pool_h(x)
        x_w = self.pool_w(x).permute(0, 1, 3, 2)

        y = torch.cat([x_h, x_w], dim=2)
        y = self.conv1(y)
        y = self.bn1(y)
        y = self.act(y)

        x_h, x_w = torch.split(y, [h, w], dim=2)
        x_w = x_w.permute(0, 1, 3, 2)

        a_h = self.conv_h(x_h).sigmoid()
        a_w = self.conv_w(x_w).sigmoid()

        out = identity * a_w * a_h
        out = torch.squeeze(out, 1)
        return out