# Reaches around 0.7945 ± 0.0059 test accuracy.

import argparse
import random

import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
from ogb.nodeproppred import PygNodePropPredDataset, Evaluator
from torch.nn import Linear as Lin
from torch_geometric.loader import NeighborSampler
from tqdm import tqdm

from models import GAT_TYPE
from utils.logger import Logger

parser = argparse.ArgumentParser(description='OGBN - products (GAT)')
parser.add_argument("--type", dest="type", default=GAT_TYPE.GAT, type=GAT_TYPE.from_string, choices=list(GAT_TYPE))
parser.add_argument('--num_heads', type=int, default=4)
parser.add_argument('--log_steps', type=int, default=1)
parser.add_argument('--lr', type=float, default=0.001)
parser.add_argument('--eval_steps', type=int, default=10)
parser.add_argument('--dataset', type=str, default='ogbn-products')
parser.add_argument('--hidden_channels', type=int, default=128)
parser.add_argument('--num_layers', type=int, default=3)
parser.add_argument('--epochs', type=int, default=100)
parser.add_argument('--runs', type=int, default=5)
parser.add_argument('--patient', type=int, default=10)
parser.add_argument('--max_loss', type=float, default=2.0)
parser.add_argument('--seed', type=int, default=123)
parser.add_argument('--dropout', type=float, default=0.5)

parser.add_argument('--gcn-mode', action='store_true', help='gcn mode')
parser.add_argument('--convolve', action='store_true', help='convolve the attention score')
parser.add_argument('--lambda-policy', type=str, choices=['learn1', 'learn2', 'learn12', 'gcn_gat'])
parser.add_argument('--share-weights-score', action='store_true', help='share W_i and W_j')
parser.add_argument('--share-weights-value', action='store_true', help='share W_j and W_v')

args = parser.parse_args()
print(args)
torch.manual_seed(args.seed)
random.seed(args.seed)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = PygNodePropPredDataset('ogbn-products')
split_idx = dataset.get_idx_split()
evaluator = Evaluator(name='ogbn-products')
data = dataset[0]

train_idx = split_idx['train']
train_loader = NeighborSampler(data.edge_index, node_idx=train_idx,
                               sizes=[10] * args.num_layers, batch_size=256,
                               shuffle=True, num_workers=6)
subgraph_loader = NeighborSampler(data.edge_index, node_idx=None, sizes=[-1],
                                  batch_size=4096, shuffle=False,
                                  num_workers=6)


