# DEEP LEARNING BASED QRS DELINEATOR
# Malina Ondrej Bc. 
# FEKT UBMI 2020/2021


# ###################################################
# #################### Imports ######################
# ###################################################

import numpy as np

import torch
import torch.nn as nn
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence

#### For testing inference and backpropagation, comment with main() func in evaluate version
from src.data_loading.EKG_Dataset import EKG_dataset
from torch import optim

# ###################################################
# ############### Class definition ##################
# ###################################################


# Torch neural network class implementing segmentation EKG network from paper referenced in README
class Ekg_SegNET(nn.Module):

    # __init__ - constructor for initialization of neural network object
    def __init__(self, act_func="relu", in_dims=3, out_dims=4, lstm_layers=2, fc_layers=1):
        super(Ekg_SegNET, self).__init__()

        # Initialize convolution and batch normalization layers
        self.Conv1D_layer = nn.Conv1d(in_channels=in_dims, out_channels=64, kernel_size=3, padding=1)
        self.BatchNorm_layer = nn.BatchNorm1d(num_features=64)

        # Initialize LSTM layers
        self.BidirectLSTM1_layer = nn.LSTM(input_size=64, hidden_size=64, num_layers=lstm_layers, bidirectional=True)

        # Initialize fc layers, and their weights to values from normal distribution
        if fc_layers == 1:
            self.Dense_layers = nn.Linear(128, out_dims)
            nn.init.normal_(self.Dense_layers.weight, std=0.02)
        elif fc_layers == 2:
            self.Dense_layers = nn.Sequential(nn.Linear(128, 64), nn.Linear(64, out_dims))
            nn.init.normal_(self.Dense_layers[0].weight, std=0.02)
            nn.init.normal_(self.Dense_layers[1].weight, std=0.02)

        # Initialize activation functions
        self.Act_func = nn.ReLU() if act_func == "relu" else nn.GELU()
        self.SoftMax_func = nn.Softmax(dim=1)

    # Forward - forward propagates input tensor x through network layers
    def forward(self, x):

        # x 
        # |x| : (batch_size, feature_dim, time_dim)

        # Create support input for pack_padded_sequence input for LSTM layers
        lengths = [x.shape[-1]] * x.shape[0]
        # |lengths| : (batch_size, time_dim)

        # Propagate through convolution layer
        x = self.Conv1D_layer(x)
        x = self.BatchNorm_layer(x)
        x = self.Act_func(x)

        # Propagate through first bidirectional LSTM layer
        x = torch.transpose(x, 1, 2)
        x = pack_padded_sequence(x, lengths, batch_first=True)

        x, (_, _) = self.BidirectLSTM1_layer(x)

        x, _ = pad_packed_sequence(x, batch_first=True)

        # Propagate through dense layer
        x = self.Dense_layers(x)
        x = torch.transpose(x, 1, 2)

        # Compute probabilities from features
        x = self.SoftMax_func(x)

        return x


# ###################################################
# #### FOR MODEL FORWARD PASS TEST PURPOSE ONLY #####
# ###################################################

# Function to test models inference and backpropagation
def main():

    # Initialize torch
    torch.set_grad_enabled(True)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Load network
    net = Ekg_SegNET()
    net.to(device)

    # Training initialization
    optimizer = optim.Adam(net.parameters(), lr=0.001, weight_decay=0.0)
    loss_func = nn.NLLLoss()

    # Class label initialization
    label_dict = {"P": 0,
                  "QRS": 1,
                  "T": 2,
                  "n": 3}

    # Load dataset
    TRAIN_DATA_PATH = r"data/data_QTDB/TRAIN/"
    train_dataset = EKG_dataset(data_path=TRAIN_DATA_PATH,
                                window_size=300,
                                window_overlap=25,
                                label_dict=label_dict)

    # Load first sample
    sample_input, sample_label = train_dataset[0]
    # Add batch dimension
    sample_input = torch.unsqueeze(sample_input, 0).to(device)
    sample_label = torch.unsqueeze(sample_label, 0).to(device)

    # Infer output and compute loss
    with torch.cuda.amp.autocast():
        pred = net(sample_input)
        loss = loss_func(torch.log(pred), sample_label)

    # Backward propagation
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()


# Testing forward and backward propagation
if __name__ == "__main__":
    # execute only if run as a script
    main()
