# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/101b_models.ResNetPlus.ipynb (unless otherwise specified).

__all__ = ['ResBlockPlus', 'ResNetPlus']

# Cell
from ..imports import *
from .utils import *
from .layers import *

# Cell
class ResBlockPlus(Module):
    def __init__(self, ni, nf, ks=[7, 5, 3], coord=False, separable=False, bn_1st=True, zero_norm=False, sa=False, se=None, act=nn.ReLU, act_kwargs={}):
        self.convblock1 = ConvBlock(ni, nf, ks[0], coord=coord, separable=separable, bn_1st=bn_1st, act=act, act_kwargs=act_kwargs)
        self.convblock2 = ConvBlock(nf, nf, ks[1], coord=coord, separable=separable, bn_1st=bn_1st, act=act, act_kwargs=act_kwargs)
        self.convblock3 = ConvBlock(nf, nf, ks[2], coord=coord, separable=separable, zero_norm=zero_norm, act=None)
        self.se = SEModule1d(nf, reduction=se, act=act) if se and nf//se > 0 else noop
        self.sa = SimpleSelfAttention(nf, ks=1) if sa else noop
        self.shortcut = BN1d(ni) if ni == nf else ConvBlock(ni, nf, 1, coord=coord, act=None)
        self.add = Add()
        self.act = act(**act_kwargs)

        self._init_cnn(self)

    def _init_cnn(self, m):
        if getattr(self, 'bias', None) is not None: nn.init.constant_(self.bias, 0)
        if isinstance(self, (nn.Conv1d,nn.Conv2d,nn.Conv3d,nn.Linear)): nn.init.kaiming_normal_(self.weight)
        for l in m.children(): self._init_cnn(l)

    def forward(self, x):
        res = x
        x = self.convblock1(x)
        x = self.convblock2(x)
        x = self.convblock3(x)
        x = self.se(x)
        x = self.sa(x)
        x = self.add(x, self.shortcut(res))
        x = self.act(x)
        return x


@delegates(ResBlockPlus.__init__)
class ResNetPlus(Module):
    def __init__(self, c_in, c_out, nf=64, sa=False, se=None, fc_dropout=0., **kwargs):
        self.resblock1 = ResBlockPlus(c_in,   nf,     se=se,   **kwargs)
        self.resblock2 = ResBlockPlus(nf,     nf * 2, se=se,   **kwargs)
        self.resblock3 = ResBlockPlus(nf * 2, nf * 2, sa=sa, **kwargs)
        self.gap = GAP1d(1)
        self.do = nn.Dropout(fc_dropout) if fc_dropout > 0 else noop
        self.fc = nn.Linear(nf * 2, c_out)

    def forward(self, x):
        x = self.resblock1(x)
        x = self.resblock2(x)
        x = self.resblock3(x)
        x = self.gap(x)
        x = self.do(x)
        return self.fc(x)