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

__all__ = ['noop', 'mish', 'Mish', 'get_act_layer', 'same_padding1d', 'Pad1d', 'Conv1dSame', 'Chomp1d', 'Conv1dCausal',
           'Conv1d', 'CoordConv1D', 'LambdaPlus', 'Flatten', 'Squeeze', 'Unsqueeze', 'YRange', 'Temp']

# Cell
from fastai2.torch_core import Module
from ..imports import *

# Cell
def noop(x): return x

# Cell
# Misra, D. (2019). Mish: A Self Regularized Non-Monotonic Neural Activation Function. arXiv preprint arXiv:1908.08681.
# https://arxiv.org/abs/1908.08681
# GitHub: https://github.com/digantamisra98/Mish
@torch.jit.script
def mish(input):
    '''Applies the mish function element-wise: mish(x) = x * tanh(softplus(x)) = x * tanh(ln(1 + exp(x)))'''
    return input * torch.tanh(F.softplus(input))

class Mish(Module):
    def forward(self, input):
        return mish(input)

# Cell
def get_act_layer(act_fn, act_kwargs={}):
    act_fn = act_fn.lower()
    assert act_fn in ['relu', 'leakyrelu', 'prelu', 'elu', 'mish', 'swish'], 'incorrect act_fn'
    if act_fn == 'relu': return nn.ReLU()
    elif act_fn == 'leakyrelu': return nn.LeakyReLU(**act_kwargs)
    elif act_fn == 'prelu': return nn.PReLU(**act_kwargs)
    elif act_fn == 'elu': return nn.ELU(**act_kwargs)
    elif act_fn == 'mish': return Mish()
    elif act_fn == 'swish': return Swish()

# Cell
def same_padding1d(seq_len, ks, stride=1, 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 Pad1d(nn.ConstantPad1d):
    def __init__(self, padding, value=0.):
        super().__init__(padding, value)


class Conv1dSame(Module):
    "Conv1d with padding='same'"

    def __init__(self, c_in, c_out, ks=3, stride=1, dilation=1, **kwargs):
        self.ks, self.stride, self.dilation = ks, stride, dilation
        self.conv1d_same = nn.Conv1d(c_in, c_out, ks, stride=stride, dilation=dilation, **kwargs)
        self.pad = Pad1d


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

# Cell
# https://github.com/locuslab/TCN/blob/master/TCN/tcn.py
class Chomp1d(Module):
    def __init__(self, chomp_size):
        self.chomp_size = chomp_size

    def forward(self, x):
        return x[:, :, :-self.chomp_size].contiguous()


class Conv1dCausal(Module):
    def __init__(self, c_in, c_out, ks, stride=1, dilation=1, **kwargs):
        padding = (ks - 1) * dilation
        self.conv = nn.Conv1d(c_in, c_out, ks, stride=stride, padding=padding, dilation=dilation, **kwargs)
        self.chomp = Chomp1d(math.ceil(padding / stride))

    def forward(self, x):
        return self.chomp(self.conv(x))

# Cell
def Conv1d(c_in, c_out, ks=3, stride=1, padding='same', dilation=1, bias=True, act_fn='relu', act_kwargs={},
           bn_before_conv=False, bn_before_act=True, bn_after_act=False, zero_bn=False, **kwargs):
    '''conv1d with default padding='same', bn and act_fn (default = 'relu')'''
    layers = []
    if bn_before_conv: layers.append(nn.BatchNorm1d(c_in))
    if padding == 'same': layers.append(Conv1dSame(c_in, c_out, ks, stride=stride, dilation=dilation, bias=bias, **kwargs))
    elif padding == 'causal': layers.append(Conv1dCausal(c_in, c_out, ks, stride=stride, dilation=dilation, bias=bias, **kwargs))
    else:
        if padding == 'valid': padding = 0
        layers.append(nn.Conv1d(c_in, c_out, ks, stride=stride, padding=padding, dilation=dilation, bias=bias, **kwargs))
    if bn_before_act: layers.append(nn.BatchNorm1d(c_out))
    if act_fn: layers.append(get_act_layer(act_fn, act_kwargs))
    if bn_after_act:
        bn = nn.BatchNorm1d(c_out)
        nn.init.constant_(bn.weight, 0. if zero_bn else 1.)
        layers.append(bn)
    return nn.Sequential(*layers)

# Cell
class CoordConv1D(Module):
    def forward(self, x):
        bs, _, seq_len = x.size()
        cc = torch.arange(seq_len, device=device, dtype=torch.float) / (seq_len - 1)
        cc = cc * 2 - 1
        cc = cc.repeat(bs, 1, 1)
        x = torch.cat([x, cc], dim=1)
        return x

# Cell
class LambdaPlus(Module):
    def __init__(self, func, *args, **kwargs): self.func,self.args,self.kwargs=func,args,kwargs
    def forward(self, x): return self.func(x, *self.args, **self.kwargs)

# Cell
class Flatten(Module):
    def forward(self, x): return x.view(x.size(0), -1)

class Squeeze(Module):
    def __init__(self, dim=-1):
        self.dim = dim
    def forward(self, x): return x.squeeze(dim=self.dim)

class Unsqueeze(Module):
    def __init__(self, dim=-1):
        self.dim = dim
    def forward(self, x): return x.unsqueeze(dim=self.dim)

class YRange(Module):
    def __init__(self, y_range:tuple):
        self.y_range = y_range
        self.sigmoid = torch.sigmoid
    def forward(self, x):
        x = self.sigmoid(x)
        return x * (self.y_range[1] - self.y_range[0]) + self.y_range[0]

class Temp(Module):
    def __init__(self, temp):
        self.temp = float(temp)
        self.temp = nn.Parameter(torch.Tensor(1).fill_(self.temp).to(device))
    def forward(self, x):
        return x.div_(self.temp)