<reponame>MohsenFazaeli/SAME<gh_stars>10-100
import copy
import numpy as np
import os
import torch
import torch.nn.functional as F
from torch_geometric.data import DataLoader
import shutil
from sklearn.metrics import roc_auc_score
from sklearn.svm import LinearSVC
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from tqdm import tqdm
from tqdm.auto import trange

from baselines.baseline_models import NodeClassificationOutputModule, GraphClassificationOutputModule, LinkPredictionOutputModule
import baselines.bl_utils as bl_ut
from baselines.bl_utils import EpochStats
from baselines.one_task_gcn import prepare_batch_for_task
import data_utils
import utils as ut


def generate_intermediate_data(og_model, dataset, device="cpu"):
    """In each Data object in the dataset adds a new field called node_embeddings, containing the node embeddings
    generated by og_model"""
    og_model.eval()
    new_data = []
    with torch.no_grad():
        for data in dataset:
            data = data.to(device)
            node_emebddings = og_model(data, return_embeddings=True)
            data = data.to("cpu")
            data.node_embeddings = node_emebddings
            new_data.append(data)
    return new_data


def get_baseline_nn_output_model(output_task, embedding_dim, hidden_dim, num_classes=None):
    if output_task == "gc":
        return GraphClassificationOutputModule(embedding_dim, hidden_dim, num_classes)
    elif output_task == "nc":
        return NodeClassificationOutputModule(embedding_dim, num_classes)
    elif output_task == "lp":
        return LinkPredictionOutputModule(embedding_dim)


def train_baseline_nn_output_model(output_model, dataloader, output_task, epochs, lr, early_stopping=False, es_tmpdir=None, val_dataloader=None, device="cpu"):
    output_model.train()
    optimizer = torch.optim.Adam(output_model.parameters(), lr=lr)

    if early_stopping:
        best_val_score = 0
        if not es_tmpdir:
            es_tmpdir = "emb_to_"+output_task+"_bst_early_stopping_tmp"
    for epoch in trange(epochs, desc="Epoch"):
        epoch_stats = EpochStats()
        for batch_idx, batch in enumerate(tqdm(dataloader, desc="Train Batch")):
            optimizer.zero_grad()

            batch = prepare_batch_for_task(batch, output_task, train=True)
            batch = batch.to(device)

            # Forward pass 
            if output_task == "gc":
                train_logit = output_model(batch.node_embeddings, batch.batch)
            elif output_task == "nc":
                train_logit = output_model(batch.node_embeddings)
            elif output_task == "lp":
                train_logit = output_model(batch.node_embeddings, batch.pos_edge_index, batch.neg_edge_index)                

            # Evaluate Loss and Accuracy
            if output_task == "gc":
                loss = F.cross_entropy(train_logit, batch.y)
                with torch.no_grad():
                    acc = ut.get_accuracy(train_logit, batch.y)
            elif output_task == "nc":
                node_labels = batch.node_y.argmax(1)
                train_mask = batch.train_mask.squeeze()
                loss = F.cross_entropy(train_logit[train_mask==1], node_labels[train_mask==1])
                with torch.no_grad():
                    acc = ut.get_accuracy(train_logit[train_mask==1], node_labels[train_mask==1])
            elif output_task == "lp":
                train_link_labels = data_utils.get_link_labels(batch.pos_edge_index, batch.neg_edge_index)
                loss = F.binary_cross_entropy_with_logits(train_logit.squeeze(), train_link_labels)
                with torch.no_grad():
                    train_labels = train_link_labels.detach().cpu().numpy()
                    train_predictions = train_logit.detach().cpu().numpy()
                    acc = roc_auc_score(train_labels, train_predictions.squeeze())

            epoch_stats.update(output_task, batch, loss, acc, True)
            
            # Backprop and update parameters
            loss.backward()
            optimizer.step()
            
        if early_stopping and epoch > 5 and epoch%5 == 0:
            model_copy = copy.deepcopy(output_model)
            tqdm.write("\nTest on Validation Set")
            val_stats = eval_baseline_nn_output_model(model_copy, val_dataloader, output_task, device=device)
            epoch_acc = val_stats[output_task]["acc"]
            if epoch_acc > best_val_score:
                best_val_score = epoch_acc
                model_copy.to("cpu")
                args = type('', (), {})()
                args.early_stopping_stats = val_stats # so it save them in file
                args.early_stopping_epoch_acc = epoch_acc
                args.early_stopping_epoch = epoch
                ut.save_model(model_copy, es_tmpdir, "best_val", args)

        task_epoch_stats = epoch_stats.get_average_stats()
        bl_ut.print_train_epoch_stats(epoch, task_epoch_stats)

    if early_stopping:
        ut.recover_early_stopping_best_weights(output_model, es_tmpdir)
       

