# -*- coding: utf-8 -*-
from sklearn import preprocessing
from torch.autograd import Variable
from models_gat import GAT

import os
import torch
import numpy as np
import argparse
import pickle
import sklearn.metrics as metrics
import cross_val
import time
import random


torch.manual_seed(0)
np.random.seed(0)
random.seed(0)

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

def evaluate(dataset, model_GAT, args, threshold_value, model_name):
    """
    Parameters
    ----------
    dataset : dataloader (dataloader for the validation/test dataset).
    model_GCN : nn model (GAT model).
    args : arguments
    threshold_value : float (threshold for adjacency matrices).
    
    Description
    ----------
    This methods performs the evaluation of the model on test/validation dataset
    
    Returns
    -------
    test accuracy.
    """
    model_GAT.eval()
    labels = []
    preds = []
    for batch_idx, data in enumerate(dataset):
        adj = Variable(data['adj'].float(), requires_grad=False).to(device)
        labels.append(data['label'].long().numpy())
        
        adj = torch.squeeze(adj)
        
        features = np.identity(adj.shape[0])
        
        features = Variable(torch.from_numpy(features).float(), requires_grad=False).cpu()
        if args.threshold in ["median", "mean"]:
            adj = torch.where(adj > threshold_value, torch.tensor([1.0]), torch.tensor([0.0]))
        
        ypred = model_GAT(features, adj)

        _, indices = torch.max(ypred, 1)
        preds.append(indices.cpu().data.numpy())

    labels = np.hstack(labels)
    preds = np.hstack(preds)
    simple_r = {'labels':labels,'preds':preds}

    with open("./gat/Labels_and_preds/"+model_name+".pickle", 'wb') as f:
      pickle.dump(simple_r, f)

    result = {'prec': metrics.precision_score(labels, preds, average='macro'),
              'recall': metrics.recall_score(labels, preds, average='macro'),
              'acc': metrics.accuracy_score(labels, preds),
              'F1': metrics.f1_score(labels, preds, average="micro")}
    if args.evaluation_method == 'model assessment':
        name = 'Test'
    if args.evaluation_method == 'model selection':
        name = 'Validation'
    print(name, " accuracy:", result['acc'])
    return result['acc']

def minmax_sc(x):
    min_max_scaler = preprocessing.MinMaxScaler()
    x = min_max_scaler.fit_transform(x)
    return x

def train(args, train_dataset, val_dataset, model_GAT, threshold_value, model_name):
    """
    Parameters
    ----------
    args : arguments
    train_dataset : dataloader (dataloader for the validation/test dataset).
    val_dataset : dataloader (dataloader for the validation/test dataset).
    model_GAT : nn model (GAT model).
    threshold_value : float (threshold for adjacency matrices).
    
    Description
    ----------
    This methods performs the training of the model on train dataset and calls evaluate() method for evaluation.
    
    Returns
    -------
    test accuracy.
    """
    params = list(model_GAT.parameters()) 
    optimizer = torch.optim.Adam(params, lr=args.lr, weight_decay=args.weight_decay)
    test_accs = []
    train_loss=[]
    val_acc=[]    

    for epoch in range(args.num_epochs):
        print("Epoch ",epoch)
        
        print("Size of Training Set:" + str(len(train_dataset)))
        print("Size of Validation Set:" + str(len(val_dataset)))    
        model_GAT.train()
        total_time = 0
        avg_loss = 0.0
        
        preds = []
        labels = []
        for batch_idx, data in enumerate(train_dataset):
            begin_time = time.time()
            
            adj = Variable(data['adj'].float(), requires_grad=False).to(device)
            label = Variable(data['label'].long()).to(device)
            #adj_id = Variable(data['id'].int()).to(device)
            
            adj = torch.squeeze(adj)

            features = np.identity(adj.shape[0])
            features = Variable(torch.from_numpy(features).float(), requires_grad=False).cpu()
            
            if args.threshold in ["median", "mean"]:
                adj = torch.where(adj > threshold_value, torch.tensor([1.0]), torch.tensor([0.0]))
            
            
            ypred = model_GAT(features, adj)
            
            _, indices = torch.max(ypred, 1)
            preds.append(indices.cpu().data.numpy())
            labels.append(data['label'].long().numpy())
            loss = model_GAT.loss(ypred, label)
            model_GAT.zero_grad()
            loss.backward()
            #nn.utils.clip_grad_norm_(model_DIFFPOOL.parameters(), args.clip)
            optimizer.step()
            
            avg_loss += loss
            elapsed = time.time() - begin_time
            total_time += elapsed
        
        if epoch == args.num_epochs-1:
            model_GAT.is_trained = True
        
        preds = np.hstack(preds)
        labels = np.hstack(labels)
        print("Train accuracy : ", np.mean( preds == labels ))
        test_acc = evaluate(val_dataset, model_GAT, args, threshold_value, model_name)
        print('Avg loss: ', avg_loss, '; epoch time: ', total_time)
        test_accs.append(test_acc)
        train_loss.append(avg_loss)
        val_acc.append(test_acc)
        

    path = './gat/weights/W_'+model_name+'.pickle'
    
    if os.path.exists(path):
        os.remove(path)
        
    os.rename('GAT_W.pickle',path)
  
    los_p = {'loss':train_loss}
    with open("./gat/training_loss/Training_loss_"+model_name+".pickle", 'wb') as f:
      pickle.dump(los_p, f)
    torch.save(model_GAT,"./gat/models/GAT_"+model_name+".pt")

    return test_acc

