from typing import Optional

import torch
from torch import Tensor
import torch.nn.functional as F
from torch.nn import ModuleList, Sequential, Linear, BatchNorm1d, ReLU, Dropout
from torch_sparse import SparseTensor
from torch_scatter import segment_csr
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 GINConv
from torch_geometric.data import Batch
from torch_geometric import seed_everything
from torch_geometric.loader import DataLoader
from torch_geometric.datasets import TUDataset


class IMDBBinary(LightningDataModule):
    def __init__(self, data_dir):
        super().__init__()
        self.data_dir = data_dir
        self.transform = T.Compose([
            T.OneHotDegree(self.num_features - 1),
            T.ToSparseTensor(),
        ])

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

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

    def prepare_data(self):
        TUDataset(self.data_dir, name='IMDB-BINARY',
                  pre_transform=self.transform)

    def setup(self, stage: Optional[str] = None):
        dataset = TUDataset(self.data_dir, name='IMDB-BINARY',
                            pre_transform=self.transform).shuffle()

        self.test_dataset = dataset[:len(dataset) // 10]
        self.val_dataset = dataset[len(dataset) // 10:len(dataset) // 5]
        self.train_dataset = dataset[len(dataset) // 5:]

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=64, shuffle=True)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=128)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=128)


class GIN(LightningModule):
    def __init__(self, in_channels: int, out_channels: int,
                 hidden_channels: int = 64, num_layers: int = 3,
                 dropout: float = 0.5):
        super().__init__()
        self.save_hyperparameters()

        self.convs = ModuleList()
        for _ in range(num_layers):
            mlp = Sequential(
                Linear(in_channels, 2 * hidden_channels),
                BatchNorm1d(2 * hidden_channels),
                ReLU(inplace=True),
                Linear(2 * hidden_channels, hidden_channels),
                BatchNorm1d(hidden_channels),
                ReLU(inplace=True),
            )
            conv = GINConv(mlp, train_eps=True)
            self.convs.append(conv)
            in_channels = hidden_channels

        self.classifier = Sequential(
            Linear(hidden_channels, hidden_channels),
            BatchNorm1d(hidden_channels),
            ReLU(inplace=True),
            Dropout(p=dropout),
            Linear(hidden_channels, out_channels),
        )

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

    def forward(self, x: Tensor, adj_t: SparseTensor, ptr: Tensor) -> Tensor:
        for conv in self.convs:
            x = conv(x, adj_t)
        x = segment_csr(x, ptr, reduce='sum')  # Global add pooling.
        return self.classifier(x)

    def training_step(self, batch: Batch, batch_idx: int):
        y_hat = self(batch.x, batch.adj_t, batch.ptr)
        train_loss = F.cross_entropy(y_hat, batch.y)
        self.train_acc(y_hat.softmax(dim=-1), batch.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):
        y_hat = self(batch.x, batch.adj_t, batch.ptr)
        self.val_acc(y_hat.softmax(dim=-1), batch.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):
        y_hat = self(batch.x, batch.adj_t, batch.ptr)
        self.test_acc(y_hat.softmax(dim=-1), batch.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 = IMDBBinary('data/TUDataset')
    model = GIN(datamodule.num_features, datamodule.num_classes)
    checkpoint_callback = ModelCheckpoint(monitor='val_acc', save_top_k=1)
    trainer = Trainer(gpus=1, max_epochs=20, callbacks=[checkpoint_callback])

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

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


if __name__ == "__main__":
    main()