def eval_baseline_nn_output_model(output_model, dataloader, output_task, device="cpu"):
    output_model.eval()
    epoch_stats = EpochStats()
    for batch_idx, batch in enumerate(tqdm(dataloader, desc="Eval Batch")):
        batch = prepare_batch_for_task(batch, output_task, train=False)
        batch = batch.to(device)
        with torch.no_grad():
            # Forward pass 
            if output_task == "gc":
                test_logit = output_model(batch.node_embeddings, batch.batch)
            elif output_task == "nc":
                test_logit = output_model(batch.node_embeddings)
            elif output_task == "lp":
                test_logit = output_model(batch.node_embeddings, batch.pos_edge_index, batch.neg_edge_index)                

            # Evaluate Loss and Accuracy
            if output_task == "gc":
                loss = F.cross_entropy(test_logit, batch.y)
                with torch.no_grad():
                    acc = ut.get_accuracy(test_logit, batch.y)
            elif output_task == "nc":
                node_labels = batch.node_y.argmax(1)
                train_mask = batch.train_mask.squeeze()
                test_mask = (train_mask==0).float()
                loss = F.cross_entropy(test_logit[test_mask==1], node_labels[test_mask==1])
                with torch.no_grad():
                    acc = ut.get_accuracy(test_logit[test_mask==1], node_labels[test_mask==1])
            elif output_task == "lp":
                test_link_labels = data_utils.get_link_labels(batch.pos_edge_index, batch.neg_edge_index)
                loss = F.binary_cross_entropy_with_logits(test_logit.squeeze(), test_link_labels)
                with torch.no_grad():
                    test_labels = test_link_labels.detach().cpu().numpy()
                    test_predictions = test_logit.detach().cpu().numpy()
                    acc = roc_auc_score(test_labels, test_predictions.squeeze())

            epoch_stats.update(output_task, batch, loss, acc, False)

    task_test_stats = epoch_stats.get_average_stats()
    bl_ut.print_test_stats(task_test_stats)
    return task_test_stats


def get_data_for_linear_classifier(data, task, shuffle=True):
    X = []
    y = []
    for d in data:
        node_embeddings = d.node_embeddings.detach().cpu().numpy()
        #print(node_embeddings.shape)
        if task == "gc":
            X.append(node_embeddings.mean(axis=0))   
            y.append(d.y.detach().cpu().numpy())
        elif task == "nc":
            #train_mask = d.train_mask.detach().cpu().numpy()
            X.append(node_embeddings)
            node_labels = d.node_y.argmax(1).detach().cpu().numpy() 
            node_labels = np.expand_dims(node_labels, axis=1)
            y.append(node_labels) 
        elif task == "lp":
            train_data_list, test_data_list = data_utils.prepare_data_for_link_prediction([d], 
                                                                                          train_ratio=0.9,
                                                                                          neg_to_pos_edge_ratio=1,
                                                                                          rnd_labeled_edges=False)

            pos_edge_idx = test_data_list[0].pos_edge_index.detach().cpu().numpy()
            neg_edge_idx = test_data_list[0].neg_edge_index.detach().cpu().numpy()
            lp_labels = data_utils.get_link_labels(test_data_list[0].pos_edge_index, test_data_list[0].neg_edge_index).detach().cpu().numpy()
 
            node_a = np.take(node_embeddings, np.concatenate((pos_edge_idx[0], neg_edge_idx[0])), axis=0)
            node_b = np.take(node_embeddings, np.concatenate((pos_edge_idx[1], neg_edge_idx[1])), axis=0)
            X.append(np.concatenate((node_a, node_b), axis=1))
            y.append(np.expand_dims(lp_labels, axis=1))
    
    X = np.vstack(X)
    y = np.vstack(y)
    if shuffle:
        perm = np.arange(X.shape[0])
        np.random.shuffle(perm)
        X = X[perm]
        y = y[perm]
    return X, y
 

