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

__all__ = ['noop', 'same_padding1d', 'ZeroPad1d', 'ConvSP1d', 'convlayer']

# Cell
import torch
import torch.nn as nn

def noop(x): return x

def same_padding1d(seq_len,ks,stride=1,dilation=1):
    assert stride > 0
    assert dilation >= 1
    effective_ks = (ks - 1) * dilation + 1
    out_dim = (seq_len + stride - 1) // stride
    p = max(0, (out_dim - 1) * stride + effective_ks - seq_len)
    padding_before = p // 2
    padding_after = p - padding_before
    return padding_before, padding_after

class ZeroPad1d(nn.ConstantPad1d):
    def __init__(self, padding):
        super().__init__(padding, 0.)

class ConvSP1d(nn.Module):
    "Conv1d padding='same'"
    def __init__(self,c_in,c_out,ks,stride=1,padding='same',dilation=1,bias=True):
        super().__init__()
        self.ks, self.stride, self.dilation = ks, stride, dilation
        self.conv = nn.Conv1d(c_in,c_out,ks,stride=stride,padding=0,dilation=dilation,bias=bias)
        self.zeropad = ZeroPad1d
        self.weight = self.conv.weight
        self.bias = self.conv.bias

    def forward(self, x):
        padding = same_padding1d(x.shape[-1],self.ks,stride=self.stride,dilation=self.dilation)
        return self.conv(self.zeropad(padding)(x))

def convlayer(c_in,c_out,ks=3,padding='same',bias=True,stride=1,
              bn_init=False,zero_bn=False,bn_before=True,act_fn=True,**kwargs):
    '''conv layer (padding="same") + bn + act'''
    if ks % 2 == 1 and padding == 'same': padding = ks // 2
    layers = [ConvSP1d(c_in,c_out, ks, bias=bias, stride=stride) if padding == 'same' else \
    nn.Conv1d(c_in,c_out, ks, stride=stride, padding=padding, bias=bias)]
    bn = nn.BatchNorm1d(c_out)
    if bn_init: nn.init.constant_(bn.weight, 0. if zero_bn else 1.)
    if bn_before: layers.append(bn)
    if act_fn: layers.append(nn.ReLU())
    if not bn_before: layers.append(bn)
    return nn.Sequential(*layers)