# -*- 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)
[((273, 293), 'torch.manual_seed', 'torch.manual_seed', (['(0)'], {}), '(0)\n', (290, 293), False, 'import torch\n'), ((294, 311), 'numpy.random.seed', 'np.random.seed', (['(0)'], {}), '(0)\n', (308, 311), True, 'import numpy as np\n'), ((312, 326), 'random.seed', 'random.seed', (['(0)'], {}), '(0)\n', (323, 326), False, 'import random\n'), ((1597, 1614), 'numpy.hstack', 'np.hstack', (['labels'], {}), '(labels)\n', (1606, 1614), True, 'import numpy as np\n'), ((1627, 1643), 'numpy.hstack', 'np.hstack', (['preds'], {}), '(preds)\n', (1636, 1643), True, 'import numpy as np\n'), ((2350, 2378), 'sklearn.preprocessing.MinMaxScaler', 'preprocessing.MinMaxScaler', ([], {}), '()\n', (2376, 2378), False, 'from sklearn import preprocessing\n'), ((3095, 3163), 'torch.optim.Adam', 'torch.optim.Adam', (['params'], {'lr': 'args.lr', 'weight_decay': 'args.weight_decay'}), '(params, lr=args.lr, weight_decay=args.weight_decay)\n', (3111, 3163), False, 'import torch\n'), ((5362, 5382), 'os.path.exists', 'os.path.exists', (['path'], {}), '(path)\n', (5376, 5382), False, 'import os\n'), ((5421, 5452), 'os.rename', 'os.rename', (['"""GAT_W.pickle"""', 'path'], {}), "('GAT_W.pickle', path)\n", (5430, 5452), False, 'import os\n'), ((5604, 5667), 'torch.save', 'torch.save', (['model_GAT', "('./gat/models/GAT_' + model_name + '.pt')"], {}), "(model_GAT, './gat/models/GAT_' + model_name + '.pt')\n", (5614, 5667), False, 'import torch\n'), ((6808, 6867), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Graph Classification"""'}), "(description='Graph Classification')\n", (6831, 6867), False, 'import argparse\n'), ((9774, 9813), 'cross_val.stratify_splits', 'cross_val.stratify_splits', (['G_list', 'args'], {}), '(G_list, args)\n', (9799, 9813), False, 'import cross_val\n'), ((11209, 11229), 'torch.manual_seed', 'torch.manual_seed', (['(0)'], {}), '(0)\n', (11226, 11229), False, 'import torch\n'), ((11234, 11251), 'numpy.random.seed', 'np.random.seed', (['(0)'], {}), '(0)\n', (11248, 11251), True, 'import numpy as np\n'), ((11256, 11270), 'random.seed', 'random.seed', (['(0)'], {}), '(0)\n', (11267, 11270), False, 'import random\n'), ((11285, 11296), 'time.time', 'time.time', ([], {}), '()\n', (11294, 11296), False, 'import time\n'), ((362, 387), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (385, 387), False, 'import torch\n'), ((1125, 1143), 'torch.squeeze', 'torch.squeeze', (['adj'], {}), '(adj)\n', (1138, 1143), False, 'import torch\n'), ((1172, 1197), 'numpy.identity', 'np.identity', (['adj.shape[0]'], {}), '(adj.shape[0])\n', (1183, 1197), True, 'import numpy as np\n'), ((1514, 1533), 'torch.max', 'torch.max', (['ypred', '(1)'], {}), '(ypred, 1)\n', (1523, 1533), False, 'import torch\n'), ((1772, 1796), 'pickle.dump', 'pickle.dump', (['simple_r', 'f'], {}), '(simple_r, f)\n', (1783, 1796), False, 'import pickle\n'), ((1820, 1875), 'sklearn.metrics.precision_score', 'metrics.precision_score', (['labels', 'preds'], {'average': '"""macro"""'}), "(labels, preds, average='macro')\n", (1843, 1875), True, 'import sklearn.metrics as metrics\n'), ((1901, 1953), 'sklearn.metrics.recall_score', 'metrics.recall_score', (['labels', 'preds'], {'average': '"""macro"""'}), "(labels, preds, average='macro')\n", (1921, 1953), True, 'import sklearn.metrics as metrics\n'), ((1976, 2013), 'sklearn.metrics.accuracy_score', 'metrics.accuracy_score', (['labels', 'preds'], {}), '(labels, preds)\n', (1998, 2013), True, 'import sklearn.metrics as metrics\n'), ((2035, 2083), 'sklearn.metrics.f1_score', 'metrics.f1_score', (['labels', 'preds'], {'average': '"""micro"""'}), "(labels, preds, average='micro')\n", (2051, 2083), True, 'import sklearn.metrics as metrics\n'), ((4915, 4931), 'numpy.hstack', 'np.hstack', (['preds'], {}), '(preds)\n', (4924, 4931), True, 'import numpy as np\n'), ((4949, 4966), 'numpy.hstack', 'np.hstack', (['labels'], {}), '(labels)\n', (4958, 4966), True, 'import numpy as np\n'), ((5392, 5407), 'os.remove', 'os.remove', (['path'], {}), '(path)\n', (5401, 5407), False, 'import os\n'), ((5578, 5599), 'pickle.dump', 'pickle.dump', (['los_p', 'f'], {}), '(los_p, f)\n', (5589, 5599), False, 'import pickle\n'), ((6090, 6104), 'pickle.load', 'pickle.load', (['f'], {}), '(f)\n', (6101, 6104), False, 'import pickle\n'), ((6204, 6218), 'pickle.load', 'pickle.load', (['f'], {}), '(f)\n', (6215, 6218), False, 'import pickle\n'), ((9818, 9842), 'random.shuffle', 'random.shuffle', (['folds[i]'], {}), '(folds[i])\n', (9832, 9842), False, 'import random\n'), ((9953, 9994), 'cross_val.datasets_splits', 'cross_val.datasets_splits', (['folds', 'args', 'i'], {}), '(folds, args, i)\n', (9978, 9994), False, 'import cross_val\n'), ((10427, 10557), 'models_gat.GAT', 'GAT', ([], {'nfeat': 'num_nodes', 'nhid': 'args.hidden', 'nclass': 'args.num_classes', 'dropout': 'args.dropout', 'nheads': 'args.nb_heads', 'alpha': 'args.alpha'}), '(nfeat=num_nodes, nhid=args.hidden, nclass=args.num_classes, dropout=\n    args.dropout, nheads=args.nb_heads, alpha=args.alpha)\n', (10430, 10557), False, 'from models_gat import GAT\n'), ((11830, 11960), 'models_gat.GAT', 'GAT', ([], {'nfeat': 'num_nodes', 'nhid': 'args.hidden', 'nclass': 'args.num_classes', 'dropout': 'args.dropout', 'nheads': 'args.nb_heads', 'alpha': 'args.alpha'}), '(nfeat=num_nodes, nhid=args.hidden, nclass=args.num_classes, dropout=\n    args.dropout, nheads=args.nb_heads, alpha=args.alpha)\n', (11833, 11960), False, 'from models_gat import GAT\n'), ((12104, 12156), 'cross_val.two_shot_loader', 'cross_val.two_shot_loader', (['train_set', 'test_set', 'args'], {}), '(train_set, test_set, args)\n', (12129, 12156), False, 'import cross_val\n'), ((3637, 3648), 'time.time', 'time.time', ([], {}), '()\n', (3646, 3648), False, 'import time\n'), ((3895, 3913), 'torch.squeeze', 'torch.squeeze', (['adj'], {}), '(adj)\n', (3908, 3913), False, 'import torch\n'), ((3938, 3963), 'numpy.identity', 'np.identity', (['adj.shape[0]'], {}), '(adj.shape[0])\n', (3949, 3963), True, 'import numpy as np\n'), ((4333, 4352), 'torch.max', 'torch.max', (['ypred', '(1)'], {}), '(ypred, 1)\n', (4342, 4352), False, 'import torch\n'), ((5002, 5026), 'numpy.mean', 'np.mean', (['(preds == labels)'], {}), '(preds == labels)\n', (5009, 5026), True, 'import numpy as np\n'), ((10117, 10181), 'cross_val.model_selection_split', 'cross_val.model_selection_split', (['train_set', 'validation_set', 'args'], {}), '(train_set, validation_set, args)\n', (10148, 10181), False, 'import cross_val\n'), ((10305, 10380), 'cross_val.model_assessment_split', 'cross_val.model_assessment_split', (['train_set', 'validation_set', 'test_set', 'args'], {}), '(train_set, validation_set, test_set, args)\n', (10337, 10380), False, 'import cross_val\n'), ((11581, 11595), 'pickle.load', 'pickle.load', (['f'], {}), '(f)\n', (11592, 11595), False, 'import pickle\n'), ((11728, 11742), 'pickle.load', 'pickle.load', (['f'], {}), '(f)\n', (11739, 11742), False, 'import pickle\n'), ((1400, 1419), 'torch.tensor', 'torch.tensor', (['[1.0]'], {}), '([1.0])\n', (1412, 1419), False, 'import torch\n'), ((1421, 1440), 'torch.tensor', 'torch.tensor', (['[0.0]'], {}), '([0.0])\n', (1433, 1440), False, 'import torch\n'), ((4743, 4754), 'time.time', 'time.time', ([], {}), '()\n', (4752, 4754), False, 'import time\n'), ((12380, 12391), 'time.time', 'time.time', ([], {}), '()\n', (12389, 12391), False, 'import time\n'), ((4182, 4201), 'torch.tensor', 'torch.tensor', (['[1.0]'], {}), '([1.0])\n', (4194, 4201), False, 'import torch\n'), ((4203, 4222), 'torch.tensor', 'torch.tensor', (['[0.0]'], {}), '([0.0])\n', (4215, 4222), False, 'import torch\n'), ((7125, 7143), 'os.scandir', 'os.scandir', (['"""data"""'], {}), "('data')\n", (7135, 7143), False, 'import os\n'), ((1235, 1261), 'torch.from_numpy', 'torch.from_numpy', (['features'], {}), '(features)\n', (1251, 1261), False, 'import torch\n'), ((3996, 4022), 'torch.from_numpy', 'torch.from_numpy', (['features'], {}), '(features)\n', (4012, 4022), False, 'import torch\n')]