class GAT(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels, num_layers, heads, dropout, convolve, lambda_policy,
                 share_weights_value, share_weights_score, gcn_mode):
        super(GAT, self).__init__()

        self.num_layers = num_layers
        self.convs = torch.nn.ModuleList()
        self.skips = torch.nn.ModuleList()

        base_layer = args.type.get_base_layer()
        kwargs = {'bias': True}

        #import pdb; pdb.set_trace()
        self.is_gat = not gcn_mode  # base_layer in (GAT_TYPE.GAT2, GAT_TYPE.GAT)
        if self.is_gat:
            kwargs.update({
                'convolve': convolve, 'lambda_policy': lambda_policy,
                'share_weights_value': share_weights_value, 'share_weights_score': share_weights_score,
                'gcn_mode': False, 'heads': heads,
            })
        else:
            kwargs.update({
                'share_weights_value': True, 'share_weights_score': True,
                'gcn_mode': True, 'convolve': False, 'heads': 1
            })
            heads = 1

        self.convs.append(base_layer(in_channels, hidden_channels // heads, **kwargs))
        for _ in range(num_layers - 2):
            self.convs.append(
                base_layer(hidden_channels, hidden_channels // heads, **kwargs))

        kwargs['heads'] = 1
        self.convs.append(
            base_layer(hidden_channels, out_channels, **kwargs))

        self.skips.append(Lin(in_channels, hidden_channels))
        for _ in range(num_layers - 2):
            self.skips.append(
                Lin(hidden_channels, hidden_channels))
        self.skips.append(Lin(hidden_channels, out_channels))

    def reset_parameters(self):
        for conv in self.convs:
            conv.reset_parameters()
        for skip in self.skips:
            skip.reset_parameters()

    @property
    def lmbda(self):
        return [float(l.lmbda) for l in self.convs] if self.is_gat else [0.] * len(self.convs)

    @property
    def lmbda2(self):
        return [float(l.lmbda2) for l in self.convs] if self.is_gat else [1.] * len(self.convs)

    def forward(self, x, adjs):
        # `train_loader` computes the k-hop neighborhood of a batch of nodes,
        # and returns, for each layer, a bipartite graph object, holding the
        # bipartite edges `edge_index`, the index `e_id` of the original edges,
        # and the size/shape `size` of the bipartite graph.
        # Target nodes are also included in the source nodes so that one can
        # easily apply skip-connections or add self-loops.
        for i, (edge_index, _, size) in enumerate(adjs):
            x_target = x[:size[1]]  # Target nodes are always placed first.
            #x = self.convs[i]((x, x_target), edge_index)
            x = checkpoint(self.convs[i], x, edge_index, size[1], preserve_rng_state=False)
            x = x + self.skips[i](x_target)
            if i != self.num_layers - 1:
                x = F.elu(x)
                x = F.dropout(x, p=args.dropout, training=self.training)
        return x.log_softmax(dim=-1)

    def inference(self, x_all):
        pbar = tqdm(total=x_all.size(0) * self.num_layers)
        pbar.set_description('Evaluating')

        # Compute representations of nodes layer by layer, using *all*
        # available edges. This leads to faster computation in contrast to
        # immediately computing the final representations of each batch.
        total_edges = 0
        for i in range(self.num_layers):
            xs = []
            for batch_size, n_id, adj in subgraph_loader:
                edge_index, _, size = adj.to(device)
                total_edges += edge_index.size(1)
                x = x_all[n_id].to(device)
                x_target = x[:size[1]]
                x = self.convs[i](x, edge_index, size[1])
                x = x + self.skips[i](x_target)
                # x = x + x_target

                if i != self.num_layers - 1:
                    x = F.elu(x)
                xs.append(x.cpu())

                pbar.update(batch_size)

            x_all = torch.cat(xs, dim=0)

        pbar.close()

        return x_all


model = GAT(dataset.num_features, args.hidden_channels, dataset.num_classes, num_layers=args.num_layers,
            heads=args.num_heads, dropout=args.dropout, convolve=args.convolve, lambda_policy=args.lambda_policy,
            share_weights_value=args.share_weights_value, share_weights_score=args.share_weights_score,
            gcn_mode=args.gcn_mode)
model = model.to(device)
print(model)

x = data.x.to(device)
y = data.y.squeeze().to(device)


def train(epoch):
    model.train()

    pbar = tqdm(total=train_idx.size(0), leave=False, dynamic_ncols=True)
    pbar.set_description(f'Epoch {epoch:02d}')

    total_loss = total_correct = 0
    for batch_size, n_id, adjs in train_loader:
        # `adjs` holds a list of `(edge_index, e_id, size)` tuples.
        adjs = [adj.to(device) for adj in adjs]

        optimizer.zero_grad()
        out = model(x[n_id], adjs)
        loss = F.nll_loss(out, y[n_id[:batch_size]])
        loss.backward()
        optimizer.step()

        total_loss += float(loss)
        total_correct += int(out.argmax(dim=-1).eq(y[n_id[:batch_size]]).sum())
        pbar.update(batch_size)

    pbar.close()

    loss = total_loss / len(train_loader)
    approx_acc = total_correct / train_idx.size(0)

    return loss, approx_acc


@torch.no_grad()
def test():
    model.eval()

    out = model.inference(x)

    y_true = y.cpu().unsqueeze(-1)
    y_pred = out.argmax(dim=-1, keepdim=True)

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

    return [train_acc, val_acc, test_acc], [model.lmbda, model.lmbda2]


logger = Logger(args.runs, args)
run = 0
while run < args.runs:
    model.reset_parameters()
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    run_success = True
    low_loss = False
    max_val_score = -float('inf')
    patient = 0
    best_val_acc = best_train_acc = final_test_acc = 0
    logger.set_time(run)
    for epoch in range(1, 1 + args.epochs):
        loss, acc = train(epoch)
        # print(f'Run: {run:02d}, Epoch: {epoch:02d}, Loss: {loss:.4f}, Approx. Train: {acc:.4f}')
        if loss <= args.max_loss:
            low_loss = True
        if epoch > args.epochs // 2 and loss > args.max_loss and low_loss is False:
            run_success = False
            logger.reset(run)
            print('Learning failed. Rerun...')
            break
        if epoch > args.epochs // 2 and epoch % args.eval_steps == 0:
            # if epoch > 50 and epoch % 10 == 0:
            result, lmbdas = test()
            train_acc, valid_acc, test_acc = result
            logger.add_result(run, result, lmbdas)
            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 valid_acc >= max_val_score:
                max_val_score = valid_acc
                patient = 0
            else:
                patient += 1
                if patient >= args.patient:
                    print("patient exhausted")
                    if low_loss is False:
                        run_success = False
                        logger.reset(run)
                        print('Learning failed. Rerun...')
                    break
        elif epoch % args.log_steps == 0:
            print(f'Run: {run + 1:02d}, '
                  f'Epoch: {epoch:02d}, '
                  f'Loss: {loss:.4f},'
                  f'Approx. Train: {acc:.4f}')
    if run_success:
        logger.print_statistics(run)
        run += 1
logger.print_statistics()