def test_linear_classifier(task, train_data, val_data, test_data):
    train_X, train_y = get_data_for_linear_classifier(train_data, task)
    #train_X, train_y = get_data_for_linear_classifier(train_data+val_data, task)
    test_X, test_y = get_data_for_linear_classifier(test_data, task)
    clf = make_pipeline(StandardScaler(), LinearSVC(random_state=0, tol=1e-5))    
    clf.fit(train_X, train_y.ravel())
    return clf.score(test_X, test_y.ravel())


def run_test(og_model, train_dataset, val_dataset, test_dataset, epochs, batch_size, lr, embedding_dim, hidden_dim, early_stopping=False, es_tmpdir=None, output_folder=None, device="cpu"):
    if es_tmpdir:
        es_tmpdir = es_tmpdir+"baseline"

    train_data = generate_intermediate_data(og_model, train_dataset, device)
    train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
    val_data = generate_intermediate_data(og_model, val_dataset, device)
    val_dataloader = DataLoader(val_data, batch_size=batch_size, shuffle=False)
    test_data = generate_intermediate_data(og_model, test_dataset, device)
    test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False)

    task_test_stats = {}
    results_str = "################## Test Embeddings ##################\n\n\n"
    for output_task in ["gc", "nc", "lp"]:
        results_str += f"######### Task: {output_task} #########\n"
        task_test_stats[output_task] = []

        results_str += "### Linear SVM ###\n"
        linear_svm_acc = test_linear_classifier(output_task, train_data, val_data, test_data)
        task_test_stats[output_task].append(linear_svm_acc)
        results_str += f"Acc: {linear_svm_acc:.4f}\n\n"

        if output_task == "gc":
            num_classes = train_dataset.num_classes
        elif output_task == "nc":
            num_classes = train_dataset[0].node_y.size(1)
        elif output_task == "lp":
            num_classes = 1

        results_str += "### Trained from scratch baseline ###\n"
        output_model = get_baseline_nn_output_model(output_task, embedding_dim, hidden_dim, num_classes)
        output_model = output_model.to(device)
        train_baseline_nn_output_model(output_model,
                                       train_dataloader,
                                       output_task,
                                       epochs,
                                       lr,
                                       early_stopping=early_stopping,
                                       es_tmpdir=es_tmpdir,
                                       val_dataloader=val_dataloader,
                                       device=device) 
        test_stats = eval_baseline_nn_output_model(output_model,
                                                   test_dataloader,
                                                   output_task,
                                                   device=device)
        task_test_stats[output_task].append(test_stats)
        results_str += f"Acc: {test_stats[output_task]['acc']:.4f}\n\n"
        if output_folder:
            ut.save_model(output_model, output_folder, "trained_baseline_nn")

        if og_model.name in ["Meta_MultitaskGCN_MAML", "Meta_MultitaskGCN_ANIL", "Baseline_MultitaskGCN"]:
            results_str += "### Eval baseline from trained multitask model ###\n"
            if "Meta" in og_model.name:
                output_layer_state_dict = getattr(og_model.output_layer, output_task+"_output_layer").state_dict()
            else:
                output_layer_state_dict = getattr(og_model, output_task+"_output_layer").state_dict()
            output_model.load_state_dict(output_layer_state_dict)
            test_stats = eval_baseline_nn_output_model(output_model,
                                                       test_dataloader,
                                                       output_task,
                                                       device=device)
            task_test_stats[output_task].append(test_stats)
            results_str += f"Acc: {test_stats[output_task]['acc']:.4f}\n\n"
            if output_folder:
                ut.save_model(output_model, output_folder, "trained_multitask_ol")

            results_str += "### Finetune from trained multitask model ###\n"
            train_baseline_nn_output_model(output_model,
                                           train_dataloader,
                                           output_task,
                                           epochs,
                                           lr,
                                           early_stopping=early_stopping,
                                           es_tmpdir=es_tmpdir,
                                           val_dataloader=val_dataloader,
                                           device=device) 
            test_stats = eval_baseline_nn_output_model(output_model,
                                                       test_dataloader,
                                                       output_task,
                                                       device=device)
            task_test_stats[output_task].append(test_stats)
            results_str += f"Acc: {test_stats[output_task]['acc']:.4f}\n\n"
            if output_folder:
                ut.save_model(output_model, output_folder, "finetuned_multitask_ol")
        results_str += "\n"
    print(results_str)
    return task_test_stats
