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

from train.metrics import accuracy_SBM as accuracy
from train.metrics import accuracy_ogb

"""
    For GCNs
"""
def train_epoch_sparse(model, optimizer, device, data_loader, epoch, framework = 'pyg'):

    model.train()
    epoch_loss = 0
    epoch_train_acc = 0
    epoch_train_acc_ogb = 0
    nb_data = 0
    gpu_mem = 0
    if framework == 'pyg':
        for iter, batch_graphs in enumerate(data_loader):
            batch_x = batch_graphs.x.to(device)  # num x feat
            edge_index = batch_graphs.edge_index.to(device)
            batch_e = batch_graphs.edge_attr.to(device)
            batch_labels = batch_graphs.y.long().to(device)
            # batch_graphs = batch_graphs.to(device)  # add to satisfy the version to put the graph to the cuda
            optimizer.zero_grad()
            try:
                batch_pos_enc = batch_graphs.pos_enc.to(device)
                sign_flip = torch.rand(batch_pos_enc.size(1)).to(device)
                sign_flip[sign_flip >= 0.5] = 1.0
                sign_flip[sign_flip < 0.5] = -1.0
                batch_pos_enc = batch_pos_enc * sign_flip.unsqueeze(0)
                batch_scores = model.forward(batch_x, edge_index, batch_e, batch_pos_enc)
            except:
                # batch_scores = model.forward(batch_graphs.x, batch_graphs.edge_index)
                batch_scores = model.forward(batch_x, edge_index, batch_e)
            loss = model.loss(batch_scores, batch_labels)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.detach().item()
            epoch_train_acc += accuracy(batch_scores, batch_labels)
            # epoch_train_acc_ogb += accuracy_ogb(batch_scores, batch_labels)
            # nb_data += batch_labels.size(0)
            # print("Number: ", iter)
        epoch_loss /= (iter + 1)
        epoch_train_acc /= (iter + 1)
        # epoch_train_acc_ogb /= nb_data

        return epoch_loss, epoch_train_acc, optimizer
    elif framework == 'dgl':
        for iter, (batch_graphs, batch_labels) in enumerate(data_loader):
            batch_x = batch_graphs.ndata['feat'].to(device)  # num x feat
            batch_e = batch_graphs.edata['feat'].to(device)
            batch_labels = batch_labels.to(device)
            batch_graphs = batch_graphs.to(device) #add to satisfy the version to put the graph to the cuda
            optimizer.zero_grad()
            try:
                batch_pos_enc = batch_graphs.ndata['pos_enc'].to(device)
                sign_flip = torch.rand(batch_pos_enc.size(1)).to(device)
                sign_flip[sign_flip>=0.5] = 1.0; sign_flip[sign_flip<0.5] = -1.0
                batch_pos_enc = batch_pos_enc * sign_flip.unsqueeze(0)
                batch_scores = model.forward(batch_graphs, batch_x, batch_e, batch_pos_enc)
            except:
                batch_scores = model.forward(batch_graphs, batch_x, batch_e)
            loss = model.loss(batch_scores, batch_labels)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.detach().item()
            epoch_train_acc += accuracy(batch_scores, batch_labels)
            # print("Number: ", iter)
        epoch_loss /= (iter + 1)
        epoch_train_acc /= (iter + 1)

        return epoch_loss, epoch_train_acc, optimizer


def evaluate_network_sparse(model, device, data_loader, epoch, framework = 'pyg'):
    
    model.eval()
    epoch_test_loss = 0
    epoch_test_acc = 0
    nb_data = 0
    if framework == 'pyg':
        with torch.no_grad():
            for iter, batch_graphs in enumerate(data_loader):
                batch_x = batch_graphs.x.to(device)  # num x feat
                edge_index = batch_graphs.edge_index.to(device)
                batch_e = batch_graphs.edge_attr.to(device)
                batch_labels = batch_graphs.y.long().to(device)
                try:
                    batch_pos_enc = batch_graphs.pos_enc.to(device)
                    batch_scores = model.forward(batch_x, edge_index,batch_e, batch_pos_enc)
                except:
                    batch_scores = model.forward(batch_x, edge_index,batch_e)
                loss = model.loss(batch_scores, batch_labels)
                epoch_test_loss += loss.detach().item()
                epoch_test_acc += accuracy(batch_scores, batch_labels)
            epoch_test_loss /= (iter + 1)
            epoch_test_acc /= (iter + 1)

        return epoch_test_loss, epoch_test_acc
    elif framework == 'dgl':
        with torch.no_grad():
            for iter, (batch_graphs, batch_labels) in enumerate(data_loader):
                batch_x = batch_graphs.ndata['feat'].to(device)
                batch_e = batch_graphs.edata['feat'].to(device)
                batch_labels = batch_labels.to(device)
                batch_graphs = batch_graphs.to(device)  # add to satisfy the version to put the graph to the cuda
                try:
                    batch_pos_enc = batch_graphs.ndata['pos_enc'].to(device)
                    batch_scores = model.forward(batch_graphs, batch_x, batch_e, batch_pos_enc)
                except:
                    batch_scores = model.forward(batch_graphs, batch_x, batch_e)
                loss = model.loss(batch_scores, batch_labels)
                epoch_test_loss += loss.detach().item()
                epoch_test_acc += accuracy(batch_scores, batch_labels)
            epoch_test_loss /= (iter + 1)
            epoch_test_acc /= (iter + 1)

        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)
#     epoch_loss /= (iter + 1)
#     epoch_train_acc /= (iter + 1)
#
#     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)
#         epoch_test_loss /= (iter + 1)
#         epoch_test_acc /= (iter + 1)
#
#     return epoch_test_loss, epoch_test_acc
