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()

[((464, 502), 'torch.nn.BCEWithLogitsLoss', 'nn.BCEWithLogitsLoss', ([], {'reduction': '"""none"""'}), "(reduction='none')\n", (484, 502), True, 'import torch.nn as nn\n'), ((2369, 2476), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""PyTorch implementation of pre-training of graph neural networks"""'}), "(description=\n    'PyTorch implementation of pre-training of graph neural networks')\n", (2392, 2476), False, 'import argparse\n'), ((5027, 5058), 'torch.manual_seed', 'torch.manual_seed', (['args.runseed'], {}), '(args.runseed)\n', (5044, 5058), False, 'import torch\n'), ((5063, 5091), 'numpy.random.seed', 'np.random.seed', (['args.runseed'], {}), '(args.runseed)\n', (5077, 5091), True, 'import numpy as np\n'), ((5207, 5232), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (5230, 5232), False, 'import torch\n'), ((5920, 5984), 'loader.MoleculeDataset', 'MoleculeDataset', (["('dataset/' + args.dataset)"], {'dataset': 'args.dataset'}), "('dataset/' + args.dataset, dataset=args.dataset)\n", (5935, 5984), False, 'from loader import MoleculeDataset\n'), ((7012, 7113), 'torch_geometric.data.DataLoader', 'DataLoader', (['train_dataset'], {'batch_size': 'args.batch_size', 'shuffle': '(True)', 'num_workers': 'args.num_workers'}), '(train_dataset, batch_size=args.batch_size, shuffle=True,\n    num_workers=args.num_workers)\n', (7022, 7113), False, 'from torch_geometric.data import DataLoader\n'), ((7129, 7231), 'torch_geometric.data.DataLoader', 'DataLoader', (['valid_dataset'], {'batch_size': 'args.batch_size', 'shuffle': '(False)', 'num_workers': 'args.num_workers'}), '(valid_dataset, batch_size=args.batch_size, shuffle=False,\n    num_workers=args.num_workers)\n', (7139, 7231), False, 'from torch_geometric.data import DataLoader\n'), ((7248, 7349), 'torch_geometric.data.DataLoader', 'DataLoader', (['test_dataset'], {'batch_size': 'args.batch_size', 'shuffle': '(False)', 'num_workers': 'args.num_workers'}), '(test_dataset, batch_size=args.batch_size, shuffle=False,\n    num_workers=args.num_workers)\n', (7258, 7349), False, 'from torch_geometric.data import DataLoader\n'), ((7379, 7542), 'model.GNN_graphpred', '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'}), '(args.num_layer, args.emb_dim, num_tasks, JK=args.JK,\n    drop_ratio=args.dropout_ratio, graph_pooling=args.graph_pooling,\n    gnn_type=args.gnn_type)\n', (7392, 7542), False, 'from model import GNN, GNN_graphpred\n'), ((8095, 8161), 'torch.optim.Adam', 'optim.Adam', (['model_param_group'], {'lr': 'args.lr', 'weight_decay': 'args.decay'}), '(model_param_group, lr=args.lr, weight_decay=args.decay)\n', (8105, 8161), True, 'import torch.optim as optim\n'), ((609, 639), 'tqdm.tqdm', 'tqdm', (['loader'], {'desc': '"""Iteration"""'}), "(loader, desc='Iteration')\n", (613, 639), False, 'from tqdm import tqdm\n'), ((1391, 1421), 'tqdm.tqdm', 'tqdm', (['loader'], {'desc': '"""Iteration"""'}), "(loader, desc='Iteration')\n", (1395, 1421), False, 'from tqdm import tqdm\n'), ((5149, 5174), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (5172, 5174), False, 'import torch\n'), ((5180, 5199), 'torch.device', 'torch.device', (['"""cpu"""'], {}), "('cpu')\n", (5192, 5199), False, 'import torch\n'), ((5242, 5282), 'torch.cuda.manual_seed_all', 'torch.cuda.manual_seed_all', (['args.runseed'], {}), '(args.runseed)\n', (5268, 5282), False, 'import torch\n'), ((6208, 6309), 'splitters.scaffold_split', 'scaffold_split', (['dataset', 'smiles_list'], {'null_value': '(0)', 'frac_train': '(0.8)', 'frac_valid': '(0.1)', 'frac_test': '(0.1)'}), '(dataset, smiles_list, null_value=0, frac_train=0.8,\n    frac_valid=0.1, frac_test=0.1)\n', (6222, 6309), False, 'from splitters import scaffold_split, random_split\n'), ((1175, 1194), 'torch.sum', 'torch.sum', (['loss_mat'], {}), '(loss_mat)\n', (1184, 1194), False, 'import torch\n'), ((1195, 1214), 'torch.sum', 'torch.sum', (['is_valid'], {}), '(is_valid)\n', (1204, 1214), False, 'import torch\n'), ((1471, 1486), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (1484, 1486), False, 'import torch\n'), ((6417, 6519), 'splitters.random_split', 'random_split', (['dataset'], {'null_value': '(0)', 'frac_train': '(0.8)', 'frac_valid': '(0.1)', 'frac_test': '(0.1)', 'seed': 'args.seed'}), '(dataset, null_value=0, frac_train=0.8, frac_valid=0.1,\n    frac_test=0.1, seed=args.seed)\n', (6429, 6519), False, 'from splitters import scaffold_split, random_split\n'), ((1900, 1927), 'numpy.sum', 'np.sum', (['(y_true[:, (i)] == 1)'], {}), '(y_true[:, (i)] == 1)\n', (1906, 1927), True, 'import numpy as np\n'), ((1933, 1961), 'numpy.sum', 'np.sum', (['(y_true[:, (i)] == -1)'], {}), '(y_true[:, (i)] == -1)\n', (1939, 1961), True, 'import numpy as np\n'), ((2034, 2101), 'sklearn.metrics.roc_auc_score', 'roc_auc_score', (['((y_true[is_valid, i] + 1) / 2)', 'y_scores[is_valid, i]'], {}), '((y_true[is_valid, i] + 1) / 2, y_scores[is_valid, i])\n', (2047, 2101), False, 'from sklearn.metrics import roc_auc_score\n'), ((1663, 1687), 'torch.cat', 'torch.cat', (['y_true'], {'dim': '(0)'}), '(y_true, dim=0)\n', (1672, 1687), False, 'import torch\n'), ((1719, 1745), 'torch.cat', 'torch.cat', (['y_scores'], {'dim': '(0)'}), '(y_scores, dim=0)\n', (1728, 1745), False, 'import torch\n'), ((6065, 6142), 'pandas.read_csv', 'pd.read_csv', (["('dataset/' + args.dataset + '/processed/smiles.csv')"], {'header': 'None'}), "('dataset/' + args.dataset + '/processed/smiles.csv', header=None)\n", (6076, 6142), True, 'import pandas as pd\n'), ((1049, 1076), 'torch.zeros', 'torch.zeros', (['loss_mat.shape'], {}), '(loss_mat.shape)\n', (1060, 1076), False, 'import torch\n'), ((6605, 6682), 'pandas.read_csv', 'pd.read_csv', (["('dataset/' + args.dataset + '/processed/smiles.csv')"], {'header': 'None'}), "('dataset/' + args.dataset + '/processed/smiles.csv', header=None)\n", (6616, 6682), True, 'import pandas as pd\n')]