import argparse

import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

import torch_geometric.transforms as T
from torch_geometric.nn import GCNConv, SAGEConv, GATv2Conv #as GATv2Layer

from ogb.nodeproppred import PygNodePropPredDataset, Evaluator

from logger import Logger

from attention import GATv2Layer, GATv1Layer


class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout):
        super(GCN, self).__init__()

        self.convs = torch.nn.ModuleList()
        self.convs.append(GCNConv(in_channels, hidden_channels, cached=True))
        self.bns = torch.nn.ModuleList()
        self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(
                GCNConv(hidden_channels, hidden_channels, cached=True))
            self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        self.convs.append(GCNConv(hidden_channels, out_channels, cached=True))

        self.dropout = dropout

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        for bn in self.bns:
            bn.reset_parameters()

    def forward(self, x, adj_t):
        for i, conv in enumerate(self.convs[:-1]):
            x = conv(x, adj_t)
            x = self.bns[i](x)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.convs[-1](x, adj_t)
        return x.log_softmax(dim=-1)


class SAGE(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout):
        super(SAGE, self).__init__()

        self.convs = torch.nn.ModuleList()
        self.convs.append(SAGEConv(in_channels, hidden_channels))
        self.bns = torch.nn.ModuleList()
        self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(SAGEConv(hidden_channels, hidden_channels))
            self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        self.convs.append(SAGEConv(hidden_channels, out_channels))

        self.dropout = dropout

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        for bn in self.bns:
            bn.reset_parameters()

    def forward(self, x, adj_t):
        for i, conv in enumerate(self.convs[:-1]):
            x = conv(x, adj_t)
            x = self.bns[i](x)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.convs[-1](x, adj_t)
        return x.log_softmax(dim=-1)


