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

__all__ = ['ResBlock', 'ResNet']

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

# Cell
class ResBlock(Module):
    def __init__(self, ni, nf, ks=[7, 5, 3]):
        self.conv1 = Conv1d(ni, nf, ks[0], padding='same', act_fn='relu')
        self.conv2 = Conv1d(nf, nf, ks[1], padding='same', act_fn='relu')
        self.conv3 = Conv1d(nf, nf, ks[2], padding='same', act_fn='relu')

        # expand channels for the sum if necessary
        self.shortcut = noop if ni == nf else Conv1d(ni, nf, ks=1, act_fn=False)
        self.act_fn = nn.ReLU()

    def forward(self, x):
        res = x
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        sc = self.shortcut(res)
        x += sc
        x = self.act_fn(x)
        return x

class ResNet(Module):
    def __init__(self,c_in, c_out):
        nf = 64
        self.block1 = ResBlock(c_in, nf, ks=[7, 5, 3])
        self.block2 = ResBlock(nf, nf * 2, ks=[7, 5, 3])
        self.block3 = ResBlock(nf * 2, nf * 2, ks=[7, 5, 3])
        self.gap = nn.AdaptiveAvgPool1d(1)
        self.squeeze = Squeeze(-1)
        self.fc = nn.Linear(nf * 2, c_out)

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        x = self.squeeze(self.gap(x))
        return self.fc(x)