import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import weight_norm


class TCN_cnntcn(nn.Module):   # original CNNTCN network
    def __init__(self, nhid, levels, kernel_size, dropout):
        super(TCN_cnntcn, self).__init__()

        self.conv2d_1 = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=(3,3), 
                                  stride=1, padding=(1,0), padding_mode='zeros')
        self.elu_1 = nn.ELU()
        self.max_pooling_1 = nn.MaxPool2d(kernel_size=(1,4))
        self.dropout_1 = nn.Dropout(dropout)      
        
        self.conv2d_2 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=(3,3), 
                                  stride=1, padding=(1,0), padding_mode='zeros')
        self.elu_2 = nn.ELU()
        self.max_pooling_2 = nn.MaxPool2d(kernel_size=(1,3))
        self.dropout_2 = nn.Dropout(dropout)

        self.conv2d_3 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=(1,9), stride=1)     
        self.elu_3 = nn.ELU()        
        self.dropout_3 = nn.Dropout(dropout)    
        
        num_channels = [nhid] * levels
        self.tcn = TemporalConvNetCNNTCN(num_inputs=16, num_channels=num_channels,
                                   kernel_size=kernel_size, dropout=dropout)
        self.linear = nn.Linear(num_channels[-1], out_features=2)

    def forward(self, x):       
        x = x[:, None, :]
        output = self.conv2d_1(x) 
        output = self.elu_1(output)
        output = self.max_pooling_1(output)
        output = self.dropout_1(output)
        output = self.conv2d_2(output)    
        output = self.elu_2(output)
        output = self.max_pooling_2(output)        
        output = self.dropout_2(output)
        output = self.conv2d_3(output)        
        output = self.elu_3(output)
        output = self.dropout_3(output)
        output = torch.squeeze(output, dim=3)
        output = self.tcn(output).transpose(1, 2)
        output = self.linear(output).double()
        output = F.log_softmax(output, dim=2).transpose(1, 2)
        return output


class TCN_simple(nn.Module):  # only TCN network
    def __init__(self, nhid, levels, kernel_size, dropout):
        super(TCN_simple, self).__init__()

        num_channels = [nhid] * levels
        self.tcn = TemporalConvNet(num_inputs=120, num_channels=num_channels, 
                                   kernel_size=kernel_size, dropout=dropout)
        self.linear = nn.Linear(num_channels[-1], out_features=2)

    def forward(self, x):
        output = self.tcn(x).transpose(1, 2)
        output = self.linear(output).double()
        output = F.log_softmax(output, dim=2).transpose(1, 2)
        return output


class TemporalBlock(nn.Module):  # Residual block with stacked convolutions
    def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2):
        super(TemporalBlock, self).__init__()
        self.conv1 = weight_norm(nn.Conv1d(n_inputs, n_outputs, kernel_size,
                                           stride=stride, padding=int(padding/2), dilation=dilation))
        self.elu1 = nn.ELU()
        self.dropout1 = nn.Dropout(dropout)
        self.conv2 = weight_norm(nn.Conv1d(n_outputs, n_outputs, kernel_size,
                                           stride=stride, padding=int(padding/2), dilation=dilation))
        self.elu2 = nn.ELU()
        self.dropout2 = nn.Dropout(dropout)
        self.net = nn.Sequential(self.conv1, self.elu1, self.dropout1,
                                 self.conv2, self.elu2, self.dropout2)
        self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None
        self.elu = nn.ELU()
        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)
        return self.elu(out + res)


class TemporalBlockCNNTCN(nn.Module):  # Residual block only with one convolution
    def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2):
        super(TemporalBlockCNNTCN, self).__init__()
        self.conv1 = weight_norm(nn.Conv1d(n_inputs, n_outputs, kernel_size,
                                           stride=stride, padding=int(padding/2), dilation=dilation))
        self.elu1 = nn.ELU()
        self.dropout1 = nn.Dropout(dropout)
        self.net = nn.Sequential(self.conv1, self.elu1, self.dropout1)
        self.downsample = nn.Conv1d(n_inputs, n_outputs, 1) if n_inputs != n_outputs else None
        self.elu = nn.ELU()
        self.init_weights()

    def init_weights(self):
        self.conv1.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)
        return self.elu(out + res)


class TemporalConvNet(nn.Module):  # Generate TCN block
    def __init__(self, num_inputs, num_channels, kernel_size, dropout):
        super(TemporalConvNet, self).__init__()
        layers = []
        num_levels = len(num_channels)
        for i in range(num_levels):
            dilation_size = 2 ** i
            in_channels = num_inputs if i == 0 else num_channels[i-1]
            out_channels = num_channels[i]
            layers += [TemporalBlock(in_channels, out_channels, kernel_size, stride=1, dilation=dilation_size,
                                     padding=(kernel_size-1) * dilation_size, dropout=dropout)]

        self.network = nn.Sequential(*layers)

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


class TemporalConvNetCNNTCN(nn.Module):  # Generate TCN block for CNNTCN
    def __init__(self, num_inputs, num_channels, kernel_size, dropout):
        super(TemporalConvNetCNNTCN, self).__init__()
        layers = []
        num_levels = len(num_channels)
        for i in range(num_levels):
            dilation_size = 2 ** i
            in_channels = num_inputs if i == 0 else num_channels[i-1]
            out_channels = num_channels[i]
            layers += [TemporalBlockCNNTCN(in_channels, out_channels, kernel_size, stride=1, dilation=dilation_size,
                                     padding=(kernel_size-1) * dilation_size, dropout=dropout)]

        self.network = nn.Sequential(*layers)

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


class BiLSTM(nn.Module):  # Bidirectional LSTM network
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(BiLSTM, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, bidirectional=True, dropout=0.1)
        self.fc = nn.Linear(hidden_size * 2, num_classes)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size)  # 2 for bidirection
        c0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size)
        # h0, c0 = h0.cuda(), c0.cuda()   # In case of GPU uncomment this!
        out, _ = self.lstm(x, (h0, c0))
        out = self.fc(out)
        y = F.log_softmax(out, dim=2)
        return y


class BiGRU(nn.Module):  # Bidirectional GRU network
    def __init__(self, input_size, hidden_size, num_layers, num_classes):
        super(BiGRU, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.gru = nn.GRU(input_size, hidden_size, num_layers, batch_first=True, bidirectional=True, dropout=0.1)
        self.fc = nn.Linear(hidden_size * 2, num_classes)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers * 2, x.size(0), self.hidden_size) # 2 for bidirection
        # h0 = h0.cuda()   # In case of GPU uncomment this!
        out, _ = self.gru(x, (h0))
        out = self.fc(out)
        y = F.log_softmax(out, dim=2)
        return y

