#
# From https://github.com/locuslab/TCN/blob/master/TCN/tcn.py
# Adapted to have strides and batchnorm
#

import torch.nn as nn
from math import ceil
from torch.nn import BatchNorm1d


class Chomp1d(nn.Module):
    def __init__(self, chomp_size):
        super(Chomp1d, self).__init__()
        self.chomp_size = chomp_size

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


class StridedTemporalBlock(nn.Module):
    def __init__(self, n_inputs, n_hidden, n_outputs, kernel_size, stride, dilation, padding, final_output,
                 dropout=0.2):
        super(StridedTemporalBlock, self).__init__()
        self.stride = stride
        self.conv1 = nn.Conv1d(n_inputs, n_hidden, kernel_size,
                               stride=stride, padding=padding, dilation=dilation)
        self.batch1 = BatchNorm1d(n_hidden)
        self.chomp1 = Chomp1d(int(padding / stride))
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(dropout)

        self.conv2 = nn.Conv1d(n_hidden, n_outputs, kernel_size,
                               stride=stride, padding=padding, dilation=dilation)
        self.batch2 = BatchNorm1d(n_outputs)
        self.chomp2 = Chomp1d(int(padding / stride))
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(dropout)

        self.net = nn.Sequential(self.conv1, self.chomp1, self.relu1, self.batch1, self.dropout1,
                                 self.conv2, self.chomp2, self.relu2, self.batch2, self.dropout2)
        self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None
        self.residual_downsample = nn.AdaptiveAvgPool1d(final_output)
        self.relu = nn.ReLU()
        self.init_weights()

    def init_weights(self):
        self.conv1.weight.data.normal_(0, 0.01)
        self.conv2.weight.data.normal_(0, 0.01)
        if self.downsample is not None:
            self.downsample.weight.data.normal_(0, 0.01)

    def forward(self, x):
        out = self.net(x)
        res = x if self.downsample is None else self.downsample(x)
        if self.stride > 1:
            res = self.residual_downsample(res)
        return self.relu(out + res)


class TemporalConvNet(nn.Module):
    def __init__(self, num_inputs, num_channels, stride, num_time_length, kernel_size, dropout=0.2):
        super(TemporalConvNet, self).__init__()
        layers = []
        num_levels = len(num_channels)
        previous_final = num_time_length
        for i in range(0, num_levels - 1, 2):
            previous_final = ceil(previous_final / (stride * 2))
            dilation_size = 2 ** i
            in_channels = num_inputs if i == 0 else num_channels[i - 1]
            hidden_channels = num_channels[i]
            out_channels = num_channels[i + 1]
            layers += [StridedTemporalBlock(in_channels, hidden_channels, out_channels, kernel_size, stride=stride,
                                            dilation=dilation_size,
                                            padding=(kernel_size - 1) * dilation_size, dropout=dropout,
                                            final_output=previous_final)]

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)
