# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/vision.models.attention_modules.ipynb.

# %% auto 0
__all__ = ['ECA', 'ShuffleAttention', 'TripletAttention']

# %% ../../../nbs/vision.models.attention_modules.ipynb 1
import math, torch
import torch.nn as nn
from fastai.torch_core import Module
from fastai.layers import ConvLayer
from torch.nn.parameter import Parameter

# %% ../../../nbs/vision.models.attention_modules.ipynb 5
# ECA modified from https://github.com/BangguWu/ECANet
# ECANet - MIT License - Copyright (c) 2019 BangguWu & Qilong Wang

class ECA(Module):
    "Efficient Channel Attention, from https://arxiv.org/abs/1910.03151."
    def __init__(self,
        nf, # number of input features
        ks:int=None, # if set, Cov1D uses a fixed kernel size instead of adaptive kernel size
        gamma:int=2, # used for adaptive kernel size, see paper for more details
        beta:int=1 # used for adaptive kernel size, see paper for more details
    ):
        if ks is None:
            ks = int(abs(math.log(nf, 2)+ beta) / gamma)
            ks = max(3, ks if ks % 2 == 1 else ks + 1)
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=ks, padding=(ks-1)//2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        y = self.avg_pool(x)
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
        y = self.sigmoid(y)
        return x * y.expand_as(x)

# %% ../../../nbs/vision.models.attention_modules.ipynb 8
class ShuffleAttention(Module):
    "Implementation of Shuffle Attention, from https://arxiv.org/abs/2102.00240"
    def __init__(self,
        nf, # number of input features
        groups=64 # number of subfeature groups, usually 32 or 64
    ):
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.sigmoid = nn.Sigmoid()
        self.groups = groups
        self.cweight = Parameter(torch.zeros(1, nf // (2 * groups), 1, 1))
        self.cbias = Parameter(torch.ones(1, nf // (2 * groups), 1, 1))
        self.sweight = Parameter(torch.zeros(1, nf // (2 * groups), 1, 1))
        self.sbias = Parameter(torch.ones(1, nf // (2 * groups), 1, 1))
        self.norm = nn.GroupNorm(nf // (2 * groups), nf // (2 * groups))

    def forward(self, x):
        b, c, h, w = x.shape
        x = x.reshape(b*self.groups, -1, h, w) # group into subfeatures
        xc, xs = x.chunk(2, dim=1) # split channels

        # channel attention
        yc = self.pool(xc)
        yc = self.cweight * yc + self.cbias
        xc = xc * self.sigmoid(yc)

        # spatial attention
        ys = self.norm(xs)
        ys = self.sweight * ys + self.sbias
        xs = xs * self.sigmoid(ys)

        # concat and shuffle
        x = torch.cat([xc, xs], dim=1)
        x = x.reshape(b, 2, -1, h, w)
        x = x.permute(0, 2, 1, 3, 4)
        return x.reshape(b, c, h, w)

# %% ../../../nbs/vision.models.attention_modules.ipynb 10
# TripletAttention lightly modified from https://github.com/landskape-ai/triplet-attention
# triplet-attention - MIT License - Copyright (c) 2020 LandskapeAI

class ZPool(Module):
    def forward(self, x):
        return torch.cat((torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1)

# %% ../../../nbs/vision.models.attention_modules.ipynb 11
class AttentionGate(Module):
    def __init__(self,
        ks:int=7 # kernel size for Conv2D
    ):
        self.compress = ZPool()
        self.conv = ConvLayer(2, 1, ks, padding=(ks-1)//2, act_cls=None)

    def forward(self, x):
        x_compress = self.compress(x)
        x_out = self.conv(x_compress)
        scale = torch.sigmoid_(x_out)
        return x * scale

# %% ../../../nbs/vision.models.attention_modules.ipynb 12
class TripletAttention(Module):
    "Lightly modified implementation of Triplet Attention, from http://arxiv.org/abs/2010.03045"
    def __init__(self,
        nf, # unused input features, for compatibility
        ks:int=7, # kernel size for AttentionGate
        no_spatial=False # exclude Spatial attention as third attention
    ):
        self.cw = AttentionGate(ks=ks)
        self.hc = AttentionGate(ks=ks)
        self.no_spatial = no_spatial
        if not no_spatial:
            self.hw = AttentionGate(ks=ks)

    def forward(self, x):
        x_perm1 = x.permute(0, 2, 1, 3).contiguous()
        x_out1 = self.cw(x_perm1)
        x_out11 = x_out1.permute(0, 2, 1, 3).contiguous()
        x_perm2 = x.permute(0, 3, 2, 1).contiguous()
        x_out2 = self.hc(x_perm2)
        x_out21 = x_out2.permute(0, 3, 2, 1).contiguous()
        if not self.no_spatial:
            x_out = self.hw(x)
            x_out = 1 / 3 * (x_out + x_out11 + x_out21)
        else:
            x_out = 1 / 2 * (x_out11 + x_out21)
        return x_out
