from typing import Optional, List, NamedTuple

import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn import ModuleList, BatchNorm1d
from torch_sparse import SparseTensor
from pytorch_lightning.metrics import Accuracy
from pytorch_lightning.callbacks import ModelCheckpoint
from pytorch_lightning import LightningDataModule, LightningModule, Trainer

import torch_geometric.transforms as T
from torch_geometric.nn import SAGEConv
from torch_geometric import seed_everything
from torch_geometric.loader import NeighborSampler
from torch_geometric.datasets import Reddit


class Batch(NamedTuple):
    x: Tensor
    y: Tensor
    adjs_t: List[SparseTensor]


class RedditDataModule(LightningDataModule):
    def __init__(self, data_dir):
        super().__init__()
        self.data_dir = data_dir
        self.transform = T.ToSparseTensor()

    @property
    def num_features(self) -> int:
        return 602

    @property
    def num_classes(self) -> int:
        return 41

    def prepare_data(self):
        Reddit(self.data_dir, pre_transform=self.transform)

    def setup(self, stage: Optional[str] = None):
        self.data = Reddit(self.data_dir, pre_transform=self.transform)[0]

    def train_dataloader(self):
        return NeighborSampler(self.data.adj_t, node_idx=self.data.train_mask,
                               sizes=[25, 10], return_e_id=False,
                               transform=self.convert_batch, batch_size=1024,
                               shuffle=True, num_workers=6,
                               persistent_workers=True)

    def val_dataloader(self):
        return NeighborSampler(self.data.adj_t, node_idx=self.data.val_mask,
                               sizes=[25, 10], return_e_id=False,
                               transform=self.convert_batch, batch_size=2048,
                               num_workers=3, persistent_workers=True)

    def test_dataloader(self):
        return NeighborSampler(self.data.adj_t, node_idx=self.data.test_mask,
                               sizes=[25, 10], return_e_id=False,
                               transform=self.convert_batch, batch_size=2048,
                               num_workers=3, persistent_workers=True)

    def convert_batch(self, batch_size, n_id, adjs):
        return Batch(
            x=self.data.x[n_id],
            y=self.data.y[n_id[:batch_size]],
            adjs_t=[adj_t for adj_t, _, _ in adjs],
        )


class GraphSAGE(LightningModule):
    def __init__(self, in_channels: int, out_channels: int,
                 hidden_channels: int = 256, num_layers: int = 2,
                 dropout: float = 0.5):
        super().__init__()
        self.save_hyperparameters()
        self.dropout = dropout

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

        self.bns = ModuleList()
        for _ in range(num_layers - 1):
            self.bns.append(BatchNorm1d(hidden_channels))

        self.train_acc = Accuracy()
        self.val_acc = Accuracy()
        self.test_acc = Accuracy()

    def forward(self, x: Tensor, adjs_t: List[SparseTensor]) -> Tensor:
        for i, adj_t in enumerate(adjs_t):
            x = self.convs[i]((x, x[:adj_t.size(0)]), adj_t)
            if i < len(adjs_t) - 1:
                x = self.bns[i](x)
                x = x.relu_()
                x = F.dropout(x, p=self.dropout, training=self.training)
        return x

    def training_step(self, batch: Batch, batch_idx: int):
        x, y, adjs_t = batch
        y_hat = self(x, adjs_t)
        train_loss = F.cross_entropy(y_hat, y)
        self.train_acc(y_hat.softmax(dim=-1), y)
        self.log('train_acc', self.train_acc, prog_bar=True, on_step=False,
                 on_epoch=True)
        return train_loss

    def validation_step(self, batch: Batch, batch_idx: int):
        x, y, adjs_t = batch
        y_hat = self(x, adjs_t)
        self.val_acc(y_hat.softmax(dim=-1), y)
        self.log('val_acc', self.val_acc, prog_bar=True, on_step=False,
                 on_epoch=True)

    def test_step(self, batch: Batch, batch_idx: int):
        x, y, adjs_t = batch
        y_hat = self(x, adjs_t)
        self.test_acc(y_hat.softmax(dim=-1), y)
        self.log('test_acc', self.test_acc, prog_bar=True, on_step=False,
                 on_epoch=True)

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.01)


def main():
    seed_everything(42)
    datamodule = RedditDataModule('data/Reddit')
    model = GraphSAGE(datamodule.num_features, datamodule.num_classes)
    checkpoint_callback = ModelCheckpoint(monitor='val_acc', save_top_k=1)
    trainer = Trainer(gpus=1, max_epochs=10, callbacks=[checkpoint_callback])

    # Uncomment to train on multiple GPUs:
    # trainer = Trainer(gpus=2, accelerator='ddp', max_epochs=10,
    #                   callbacks=[checkpoint_callback])

    trainer.fit(model, datamodule=datamodule)
    trainer.test()


if __name__ == "__main__":
    main()
