import torch.nn as nn
import torch.nn.functional as F
import torch
from torch_geometric.nn import GINConv
from .base_model import BaseModel
from torch_sparse import SparseTensor


class GIN(BaseModel):

    def __init__(self, nfeat, nhid, nclass, nlayers=2, dropout=0., lr=0.01,
                 with_bn=False, weight_decay=5e-4, device=None, args=None):

        super(GIN, self).__init__()
        from utils import eval_acc, eval_f1, eval_rocauc
        if args.dataset == 'twitch-e':
            self.eval_func = eval_rocauc
        elif args.dataset == 'elliptic':
            self.eval_func = eval_f1
        elif args.dataset in ['cora', 'amazon-photo', 'ogb-arxiv', 'fb100']:
            self.eval_func = eval_acc
        else:
            raise NotImplementedError

        assert device is not None, "Please specify 'device'!"
        self.device = device

        self.convs = nn.ModuleList()
        self.bns = nn.ModuleList()

        self.convs.append(GINConv(nn.Sequential(
            nn.Linear(nfeat, nhid),
            nn.ReLU(),
            nn.Linear(nhid, nhid),
            nn.ReLU(),
            nn.BatchNorm1d(nhid) if with_bn else nn.Identity()
        ), train_eps=True))

        for _ in range(nlayers - 2):
            self.convs.append(GINConv(nn.Sequential(
                nn.Linear(nhid, nhid),
                nn.ReLU(),
                nn.Linear(nhid, nhid),
                nn.ReLU(),
                nn.BatchNorm1d(nhid) if with_bn else nn.Identity()
            ), train_eps=True))

        self.convs.append(GINConv(nn.Sequential(
            nn.Linear(nhid, nhid),
            nn.ReLU(),
            nn.Linear(nhid, nclass),
        ), train_eps=True))

        self.dropout = dropout
        self.lr = lr
        self.weight_decay = weight_decay
        self.output = None
        self.best_model = None
        self.best_output = None
        self.activation = F.relu
        self.name = 'GIN'
        self.args = args
        self.with_bn = with_bn

    def forward(self, x, edge_index, edge_weight=None):
        for i, conv in enumerate(self.convs[:-1]):
            x = conv(x, edge_index)
            x = self.activation(x)
            x = F.dropout(x, p=self.dropout, training=self.training)
        x = self.convs[-1](x, edge_index)
        return x

    def initialize(self):
        self.reset_parameters()

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

    def get_embed(self, x, edge_index, edge_weight=None):
        x, edge_index, edge_weight = self._ensure_contiguousness(x, edge_index, edge_weight)
        for ii, layer in enumerate(self.convs):
            if ii == len(self.convs) - 1:
                return x
            x = layer(x, edge_index)
            if ii != len(self.convs) - 1:
                x = F.relu(x)
        return x


if __name__ == "__main__":
    from deeprobust.graph.data import Dataset, Dpr2Pyg
    data = Dataset(root='/tmp/', name='cora')
    adj, features, labels = data.adj, data.features, data.labels
    idx_train, idx_val, idx_test = data.idx_train, data.idx_val, data.idx_test
    gin = GIN(nfeat=features.shape[1],
              nhid=8,
              nclass=labels.max().item() + 1,
              dropout=0.5, device='cpu', args=args)
    gin = gin.to('cpu')
    pyg_data = Dpr2Pyg(data)
    gin.fit(pyg_data, verbose=True) # train with earlystopping
    gin.test()
    print(gin.predict())