import argparse

from loader import MoleculeDataset
from torch_geometric.data import DataLoader

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

from tqdm import tqdm
import numpy as np

from model import GNN, GNN_graphpred
from sklearn.metrics import roc_auc_score

from splitters import scaffold_split, random_split
import pandas as pd

import os
import shutil

from tensorboardX import SummaryWriter

criterion = nn.BCEWithLogitsLoss(reduction = "none")

def train(args, model, device, loader, optimizer):
    model.train()

    for step, batch in enumerate(tqdm(loader, desc="Iteration")):
        batch = batch.to(device)
        pred = model(batch.x, batch.edge_index, batch.edge_attr, batch.batch)
        y = batch.y.view(pred.shape).to(torch.float64)

        #Whether y is non-null or not.
        is_valid = y**2 > 0
        #Loss matrix
        loss_mat = criterion(pred.double(), (y+1)/2)
        #loss matrix after removing null target
        loss_mat = torch.where(is_valid, loss_mat, torch.zeros(loss_mat.shape).to(loss_mat.device).to(loss_mat.dtype))
            
        optimizer.zero_grad()
        loss = torch.sum(loss_mat)/torch.sum(is_valid)
        loss.backward()

        optimizer.step()


def eval(args, model, device, loader):
    model.eval()
    y_true = []
    y_scores = []

    for step, batch in enumerate(tqdm(loader, desc="Iteration")):
        batch = batch.to(device)

        with torch.no_grad():
            pred = model(batch.x, batch.edge_index, batch.edge_attr, batch.batch)

        y_true.append(batch.y.view(pred.shape))
        y_scores.append(pred)

    y_true = torch.cat(y_true, dim = 0).cpu().numpy()
    y_scores = torch.cat(y_scores, dim = 0).cpu().numpy()

    roc_list = []
    for i in range(y_true.shape[1]):
        #AUC is only defined when there is at least one positive data.
        if np.sum(y_true[:,i] == 1) > 0 and np.sum(y_true[:,i] == -1) > 0:
            is_valid = y_true[:,i]**2 > 0
            roc_list.append(roc_auc_score((y_true[is_valid,i] + 1)/2, y_scores[is_valid,i]))

    if len(roc_list) < y_true.shape[1]:
        print("Some target is missing!")
        print("Missing ratio: %f" %(1 - float(len(roc_list))/y_true.shape[1]))

    return sum(roc_list)/len(roc_list) #y_true.shape[1]



