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

# Imports
import os
import numpy as np
import pandas as pd
import random
import matplotlib.pyplot as plt

import torch
from torch.utils.data import DataLoader

from src.architecture.model import Ekg_SegNET
from src.data_loading.EKG_Dataset import EKG_dataset
from src.utils.utilities import evaluate_model_metrics, predict_patient_sig

numbers = range(1,803)

for val in numbers:    
    # Path to TEST DATA and MODEL to test
    num_pat = str (val)
    TEST_DATA_PATH = r"data/data_THEW/DELINEATOR/input/" + num_pat
    print (TEST_DATA_PATH)
    PATH_TO_MODEL_WEIGHTS = r"data/data_THEW/DELINEATOR/LSTM_EKG_Model.pt"
    
    
    # Initialization
    device_string = None
    num_workers = 0
    batch_size = 400 # 200 400 600
    window_size = 300 # 300 500 700
    window_overlap = 50 # 25 50 100
    out_dims = 2

    act_func = "relu" # "gelu"
    lstm_layers = 2 # 1 2 3
    fc_layers = 1 # 1 2
    
    # Class label initialization
    # label_dict = {"P": 0,
    #               "QRS": 1,
    #               "T": 2,
    #               "n": 3}
    
    label_dict = {"QRS": 1,
                  "n": 0}
    
    # Initialize seed
    seed = 42
    
    os.environ['PYTHONHASHSEED'] = '0'
    np.random.seed(seed)
    random.seed(seed)
    torch.random.manual_seed(seed)
    torch.cuda.manual_seed(seed=seed)
    torch.backends.cudnn.deterministic = True
    
    # device_string = "cpu"
    
    # Pick device to compute on
    if not device_string:
       device_string = "cuda" if torch.cuda.is_available() else "cpu"
    
    device = torch.device(device_string)
    
    # Initialize TEST dataset and loader
    patient_dataset = EKG_dataset(data_path=TEST_DATA_PATH,  window_size=window_size, window_overlap=window_overlap, label_dict=label_dict, shuffle=False)
    patient_loader = DataLoader(dataset=patient_dataset,
                              batch_size=batch_size,
                              num_workers=num_workers)
    
    # Initialize model and load weiths from path PATH_TO_MODEL_WEIGHTS
    net = Ekg_SegNET(act_func = act_func, out_dims=out_dims, lstm_layers=lstm_layers, fc_layers=fc_layers)
    net.to(device)
    net.load_state_dict(torch.load(PATH_TO_MODEL_WEIGHTS, map_location=device))
    net.eval() # turn off Batch Normalization layers and gradient
    
    patient_sig_dict = predict_patient_sig(model=net, dataset_loader=patient_loader, label_dict=label_dict, device=device, overlap=window_overlap)
    df_patient_sig = pd.DataFrame(patient_sig_dict)
    name = r"data/data_THEW/DELINEATOR/results_"+ num_pat +".csv"
    df_patient_sig.to_csv(name)

