import dgl
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn import GraphConv
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--data', type=str, default = "cora", help = "benchmark dataset")
parser.add_argument('--h_feats', type=int, default = 16, help = "number of hidden channels")
parser.add_argument('--lr', type=float, default = 0.01, help = "learning rate")
parser.add_argument('--epochs', type=int, default = 100, help = "epochs")
args = parser.parse_args()

data = str(args.data)
h_feats = int(args.h_feats)
lr = float(args.lr)
epochs = int(args.epochs)

if data == "cora":
    dataset = dgl.data.CoraGraphDataset()
elif data == "citeseer":
    dataset = dgl.data.CiteseerGraphDataset()
elif data == "pubmed":
    dataset = dgl.data.PubmedGraphDataset()

print('Number of categories:', dataset.num_classes)

g = dataset[0]


class GCN(nn.Module):
    def __init__(self, in_feats, h_feats, num_classes):
        super(GCN, self).__init__()
        self.conv1 = GraphConv(in_feats, h_feats)
        self.conv2 = GraphConv(h_feats, num_classes)

    def forward(self, g, in_feat):
        h = self.conv1(g, in_feat)
        h = F.relu(h)
        h = self.conv2(g, h)
        return h


def train(g, model):
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    best_val_acc = 0
    best_test_acc = 0

    features = g.ndata['feat']
    labels = g.ndata['label']
    train_mask = g.ndata['train_mask']
    val_mask = g.ndata['val_mask']
    test_mask = g.ndata['test_mask']
    for e in range(epochs):
        # Forward
        logits = model(g, features)

        # Compute prediction
        pred = logits.argmax(1)

        # Compute loss
        # Note that you should only compute the losses of the nodes in the training set.
        loss = F.cross_entropy(logits[train_mask], labels[train_mask])

        # Compute accuracy on training/validation/test
        train_acc = (pred[train_mask] == labels[train_mask]).float().mean()
        val_acc = (pred[val_mask] == labels[val_mask]).float().mean()
        test_acc = (pred[test_mask] == labels[test_mask]).float().mean()

        # Save the best validation accuracy and the corresponding test accuracy.
        if best_val_acc < val_acc:
            best_val_acc = val_acc
            best_test_acc = test_acc

        # Backward
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if e % (epochs/20) == 0:
            print('In epoch {}, loss: {:.3f}, val acc: {:.3f} (best {:.3f}), test acc: {:.3f} (best {:.3f})'.format(
                e, loss, val_acc, best_val_acc, test_acc, best_test_acc))

model = GCN(g.ndata['feat'].shape[1], h_feats, dataset.num_classes)
train(g, model)