def load_data(args):
    """
    Parameters
    ----------
    args : arguments
    Description
    ----------
    This methods loads the adjacency matrices representing the args.view -th view in dataset
    
    Returns
    -------
    List of dictionaries{adj, label, id}
    """
    #Load graphs and labels
    with open('data/'+args.dataset+'/'+args.dataset+'_edges','rb') as f:
        multigraphs = pickle.load(f)        
    with open('data/'+args.dataset+'/'+args.dataset+'_labels','rb') as f:
        labels = pickle.load(f)
    adjacencies = [multigraphs[i][:,:,args.view] for i in range(len(multigraphs))]
    #Normalize inputs
    if args.NormalizeInputGraphs==True:
        for subject in range(len(adjacencies)):
            adjacencies[subject] = minmax_sc(adjacencies[subject])
    
    #Create List of Dictionaries
    G_list=[]
    for i in range(len(labels)):
        G_element = {"adj":   adjacencies[i],"label": labels[i],"id":  i,}
        G_list.append(G_element)
    return G_list

def arg_parse(dataset, view, num_shots=2, cv_number=5):
    """
    arguments definition method
    """
    parser = argparse.ArgumentParser(description='Graph Classification')
    
    
    
    parser.add_argument('--mode', type=str, default='train', choices=['train', 'test'])
    parser.add_argument('--v', type=str, default=1)
    parser.add_argument('--data', type=str, default='Sample_dataset', choices = [ f.path[5:] for f in os.scandir("data") if f.is_dir() ])
    
    
    
    
    parser.add_argument('--dataset', type=str, default=dataset,
                        help='Dataset')
    parser.add_argument('--view', type=int, default=view,
                        help = 'view index in the dataset')
    parser.add_argument('--num_epochs', type=int, default=1, #50
                        help='Training Epochs')
    parser.add_argument('--num_shots', type=int, default=num_shots, #100
                        help='number of shots')
    parser.add_argument('--cv_number', type=int, default=cv_number,
                        help='number of validation folds.')
    parser.add_argument('--NormalizeInputGraphs', default=False, action='store_true',
                        help='Normalize Input adjacency matrices of graphs')
    parser.add_argument('--evaluation_method', type=str, default='model assessment',
                        help='evaluation method, possible values : model selection, model assessment')
    parser.add_argument('--threshold', dest='threshold', default='mean',
                        help='threshold the graph adjacency matrix. Possible values: no_threshold, median, mean')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='Disables CUDA training.')
    parser.add_argument('--num-classes', dest='num_classes', type=int, default=2,
                        help='Number of label classes')
    parser.add_argument('--lr', type=float, default=0.001,
                        help='Initial learning rate.')
    parser.add_argument('--weight_decay', type=float, default=5e-4,
                        help='Weight decay (L2 loss on parameters).')
    parser.add_argument('--hidden', type=int, default=8, 
                        help='Number of hidden units.')
    parser.add_argument('--nb_heads', type=int, default=8, 
                        help='Number of head attentions.')
    parser.add_argument('--dropout', type=float, default=0.8, 
                        help='Dropout rate (1 - keep probability).')
    parser.add_argument('--alpha', type=float, default=0.2, 
                        help='Alpha for the leaky_relu.')
    
    return parser.parse_args()

