import torch
import torch.nn as nn
from torch.utils.data import TensorDataset,DataLoader
from torch.utils.tensorboard import SummaryWriter

from tqdm import tqdm
import os
import random
import argparse

from HRnet import IncUNet
from create_dataset import data_preprocess
from utils import testDataEval,save_model

def train(args):
    
    PATH = 'data'
    X_train = torch.load(PATH + '/ecgtoHR_train_data.pt')
    y_train = torch.load(PATH + '/ecgtoHR_train_labels.pt')

    X_test = torch.load(PATH + '/ecgtoHR_test_data.pt')
    y_test = torch.load(PATH + '/ecgtoHR_test_labels.pt')
    
    BATCH_SIZE= 64
    NUM_EPOCHS = 400
    best_loss = 1000

    train = TensorDataset(X_train,y_train)
    val = TensorDataset(X_test,y_test)
    trainLoader = DataLoader(train,batch_size = BATCH_SIZE,shuffle = True)
    valLoader = DataLoader(val, batch_size= BATCH_SIZE, shuffle=True)

    model = IncUNet((1,1,5000))
    model.cuda()
    criterion = torch.nn.SmoothL1Loss()
    optim = torch.optim.Adam(model.parameters(),lr = 0.001)
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optim,milestones=[100,200], gamma=0.1)

    NUM_EPOCHS = 400
    best_loss = 1000

    writer = SummaryWriter()

    if not(os.path.isdir("Saved_Model")):
        os.mkdir("Saved_Model")

    for epoch in tqdm(range(NUM_EPOCHS)):
        
        model.train()
        totalLoss = 0
        
        for step,(x,y) in enumerate(trainLoader):
            
            print('.',end = " ")
            ecg= x.unsqueeze(1).cuda()
            HR = y.unsqueeze(1).cuda()
            HR_pred = model(ecg)
            optim.zero_grad()
            loss = criterion(HR_pred,HR)
            totalLoss += loss.cpu().item()
            loss.backward()
            optim.step()
            
        print ('')
        print ("Epoch:{} Train Loss:{}".format(epoch + 1,totalLoss/(step+1)))
        
        totalTestLoss = testDataEval(model, valLoader, criterion)
        scheduler.step()
        
        if best_loss > totalTestLoss:
            print ("........Saving Best Model........")
            best_loss = totalTestLoss
            save_model("Saved_Model", epoch, model, optim, best_loss )
        
        writer.add_scalar("Loss/test",totalTestLoss, epoch )
        writer.add_scalar("Loss/train",totalLoss/(step+1),epoch )
            
    writer.close()


if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--preprocess_data',action = 'store_true', help = 'Use if True')
    parser.add_argument('--sampling_freq',default = 125, type = int, help = 'Sampling Frequency')
    parser.add_argument('--upsample_freq',default = 500, type = int, help = 'Resampling Frequency')
    parser.add_argument('--window_length',default = 5,type = int,help = 'Window Length in seconds')
    parser.add_argument('--data_path',help = 'Path to dataset')
    
    parser.add_argument('--random_seed',default = 5,type = int,help = 'Random Seed initializer')
    args = parser.parse_args()
    
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    random.seed(args.random_seed)
    torch.manual_seed(args.random_seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(args.random_seed) 

    if args.preprocess_data:
        data_preprocess(args)

    train(args)