def main():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch implementation of pre-training of graph neural networks')
    parser.add_argument('--device', type=int, default=0,
                        help='which gpu to use if any (default: 0)')
    parser.add_argument('--batch_size', type=int, default=32,
                        help='input batch size for training (default: 32)')
    parser.add_argument('--epochs', type=int, default=100,
                        help='number of epochs to train (default: 100)')
    parser.add_argument('--lr', type=float, default=0.001,
                        help='learning rate (default: 0.001)')
    parser.add_argument('--lr_scale', type=float, default=1,
                        help='relative learning rate for the feature extraction layer (default: 1)')
    parser.add_argument('--decay', type=float, default=0,
                        help='weight decay (default: 0)')
    parser.add_argument('--num_layer', type=int, default=5,
                        help='number of GNN message passing layers (default: 5).')
    parser.add_argument('--emb_dim', type=int, default=300,
                        help='embedding dimensions (default: 300)')
    parser.add_argument('--dropout_ratio', type=float, default=0.5,
                        help='dropout ratio (default: 0.5)')
    parser.add_argument('--graph_pooling', type=str, default="mean",
                        help='graph level pooling (sum, mean, max, set2set, attention)')
    parser.add_argument('--JK', type=str, default="last",
                        help='how the node features across layers are combined. last, sum, max or concat')
    parser.add_argument('--gnn_type', type=str, default="gin")
    parser.add_argument('--dataset', type=str, default = 'sider', help='root directory of dataset. For now, only classification.')
    parser.add_argument('--input_model_file', type=str, default = '../motif_based_pretrain/saved_model/motif_pretrain.pth', help='filename to read the model (if there is any)')
    parser.add_argument('--filename', type=str, default = '', help='output filename')
    parser.add_argument('--seed', type=int, default=42, help = "Seed for splitting the dataset.")
    parser.add_argument('--runseed', type=int, default=0, help = "Seed for minibatch selection, random initialization.")
    parser.add_argument('--split', type = str, default="scaffold", help = "random or scaffold or random_scaffold")
    parser.add_argument('--eval_train', type=int, default = 1, help='evaluating training or not')
    parser.add_argument('--num_workers', type=int, default = 4, help='number of workers for dataset loading')
    args = parser.parse_args()


    torch.manual_seed(args.runseed)
    np.random.seed(args.runseed)
    device = torch.device("cuda:" + str(args.device)) if torch.cuda.is_available() else torch.device("cpu")
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.runseed)

    #Bunch of classification tasks
    if args.dataset == "tox21":
        num_tasks = 12
    elif args.dataset == "hiv":
        num_tasks = 1
    elif args.dataset == "pcba":
        num_tasks = 128
    elif args.dataset == "muv":
        num_tasks = 17
    elif args.dataset == "bace":
        num_tasks = 1
    elif args.dataset == "bbbp":
        num_tasks = 1
    elif args.dataset == "toxcast":
        num_tasks = 617
    elif args.dataset == "sider":
        num_tasks = 27
    elif args.dataset == "clintox":
        num_tasks = 2
    else:
        raise ValueError("Invalid dataset name.")

    #set up dataset
    dataset = MoleculeDataset("dataset/" + args.dataset, dataset=args.dataset)

    print(dataset)
    
    if args.split == "scaffold":
        smiles_list = pd.read_csv('dataset/' + args.dataset + '/processed/smiles.csv', header=None)[0].tolist()
        train_dataset, valid_dataset, test_dataset = scaffold_split(dataset, smiles_list, null_value=0, frac_train=0.8,frac_valid=0.1, frac_test=0.1)
        print("scaffold")
    elif args.split == "random":
        train_dataset, valid_dataset, test_dataset = random_split(dataset, null_value=0, frac_train=0.8,frac_valid=0.1, frac_test=0.1, seed = args.seed)
        print("random")
    elif args.split == "random_scaffold":
        smiles_list = pd.read_csv('dataset/' + args.dataset + '/processed/smiles.csv', header=None)[0].tolist()
        train_dataset, valid_dataset, test_dataset = random_scaffold_split(dataset, smiles_list, null_value=0, frac_train=0.8,frac_valid=0.1, frac_test=0.1, seed = args.seed)
        print("random scaffold")
    else:
        raise ValueError("Invalid split option.")

    print(train_dataset[0])

    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers = args.num_workers)
    val_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers = args.num_workers)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers = args.num_workers)

    #set up model
    model = GNN_graphpred(args.num_layer, args.emb_dim, num_tasks, JK = args.JK, drop_ratio = args.dropout_ratio, graph_pooling = args.graph_pooling, gnn_type = args.gnn_type)
    if not args.input_model_file == "":
        model.from_pretrained(args.input_model_file)
    
    model.to(device)

    #set up optimizer
    #different learning rate for different part of GNN
    model_param_group = []
    model_param_group.append({"params": model.gnn.parameters()})
    if args.graph_pooling == "attention":
        model_param_group.append({"params": model.pool.parameters(), "lr":args.lr*args.lr_scale})
    model_param_group.append({"params": model.graph_pred_linear.parameters(), "lr":args.lr*args.lr_scale})
    optimizer = optim.Adam(model_param_group, lr=args.lr, weight_decay=args.decay)
    print(optimizer)

    for epoch in range(1, args.epochs+1):
        print("====epoch " + str(epoch))
        
        train(args, model, device, train_loader, optimizer)

        print("====Evaluation")
        if args.eval_train:
            train_acc = eval(args, model, device, train_loader)
        else:
            print("omit the training accuracy computation")
            train_acc = 0
        val_acc = eval(args, model, device, val_loader)
        test_acc = eval(args, model, device, test_loader)

        print("train: %f val: %f test: %f" %(train_acc, val_acc, test_acc))

if __name__ == "__main__":
    main()

