import numpy as np
import torch
import csv 
import os

device = 'cuda' if torch.cuda.is_available() else 'cpu'


def validate(model, val_loader, criterion): 
    model.eval()
    val_loss = []
    total_predictions = 0.0
    correct_predictions = 0.0
    for idx, (X,Y) in enumerate(val_loader):

        X = X.to(device)
        Y = Y.to(device)
            
        preds = model(X)
        
        loss = criterion(preds,Y)
        
        val_loss.append(loss.item())

        _, predicted = torch.max(preds, 1)
        total_predictions += Y.size(0)
        correct_predictions += (predicted == Y).sum().item()

    acc = (correct_predictions/total_predictions)*100.0

    return np.mean(val_loss), acc


def train(model, num_epochs, criterion, train_loader, val_loader, optimizer, scheduler, verbose = True):
    loss_train = []
    loss_val = []
    acc_train = []
    acc_val = []
    prev_dev_acc = 0
    for epoch in range(num_epochs): 
        model.train()
        epoch_loss = []
        total_predictions = 0.0
        correct_predictions = 0.0
        for idx, (X, Y) in enumerate(train_loader):
            X = X.to(device)
            Y = Y.to(device)
            
            optimizer.zero_grad()
            
            preds = model(X)
            
            loss = criterion(preds, Y)

            _, predicted = torch.max(preds, 1)
            total_predictions += Y.size(0)
            correct_predictions += (predicted == Y).sum().item()
            
            loss.backward()
            
            optimizer.step()

            epoch_loss.append(loss.item())

        loss_train.append(np.mean(epoch_loss))
        acc = (correct_predictions/total_predictions)*100.0
        acc_train.append(acc)

        epoch_val_loss, epoch_val_acc = validate(model, val_loader, criterion)
        loss_val.append(epoch_val_loss)
        acc_val.append(epoch_val_acc)
        if epoch_val_acc > prev_dev_acc:
            
            dir_name = "results/"
            test = os.listdir(dir_name)

            for item in test:
                if item.endswith(".pth"):
                    os.remove(os.path.join(dir_name, item))
            
            path = 'results/model_parameters_' + str(epoch_val_acc)[0: 5] + '.pth'
            torch.save(model.state_dict(), path)
            print('Saving model parameters...')
            print('Validation accuracy: ', epoch_val_acc)
            prev_dev_acc = epoch_val_acc
        

        if scheduler != None:
            scheduler.step()

        if verbose:
            print('EPOCH: ' + str(epoch))
            print('TRAIN_LOSS: ' + str(loss_train[-1]))
            print('TRAIN_ACC: ' + str(acc_train[-1]))
            print('VAL_LOSS: ' + str(loss_val[-1]))
            print('VAL_ACC: ' + str(acc_val[-1]))
            print('+'*25)


    return loss_train, loss_val, acc_train, acc_val



def evaluate(model, loader):
    observations = []
    heading = ['Predictions', 'Labels']
    model.eval()
    for idx, (X, Y) in enumerate(loader):
        X = X.to(device)
        Y = Y.to(device)      
        preds = model(X)
        _, predicted = torch.max(preds, 1)
        observations.append([predicted.item(), Y.item()])
    
    return np.array(observations)