"""
    Utility functions for training one epoch 
    and evaluating one epoch
"""
import torch
import torch.nn as nn
import math

from train.metrics import accuracy_TU as accuracy

"""
    For GCNs
"""
def train_epoch_sparse(model, optimizer, device, dataset, train_idx):
    model.train()
    epoch_loss = 0
    epoch_train_acc = 0
    nb_data = 0
    gpu_mem = 0
    # for iter, (batch_graphs, batch_labels) in enumerate(data_loader):
    batch_x = dataset.dataset[0].x.to(device)
    batch_e = dataset.edge_attr.to(device)
    batch_labels = dataset.dataset[0].y.long().to(device)
    edge_index = dataset.dataset[0].edge_index.long().to(device)
    train_idx = train_idx.to(device)
    optimizer.zero_grad()
    batch_scores = model.forward(batch_x, edge_index, batch_e)[train_idx]
    loss = model.loss(batch_scores, batch_labels[train_idx]).to(torch.float)
    loss.backward()
    optimizer.step()
    epoch_loss = loss.detach().item()
    epoch_train_acc = accuracy(batch_scores, batch_labels[train_idx]) / train_idx.size(0)

    return epoch_loss, epoch_train_acc, optimizer

def evaluate_network_sparse(model, device, dataset, val_idx):
    model.eval()
    epoch_test_loss = 0
    epoch_test_acc = 0
    nb_data = 0
    with torch.no_grad():
        batch_x = dataset.dataset[0].x.to(device)
        batch_e = dataset.edge_attr.to(device)
        batch_labels = dataset.dataset[0].y.long().to(device)
        edge_index = dataset.dataset[0].edge_index.long().to(device)
        val_idx = val_idx.to(device)
            
        batch_scores = model.forward(batch_x, edge_index, batch_e)[val_idx]
        loss = model.loss(batch_scores, batch_labels[val_idx]).to(torch.float)
        epoch_test_loss = loss.detach().item()
        epoch_test_acc = accuracy(batch_scores, batch_labels[val_idx]) / val_idx.size(0)

    return epoch_test_loss, epoch_test_acc




# """
#     For WL-GNNs
# """
# def train_epoch_dense(model, optimizer, device, data_loader, epoch, batch_size):
#     model.train()
#     epoch_loss = 0
#     epoch_train_acc = 0
#     nb_data = 0
#     gpu_mem = 0
#     optimizer.zero_grad()
#     for iter, (x_with_node_feat, labels) in enumerate(data_loader):
#         x_with_node_feat = x_with_node_feat.to(device)
#         labels = labels.to(device)
#
#         scores = model.forward(x_with_node_feat)
#         loss = model.loss(scores, labels)
#         loss.backward()
#
#         if not (iter%batch_size):
#             optimizer.step()
#             optimizer.zero_grad()
#
#         epoch_loss += loss.detach().item()
#         epoch_train_acc += accuracy(scores, labels)
#         nb_data += labels.size(0)
#     epoch_loss /= (iter + 1)
#     epoch_train_acc /= nb_data
#
#     return epoch_loss, epoch_train_acc, optimizer
#
# def evaluate_network_dense(model, device, data_loader, epoch):
#     model.eval()
#     epoch_test_loss = 0
#     epoch_test_acc = 0
#     nb_data = 0
#     with torch.no_grad():
#         for iter, (x_with_node_feat, labels) in enumerate(data_loader):
#             x_with_node_feat = x_with_node_feat.to(device)
#             labels = labels.to(device)
#
#             scores = model.forward(x_with_node_feat)
#             loss = model.loss(scores, labels)
#             epoch_test_loss += loss.detach().item()
#             epoch_test_acc += accuracy(scores, labels)
#             nb_data += labels.size(0)
#         epoch_test_loss /= (iter + 1)
#         epoch_test_acc /= nb_data
#
#     return epoch_test_loss, epoch_test_acc




def check_patience(all_losses, best_loss, best_epoch, curr_loss, curr_epoch, counter):
    if curr_loss < best_loss:
        counter = 0
        best_loss = curr_loss
        best_epoch = curr_epoch
    else:
        counter += 1
    return best_loss, best_epoch, counter