def benchmark_task(args, model_name):
    """
    Parameters
    ----------
    args : Arguments
    Description
    ----------
    Initiates the model and performs train/test or train/validation splits and calls train() to execute training and evaluation.
    Returns
    -------
    test_accs : test accuracies (list)

    """
    G_list = load_data(args)
    num_nodes = G_list[0]['adj'].shape[0]
    test_accs = []
    folds = cross_val.stratify_splits(G_list,args)
    [random.shuffle(folds[i]) for i in range(len(folds))]
    for i in range(args.cv_number):
        train_set, validation_set, test_set = cross_val.datasets_splits(folds, args, i)
        
        if args.evaluation_method =='model selection':
            train_dataset, val_dataset, threshold_value = cross_val.model_selection_split(train_set, validation_set, args)
        
        if args.evaluation_method =='model assessment':
            train_dataset, val_dataset, threshold_value = cross_val.model_assessment_split(train_set, validation_set, test_set, args)

        print("CV : ",i)
        model_GAT = GAT(nfeat=num_nodes, 
                nhid=args.hidden, 
                nclass=args.num_classes, 
                dropout=args.dropout, 
                nheads=args.nb_heads, 
                alpha=args.alpha)
        
        test_acc = train(args, train_dataset, val_dataset, model_GAT, threshold_value, model_name+"_CV_"+str(i)+"_view_"+str(args.view))
        test_accs.append(test_acc)
    return test_accs

def test_scores(dataset, view, model_name, cv_number):
    
    args = arg_parse(dataset, view, cv_number=cv_number)
    print("Main : ",args)
    test_accs = benchmark_task(args, model_name)
    print("test accuracies ",test_accs)
    return test_accs
    

def two_shot_trainer(dataset, view, num_shots):
    args = arg_parse(dataset, view, num_shots=num_shots)
    torch.manual_seed(0)
    np.random.seed(0)
    random.seed(0)  
    start = time.time()
    
    for i in range(args.num_shots):
        model = "gat"
        model_name = "Few_Shot_"+dataset+"_"+model + str(i)
        print("Shot : ",i)
        with open('./Two_shot_samples_views/'+dataset+'_view_'+str(view)+'_shot_'+str(i)+'_train','rb') as f:
            train_set = pickle.load(f)
        with open('./Two_shot_samples_views/'+dataset+'_view_'+str(view)+'_shot_'+str(i)+'_test','rb') as f:
            test_set = pickle.load(f)
        
        num_nodes = train_set[0]['adj'].shape[0]
        
        model_GAT = GAT(nfeat=num_nodes, 
                nhid=args.hidden, 
                nclass=args.num_classes, 
                dropout=args.dropout, 
                nheads=args.nb_heads, 
                alpha=args.alpha)
        
        train_dataset, val_dataset, threshold_value = cross_val.two_shot_loader(train_set, test_set, args)
        
        test_acc = train(args, train_dataset, val_dataset, model_GAT, threshold_value, model_name+"_view_"+str(view))
        
        print("Test accuracy:"+str(test_acc))
        print('load data using ------>', time.time()-start)