class GAT(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers,
                 dropout, args):
        super(GAT, self).__init__()

        gat_layer = {'gatv1': GATv1Layer, 'gatv2': GATv2Layer}[args.model]
        Layer = lambda i, o: gat_layer(i, o, convolve=args.convolve, lambda_policy=args.lambda_policy, heads=args.heads, gcn_mode=args.gcn_mode, share_weights_score=True, share_weights_value=True)

        self.convs = torch.nn.ModuleList()
        self.convs.append(Layer(in_channels, hidden_channels // args.heads))
        self.bns = torch.nn.ModuleList()
        self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        for _ in range(num_layers - 2):
            self.convs.append(Layer(hidden_channels, hidden_channels // args.heads))
            self.bns.append(torch.nn.BatchNorm1d(hidden_channels))
        self.convs.append(Layer(hidden_channels, out_channels // args.heads))

        self.dropout = dropout
        self.dummy = torch.nn.Parameter(torch.zeros([]), requires_grad=True)

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        for bn in self.bns:
            bn.reset_parameters()

    def forward(self, x, adj_t):
        x = x + 0.*self.dummy
        for i, conv in enumerate(self.convs[:-1]):
            #x = conv(x, adj_t)
            x = checkpoint(conv, x, adj_t, preserve_rng_state=False)
            x = self.bns[i](x)
            #x = checkpoint(self.bns[i], x)
            x = F.relu(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        #x = self.convs[-1](x, adj_t)
        x = checkpoint(self.convs[-1], x, adj_t, preserve_rng_state=False)
        return x.log_softmax(dim=-1)

    @property
    def lmbda(self):
        result = []
        for conv in self.convs:
            result.append(float(conv.lmbda.mean()))
        return result

    @property
    def lmbda2(self):
        result = []
        for conv in self.convs:
            result.append(float(conv.lmbda2.mean()))
        return result


def train(model, data, train_idx, optimizer):
    model.train()

    optimizer.zero_grad()
    #out = checkpoint(model, data.x, data.adj_t)[train_idx]
    out = model(data.x, data.adj_t)[train_idx]
    loss = F.nll_loss(out, data.y.squeeze(1)[train_idx])
    loss.backward()
    optimizer.step()

    return loss.item()


@torch.no_grad()
def test(model, data, split_idx, evaluator):
    model.eval()

    out = model(data.x, data.adj_t)
    y_pred = out.argmax(dim=-1, keepdim=True)

    train_acc = evaluator.eval({
        'y_true': data.y[split_idx['train']],
        'y_pred': y_pred[split_idx['train']],
    })['acc']
    valid_acc = evaluator.eval({
        'y_true': data.y[split_idx['valid']],
        'y_pred': y_pred[split_idx['valid']],
    })['acc']
    test_acc = evaluator.eval({
        'y_true': data.y[split_idx['test']],
        'y_pred': y_pred[split_idx['test']],
    })['acc']

    return train_acc, valid_acc, test_acc


def main():
    parser = argparse.ArgumentParser(description='OGBN-Arxiv (GNN)')
    parser.add_argument('--device', type=int, default=0)
    parser.add_argument('--log_steps', type=int, default=1)
    parser.add_argument('--model', type=str, default='gcn')
    parser.add_argument('--num_layers', type=int, default=3)
    parser.add_argument('--hidden_channels', type=int, default=128)
    parser.add_argument('--dropout', type=float, default=0.5)
    parser.add_argument('--lr', type=float, default=0.01)
    parser.add_argument('--epochs', type=int, default=500)
    parser.add_argument('--runs', type=int, default=1)
    parser.add_argument('--lambda-policy', type=str, choices=['learn1', 'learn2', 'learn12', 'gcn_gat'])
    parser.add_argument('--convolve', action='store_true')
    parser.add_argument('--heads', type=int, default=1)
    parser.add_argument('--gcn-mode', action='store_true')
    parser.add_argument('--noise-level', type=float, default=0.0)
    args = parser.parse_args()
    print(args)

    device = f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu'
    #device = torch.device('cpu')
    device = torch.device(device)

    dataset = PygNodePropPredDataset(name='ogbn-arxiv',
                                     transform=T.ToSparseTensor())

    data = dataset[0]
    data.adj_t = data.adj_t.to_symmetric()

    data.x += torch.randn_like(data.x)*args.noise_level
    data = data.to(device)

    split_idx = dataset.get_idx_split()
    train_idx = split_idx['train'].to(device)

    if args.model == 'sage':
        model = SAGE(data.num_features, args.hidden_channels,
                     dataset.num_classes, args.num_layers,
                     args.dropout).to(device)
    elif 'gat' in args.model:
        model = GAT(data.num_features, args.hidden_channels,
                     dataset.num_classes, args.num_layers,
                     args.dropout, args).to(device)
    else:
        model = GCN(data.num_features, args.hidden_channels,
                    dataset.num_classes, args.num_layers,
                    args.dropout).to(device)

    def count_parameters(model):
        return sum(p.numel() for p in model.parameters() if p.requires_grad)

    print('No. parameters', count_parameters(model))

    evaluator = Evaluator(name='ogbn-arxiv')
    logger = Logger(args.runs, args)

    for run in range(args.runs):
        model.reset_parameters()
        optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
        for epoch in range(1, 1 + args.epochs):
            loss = train(model, data, train_idx, optimizer)
            result = test(model, data, split_idx, evaluator)
            logger.add_result(run, result)

            if epoch % args.log_steps == 0:
                train_acc, valid_acc, test_acc = result
                print(f'Run: {run + 1:02d}, '
                      f'Epoch: {epoch:02d}, '
                      f'Loss: {loss:.4f}, '
                      f'Train: {100 * train_acc:.2f}%, '
                      f'Valid: {100 * valid_acc:.2f}% '
                      f'Test: {100 * test_acc:.2f}%')

                if 'gat' in args.model:
                    print('Lambda 1', model.lmbda)
                    print('Lambda 2', model.lmbda2)

        logger.print_statistics(run)
    logger.print_statistics()


if __name__ == "__main__":
    torch.cuda.empty_cache()
    main()