[((25, 12, 25, 52), 'torch.nn.BCEWithLogitsLoss', 'nn.BCEWithLogitsLoss', (), '', True, 'import torch.nn as nn\n'), ((83, 13, 83, 115), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((118, 4, 118, 35), 'torch.manual_seed', 'torch.manual_seed', ({(118, 22, 118, 34): 'args.runseed'}, {}), '(args.runseed)', False, 'import torch\n'), ((119, 4, 119, 32), 'numpy.random.seed', 'np.random.seed', ({(119, 19, 119, 31): 'args.runseed'}, {}), '(args.runseed)', True, 'import numpy as np\n'), ((121, 7, 121, 32), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((147, 14, 147, 78), 'loader.MoleculeDataset', 'MoleculeDataset', (), '', False, 'from loader import MoleculeDataset\n'), ((167, 19, 167, 118), 'torch_geometric.data.DataLoader', 'DataLoader', (), '', False, 'from torch_geometric.data import DataLoader\n'), ((168, 17, 168, 117), 'torch_geometric.data.DataLoader', 'DataLoader', (), '', False, 'from torch_geometric.data import DataLoader\n'), ((169, 18, 169, 117), 'torch_geometric.data.DataLoader', 'DataLoader', (), '', False, 'from torch_geometric.data import DataLoader\n'), ((172, 12, 172, 175), 'model.GNN_graphpred', 'GNN_graphpred', (), '', False, 'from model import GNN, GNN_graphpred\n'), ((185, 16, 185, 82), 'torch.optim.Adam', 'optim.Adam', (), '', True, 'import torch.optim as optim\n'), ((30, 33, 30, 63), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n'), ((54, 33, 54, 63), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n'), ((120, 57, 120, 82), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((120, 88, 120, 107), 'torch.device', 'torch.device', ({(120, 101, 120, 106): '"""cpu"""'}, {}), "('cpu')", False, 'import torch\n'), ((122, 8, 122, 48), 'torch.cuda.manual_seed_all', 'torch.cuda.manual_seed_all', ({(122, 35, 122, 47): 'args.runseed'}, {}), '(args.runseed)', False, 'import torch\n'), ((153, 53, 153, 149), 'splitters.scaffold_split', 'scaffold_split', (), '', False, 'from splitters import scaffold_split, random_split\n'), ((43, 15, 43, 34), 'torch.sum', 'torch.sum', ({(43, 25, 43, 33): 'loss_mat'}, {}), '(loss_mat)', False, 'import torch\n'), ((43, 35, 43, 54), 'torch.sum', 'torch.sum', ({(43, 45, 43, 53): 'is_valid'}, {}), '(is_valid)', False, 'import torch\n'), ((57, 13, 57, 28), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((156, 53, 156, 152), 'splitters.random_split', 'random_split', (), '', False, 'from splitters import scaffold_split, random_split\n'), ((69, 11, 69, 35), 'numpy.sum', 'np.sum', ({(69, 18, 69, 34): '(y_true[:, (i)] == 1)'}, {}), '(y_true[:, (i)] == 1)', True, 'import numpy as np\n'), ((69, 44, 69, 69), 'numpy.sum', 'np.sum', ({(69, 51, 69, 68): '(y_true[:, (i)] == -1)'}, {}), '(y_true[:, (i)] == -1)', True, 'import numpy as np\n'), ((71, 28, 71, 91), 'sklearn.metrics.roc_auc_score', 'roc_auc_score', ({(71, 42, 71, 68): '((y_true[is_valid, i] + 1) / 2)', (71, 70, 71, 90): 'y_scores[is_valid, i]'}, {}), '((y_true[is_valid, i] + 1) / 2, y_scores[is_valid, i])', False, 'from sklearn.metrics import roc_auc_score\n'), ((63, 13, 63, 39), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((64, 15, 64, 43), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((152, 22, 152, 99), 'pandas.read_csv', 'pd.read_csv', (), '', True, 'import pandas as pd\n'), ((40, 51, 40, 78), 'torch.zeros', 'torch.zeros', ({(40, 63, 40, 77): 'loss_mat.shape'}, {}), '(loss_mat.shape)', False, 'import torch\n'), ((159, 22, 159, 99), 'pandas.read_csv', 'pd.read_csv', (), '', True, 'import pandas as pd\n')]