[((17, 0, 17, 20), 'torch.manual_seed', 'torch.manual_seed', ({(17, 18, 17, 19): '(0)'}, {}), '(0)', False, 'import torch\n'), ((18, 0, 18, 17), 'numpy.random.seed', 'np.random.seed', ({(18, 15, 18, 16): '(0)'}, {}), '(0)', True, 'import numpy as np\n'), ((19, 0, 19, 14), 'random.seed', 'random.seed', ({(19, 12, 19, 13): '(0)'}, {}), '(0)', False, 'import random\n'), ((60, 13, 60, 30), 'numpy.hstack', 'np.hstack', ({(60, 23, 60, 29): 'labels'}, {}), '(labels)', True, 'import numpy as np\n'), ((61, 12, 61, 28), 'numpy.hstack', 'np.hstack', ({(61, 22, 61, 27): 'preds'}, {}), '(preds)', True, 'import numpy as np\n'), ((79, 21, 79, 49), 'sklearn.preprocessing.MinMaxScaler', 'preprocessing.MinMaxScaler', ({}, {}), '()', False, 'from sklearn import preprocessing\n'), ((102, 16, 102, 84), 'torch.optim.Adam', 'torch.optim.Adam', (), '', False, 'import torch\n'), ((164, 7, 164, 27), 'os.path.exists', 'os.path.exists', ({(164, 22, 164, 26): 'path'}, {}), '(path)', False, 'import os\n'), ((167, 4, 167, 34), 'os.rename', 'os.rename', ({(167, 14, 167, 28): '"""GAT_W.pickle"""', (167, 29, 167, 33): 'path'}, {}), "('GAT_W.pickle', path)", False, 'import os\n'), ((172, 4, 172, 62), 'torch.save', 'torch.save', ({(172, 15, 172, 24): 'model_GAT', (172, 25, 172, 61): "('./gat/models/GAT_' + model_name + '.pt')"}, {}), "(model_GAT, './gat/models/GAT_' + model_name + '.pt')", False, 'import torch\n'), ((211, 13, 211, 72), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((273, 12, 273, 50), 'cross_val.stratify_splits', 'cross_val.stratify_splits', ({(273, 38, 273, 44): 'G_list', (273, 45, 273, 49): 'args'}, {}), '(G_list, args)', False, 'import cross_val\n'), ((307, 4, 307, 24), 'torch.manual_seed', 'torch.manual_seed', ({(307, 22, 307, 23): '(0)'}, {}), '(0)', False, 'import torch\n'), ((308, 4, 308, 21), 'numpy.random.seed', 'np.random.seed', ({(308, 19, 308, 20): '(0)'}, {}), '(0)', True, 'import numpy as np\n'), ((309, 4, 309, 18), 'random.seed', 'random.seed', ({(309, 16, 309, 17): '(0)'}, {}), '(0)', False, 'import random\n'), ((310, 12, 310, 23), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((21, 34, 21, 59), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((47, 14, 47, 32), 'torch.squeeze', 'torch.squeeze', ({(47, 28, 47, 31): 'adj'}, {}), '(adj)', False, 'import torch\n'), ((49, 19, 49, 44), 'numpy.identity', 'np.identity', ({(49, 31, 49, 43): 'adj.shape[0]'}, {}), '(adj.shape[0])', True, 'import numpy as np\n'), ((57, 21, 57, 40), 'torch.max', 'torch.max', ({(57, 31, 57, 36): 'ypred', (57, 38, 57, 39): '1'}, {}), '(ypred, 1)', False, 'import torch\n'), ((65, 6, 65, 30), 'pickle.dump', 'pickle.dump', ({(65, 18, 65, 26): 'simple_r', (65, 28, 65, 29): 'f'}, {}), '(simple_r, f)', False, 'import pickle\n'), ((67, 22, 67, 77), 'sklearn.metrics.precision_score', 'metrics.precision_score', (), '', True, 'import sklearn.metrics as metrics\n'), ((68, 24, 68, 76), 'sklearn.metrics.recall_score', 'metrics.recall_score', (), '', True, 'import sklearn.metrics as metrics\n'), ((69, 21, 69, 58), 'sklearn.metrics.accuracy_score', 'metrics.accuracy_score', ({(69, 44, 69, 50): 'labels', (69, 52, 69, 57): 'preds'}, {}), '(labels, preds)', True, 'import sklearn.metrics as metrics\n'), ((70, 20, 70, 68), 'sklearn.metrics.f1_score', 'metrics.f1_score', (), '', True, 'import sklearn.metrics as metrics\n'), ((152, 16, 152, 32), 'numpy.hstack', 'np.hstack', ({(152, 26, 152, 31): 'preds'}, {}), '(preds)', True, 'import numpy as np\n'), ((153, 17, 153, 34), 'numpy.hstack', 'np.hstack', ({(153, 27, 153, 33): 'labels'}, {}), '(labels)', True, 'import numpy as np\n'), ((165, 8, 165, 23), 'os.remove', 'os.remove', ({(165, 18, 165, 22): 'path'}, {}), '(path)', False, 'import os\n'), ((171, 6, 171, 27), 'pickle.dump', 'pickle.dump', ({(171, 18, 171, 23): 'los_p', (171, 25, 171, 26): 'f'}, {}), '(los_p, f)', False, 'import pickle\n'), ((191, 22, 191, 36), 'pickle.load', 'pickle.load', ({(191, 34, 191, 35): 'f'}, {}), '(f)', False, 'import pickle\n'), ((193, 17, 193, 31), 'pickle.load', 'pickle.load', ({(193, 29, 193, 30): 'f'}, {}), '(f)', False, 'import pickle\n'), ((274, 5, 274, 29), 'random.shuffle', 'random.shuffle', ({(274, 20, 274, 28): 'folds[i]'}, {}), '(folds[i])', False, 'import random\n'), ((276, 46, 276, 87), 'cross_val.datasets_splits', 'cross_val.datasets_splits', ({(276, 72, 276, 77): 'folds', (276, 79, 276, 83): 'args', (276, 85, 276, 86): 'i'}, {}), '(folds, args, i)', False, 'import cross_val\n'), ((285, 20, 290, 33), 'models_gat.GAT', 'GAT', (), '', False, 'from models_gat import GAT\n'), ((323, 20, 328, 33), 'models_gat.GAT', 'GAT', (), '', False, 'from models_gat import GAT\n'), ((330, 54, 330, 106), 'cross_val.two_shot_loader', 'cross_val.two_shot_loader', ({(330, 80, 330, 89): 'train_set', (330, 91, 330, 99): 'test_set', (330, 101, 330, 105): 'args'}, {}), '(train_set, test_set, args)', False, 'import cross_val\n'), ((119, 25, 119, 36), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((125, 18, 125, 36), 'torch.squeeze', 'torch.squeeze', ({(125, 32, 125, 35): 'adj'}, {}), '(adj)', False, 'import torch\n'), ((127, 23, 127, 48), 'numpy.identity', 'np.identity', ({(127, 35, 127, 47): 'adj.shape[0]'}, {}), '(adj.shape[0])', True, 'import numpy as np\n'), ((136, 25, 136, 44), 'torch.max', 'torch.max', ({(136, 35, 136, 40): 'ypred', (136, 42, 136, 43): '1'}, {}), '(ypred, 1)', False, 'import torch\n'), ((154, 35, 154, 61), 'numpy.mean', 'np.mean', ({(154, 44, 154, 59): '(preds == labels)'}, {}), '(preds == labels)', True, 'import numpy as np\n'), ((279, 58, 279, 122), 'cross_val.model_selection_split', 'cross_val.model_selection_split', ({(279, 90, 279, 99): 'train_set', (279, 101, 279, 115): 'validation_set', (279, 117, 279, 121): 'args'}, {}), '(train_set, validation_set, args)', False, 'import cross_val\n'), ((282, 58, 282, 133), 'cross_val.model_assessment_split', 'cross_val.model_assessment_split', ({(282, 91, 282, 100): 'train_set', (282, 102, 282, 116): 'validation_set', (282, 118, 282, 126): 'test_set', (282, 128, 282, 132): 'args'}, {}), '(train_set, validation_set, test_set, args)', False, 'import cross_val\n'), ((317, 24, 317, 38), 'pickle.load', 'pickle.load', ({(317, 36, 317, 37): 'f'}, {}), '(f)', False, 'import pickle\n'), ((319, 23, 319, 37), 'pickle.load', 'pickle.load', ({(319, 35, 319, 36): 'f'}, {}), '(f)', False, 'import pickle\n'), ((53, 53, 53, 72), 'torch.tensor', 'torch.tensor', ({(53, 66, 53, 71): '[1.0]'}, {}), '([1.0])', False, 'import torch\n'), ((53, 74, 53, 93), 'torch.tensor', 'torch.tensor', ({(53, 87, 53, 92): '[0.0]'}, {}), '([0.0])', False, 'import torch\n'), ((146, 22, 146, 33), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((335, 41, 335, 52), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((131, 57, 131, 76), 'torch.tensor', 'torch.tensor', ({(131, 70, 131, 75): '[1.0]'}, {}), '([1.0])', False, 'import torch\n'), ((131, 78, 131, 97), 'torch.tensor', 'torch.tensor', ({(131, 91, 131, 96): '[0.0]'}, {}), '([0.0])', False, 'import torch\n'), ((217, 102, 217, 120), 'os.scandir', 'os.scandir', ({(217, 113, 217, 119): '"""data"""'}, {}), "('data')", False, 'import os\n'), ((51, 28, 51, 54), 'torch.from_numpy', 'torch.from_numpy', ({(51, 45, 51, 53): 'features'}, {}), '(features)', False, 'import torch\n'), ((128, 32, 128, 58), 'torch.from_numpy', 'torch.from_numpy', ({(128, 49, 128, 57): 'features'}, {}), '(features)', False, 'import torch\n')]