import torch
import torch.nn as nn
import torch.nn.functional as F

__all__ = ["GNN", "LinkPredictor", "GNNAsymm"]

class GNNLayer(nn.Module):
    """Graph Neural Network (GNN) / Message Passing Neural Network (MPNN) Layer.

    Parameters
    ----------
    hidden_X : int
        Hidden size for the node attributes.
    hidden_Y : int
        Hidden size for the node label.
    hidden_t : int
        Hidden size for the normalized time step.
    dropout : float
        Dropout rate.
    """
    def __init__(self,
                 hidden_X,
                 hidden_Y,
                 hidden_t,
                 dropout):
        super().__init__()

        self.update_X = nn.Sequential(
            nn.Linear(hidden_X + hidden_Y + hidden_t, hidden_X),
            nn.ReLU(),
            nn.LayerNorm(hidden_X),
            nn.Dropout(dropout)
        )
        self.update_Y = nn.Sequential(
            nn.Linear(hidden_Y, hidden_Y),
            nn.ReLU(),
            nn.LayerNorm(hidden_Y),
            nn.Dropout(dropout)
        )

    def forward(self, A, h_X, h_Y, h_t):
        """
        Parameters
        ----------
        A : dglsp.SparseMatrix
            Adjacency matrix.
        h_X : torch.Tensor of shape (|V|, hidden_X)
            Hidden representations for the node attributes.
        h_Y : torch.Tensor of shape (|V|, hidden_Y)
            Hidden representations for the node label.
        h_t : torch.Tensor of shape (|V|, hidden_t)
            Hidden representations for the normalized time step.

        Returns
        -------
        h_X : torch.Tensor of shape (|V|, hidden_X)
            Updated hidden representations for the node attributes.
        h_Y : torch.Tensor of shape (|V|, hidden_Y)
            Updated hidden representations for the node label.
        """
        h_aggr_X = A @ torch.cat([h_X, h_Y], dim=1)
        h_aggr_Y = A @ h_Y

        num_nodes = h_X.size(0)
        h_t_expand = h_t.expand(num_nodes, -1)
        h_aggr_X = torch.cat([h_aggr_X, h_t_expand], dim=1)

        h_X = self.update_X(h_aggr_X)
        h_Y = self.update_Y(h_aggr_Y)

        return h_X, h_Y

class GNNTower(nn.Module):
    """Graph Neural Network (GNN) / Message Passing Neural Network (MPNN).

    Parameters
    ----------
    num_attrs_X : int
        Number of node attributes.
    num_classes_X : int
        Number of classes for each node attribute.
    num_classes_Y : int
        Number of classes for node label.
    hidden_t : int
        Hidden size for the normalized time step.
    hidden_X : int
        Hidden size for the node attributes.
    hidden_Y : int
        Hidden size for the node label.
    out_size : int
        Output size of the final MLP layer.
    num_gnn_layers : int
        Number of GNN/MPNN layers.
    dropout : float
        Dropout rate.
    node_mode : bool
        Whether the encoder is used for node attribute prediction or structure
        prediction.
    """
    def __init__(self,
                 num_attrs_X,
                 num_classes_X,
                 num_classes_Y,
                 hidden_t,
                 hidden_X,
                 hidden_Y,
                 out_size,
                 num_gnn_layers,
                 dropout,
                 node_mode):
        super().__init__()

        in_X = num_attrs_X * num_classes_X
        self.num_attrs_X = num_attrs_X
        self.num_classes_X = num_classes_X

        self.mlp_in_t = nn.Sequential(
            nn.Linear(1, hidden_t),
            nn.ReLU(),
            nn.Linear(hidden_t, hidden_t),
            nn.ReLU())
        self.mlp_in_X = nn.Sequential(
            nn.Linear(in_X, hidden_X),
            nn.ReLU(),
            nn.Linear(hidden_X, hidden_X),
            nn.ReLU()
        )
        self.emb_Y = nn.Embedding(num_classes_Y, hidden_Y)

        self.gnn_layers = nn.ModuleList([
            GNNLayer(hidden_X,
                     hidden_Y,
                     hidden_t,
                     dropout)
            for _ in range(num_gnn_layers)])

        # +1 for the input attributes
        hidden_cat = (num_gnn_layers + 1) * (hidden_X + hidden_Y) + hidden_t
        self.mlp_out = nn.Sequential(
            nn.Linear(hidden_cat, hidden_cat),
            nn.ReLU(),
            nn.Linear(hidden_cat, out_size)
        )

        self.node_mode = node_mode

    def forward(self,
                t_float,
                X_t_one_hot,
                Y_real,
                A_t):
        # Input projection.
        # (1, hidden_t)
        h_t = self.mlp_in_t(t_float).unsqueeze(0)
        h_X = self.mlp_in_X(X_t_one_hot)
#        h_Y = self.emb_Y(Y_real)
        h_Y = torch.mm(Y_real.softmax(-1), self.emb_Y(torch.LongTensor([x for x in range(0,Y_real.size()[-1])]).cuda(1)))

        h_X_list = [h_X]
        h_Y_list = [h_Y]
        for gnn in self.gnn_layers:
            h_X, h_Y = gnn(A_t, h_X, h_Y, h_t)
            h_X_list.append(h_X)
            h_Y_list.append(h_Y)

        # (|V|, hidden_t)
        h_t = h_t.expand(h_X.size(0), -1)
        h_cat = torch.cat(h_X_list + h_Y_list + [h_t], dim=1)

        if self.node_mode:
            # (|V|, F * C_X)
            logit = self.mlp_out(h_cat)
            # (|V|, F, C_X)
            logit = logit.reshape(Y_real.size(0), self.num_attrs_X, -1)

            return logit
        else:
            return self.mlp_out(h_cat)

class LinkPredictor(nn.Module):
    """Model for structure prediction.

    Parameters
    ----------
    num_attrs_X : int
        Number of node attributes.
    num_classes_X : int
        Number of classes for each node attribute.
    num_classes_Y : int
        Number of classes for node label.
    num_classes_E : int
        Number of edge classes.
    hidden_t : int
        Hidden size for the normalized time step.
    hidden_X : int
        Hidden size for the node attributes.
    hidden_Y : int
        Hidden size for the node label.
    hidden_E : int
        Hidden size for the edges.
    num_gnn_layers : int
        Number of GNN/MPNN layers.
    dropout : float
        Dropout rate.
    """
    def __init__(self,
                 num_attrs_X,
                 num_classes_X,
                 num_classes_Y,
                 num_classes_E,
                 hidden_t,
                 hidden_X,
                 hidden_Y,
                 hidden_E,
                 num_gnn_layers,
                 dropout):
        super().__init__()

        self.gnn_encoder = GNNTower(num_attrs_X,
                                    num_classes_X,
                                    num_classes_Y,
                                    hidden_t,
                                    hidden_X,
                                    hidden_Y,
                                    hidden_E,
                                    num_gnn_layers,
                                    dropout,
                                    node_mode=False)
        self.mlp_out = nn.Sequential(
            nn.Linear(hidden_E, hidden_E),
            nn.ReLU(),
            nn.Linear(hidden_E, num_classes_E)
        )

    def forward(self,
                t_float,
                X_t_one_hot,
                Y_real,
                A_t,
                src,
                dst):
        # (|V|, hidden_E)
        h = self.gnn_encoder(t_float,
                             X_t_one_hot,
                             Y_real,
                             A_t)
        # (|E|, hidden_E)
        h = h[src] * h[dst]
        # (|E|, num_classes_E)
        logit = self.mlp_out(h)

        return logit

class GNN(nn.Module):
    """P(X|Y, X^t, A^t) + P(A|Y, X^t, A^t)

    Parameters
    ----------
    num_attrs_X : int
        Number of node attributes.
    num_classes_X : int
        Number of classes for each node attribute.
    num_classes_Y : int
        Number of classes for node label.
    num_classes_E : int
        Number of edge classes.
    gnn_X_config : dict
        Configuration of the GNN for reconstructing node attributes.
    gnn_E_config : dict
        Configuration of the GNN for reconstructing edges.
    """
    def __init__(self,
                 num_attrs_X,
                 num_classes_X,
                 num_classes_Y,
                 num_classes_E,
                 gnn_X_config,
                 gnn_E_config):
        super().__init__()

        self.pred_X = GNNTower(num_attrs_X,
                               num_classes_X,
                               num_classes_Y,
                               out_size=num_attrs_X * num_classes_X,
                               node_mode=True,
                               **gnn_X_config)

        self.pred_E = LinkPredictor(num_attrs_X,
                                    num_classes_X,
                                    num_classes_Y,
                                    num_classes_E,
                                    **gnn_E_config)

    def forward(self,
                t_float,
                X_t_one_hot,
                Y,
                A_t,
                batch_src,
                batch_dst):
        """
        Parameters
        ----------
        t_float : torch.Tensor of shape (1)
            Sampled timestep divided by self.T.
        X_t_one_hot : torch.Tensor of shape (|V|, 2 * F)
            One-hot encoding of the sampled node attributes.
        Y : torch.Tensor of shape (|V|)
            Categorical node labels.
        A_t : dglsp.SparseMatrix
            Row-normalized sampled adjacency matrix.
        batch_src : torch.LongTensor of shape (B)
            Source node IDs for a batch of candidate edges (node pairs).
        batch_dst : torch.LongTensor of shape (B)
            Destination node IDs for a batch of candidate edges (node pairs).

        Returns
        -------
        logit_X : torch.Tensor of shape (|V|, F, 2)
            Predicted logits for the node attributes.
        logit_E : torch.Tensor of shape (B, 2)
            Predicted logits for the edge existence.
        """
        logit_X = self.pred_X(t_float,
                              X_t_one_hot,
                              Y,
                              A_t)

        logit_E = self.pred_E(t_float,
                              X_t_one_hot,
                              Y,
                              A_t,
                              batch_src,
                              batch_dst)

        return logit_X, logit_E

class MLPLayer(nn.Module):
    """
    Parameters
    ----------
    hidden_X : int
        Hidden size for the node attributes.
    hidden_Y : int
        Hidden size for the node labels.
    hidden_t : int
        Hidden size for the normalized time step.
    dropout : float
        Dropout rate.
    """
    def __init__(self,
                 hidden_X,
                 hidden_Y,
                 hidden_t,
                 dropout):
        super().__init__()

        self.update_X = nn.Sequential(
            nn.Linear(hidden_X + hidden_Y + hidden_t, hidden_X),
            nn.ReLU(),
            nn.LayerNorm(hidden_X),
            nn.Dropout(dropout)
        )
        self.update_Y = nn.Sequential(
            nn.Linear(hidden_Y, hidden_Y),
            nn.ReLU(),
            nn.LayerNorm(hidden_Y),
            nn.Dropout(dropout)
        )

    def forward(self, h_X, h_Y, h_t):
        """
        Parameters
        ----------
        h_X : torch.Tensor of shape (|V|, hidden_X)
            Hidden representations for the node attributes.
        h_Y : torch.Tensor of shape (|V|, hidden_Y)
            Hidden representations for the node labels.
        h_t : torch.Tensor of shape (1, hidden_t)
            Hidden representations for the normalized time step.

        Returns
        -------
        h_X : torch.Tensor of shape (|V|, hidden_X)
            Updated hidden representations for the node attributes.
        h_Y : torch.Tensor of shape (|V|, hidden_Y)
            Updated hidden representations for the node labels.
        """
        num_nodes = h_X.size(0)
        h_t_expand = h_t.expand(num_nodes, -1)
        h_X = torch.cat([h_X, h_Y, h_t_expand], dim=1)

        h_X = self.update_X(h_X)
        h_Y = self.update_Y(h_Y)

        return h_X, h_Y

class MLPTower(nn.Module):
    def __init__(self,
                 num_attrs_X,
                 num_classes_X,
                 num_classes_Y,
                 hidden_t,
                 hidden_X,
                 hidden_Y,
                 num_mlp_layers,
                 dropout):
        super().__init__()

        in_X = num_attrs_X * num_classes_X
        self.num_attrs_X = num_attrs_X
        self.num_classes_X = num_classes_X

        self.mlp_in_t = nn.Sequential(
            nn.Linear(1, hidden_t),
            nn.ReLU(),
            nn.Linear(hidden_t, hidden_t),
            nn.ReLU())
        self.mlp_in_X = nn.Sequential(
            nn.Linear(in_X, hidden_X),
            nn.ReLU(),
            nn.Linear(hidden_X, hidden_X),
            nn.ReLU()
        )
        self.emb_Y = nn.Embedding(num_classes_Y, hidden_Y)

        self.mlp_layers = nn.ModuleList([
            MLPLayer(hidden_X,
                     hidden_Y,
                     hidden_t,
                     dropout)
            for _ in range(num_mlp_layers)])

        # +1 for the input features
        hidden_cat = (num_mlp_layers + 1) * (hidden_X + hidden_Y) + hidden_t
        self.mlp_out = nn.Sequential(
            nn.Linear(hidden_cat, hidden_cat),
            nn.ReLU(),
            nn.Linear(hidden_cat, in_X)
        )

    def forward(self,
                t_float,
                X_t_one_hot,
                Y_real):
        # Input projection.
        h_t = self.mlp_in_t(t_float).unsqueeze(0)
        h_X = self.mlp_in_X(X_t_one_hot)
#        h_Y = self.emb_Y(Y_real)
        h_Y = torch.mm(Y_real.cuda(1).softmax(-1), self.emb_Y(torch.LongTensor([x for x in range(0,Y_real.size()[-1])]).cuda(1)).cuda(1))
                       

        h_X_list = [h_X]
        h_Y_list = [h_Y]
        for mlp in self.mlp_layers:
            h_X, h_Y = mlp(h_X, h_Y, h_t)
            h_X_list.append(h_X)
            h_Y_list.append(h_Y)

        h_t = h_t.expand(h_X.size(0), -1)
        h_cat = torch.cat(h_X_list + h_Y_list + [h_t], dim=1)

        logit = self.mlp_out(h_cat)
        # (|V|, F, C)
        logit = logit.reshape(Y_real.size(0), self.num_attrs_X, -1)

        return logit

class GNNAsymm(nn.Module):
    """P(X|Y, X_t) + P(A|Y, X, A_t)

    Parameters
    ----------
    num_attrs_X : int
        Number of node attributes.
    num_classes_X : int
        Number of classes for each node attribute.
    num_classes_Y : int
        Number of classes for node label.
    num_classes_E : int
        Number of edge classes.
    mlp_X_config : dict
        Configuration of the MLP for reconstructing node attributes.
    gnn_E_config : dict
        Configuration of the GNN for reconstructing edges.
    """
    def __init__(self,
                 num_attrs_X,
                 num_classes_X,
                 num_classes_Y,
                 num_classes_E,
                 data_train_mask,
                 mlp_X_config,
                 gnn_E_config):
        super().__init__()

        self.pred_X = MLPTower(num_attrs_X,
                               num_classes_X,
                               num_classes_Y,
                               **mlp_X_config)

        self.pred_E = LinkPredictor(num_attrs_X,
                                    num_classes_X,
                                    num_classes_Y,
                                    num_classes_E,
                                    **gnn_E_config)
        self.data_train_mask = data_train_mask
        self.classifier = nn.Linear(128, num_classes_Y)

    def cont_loss(self, h, label):
        y_max = label.max()+1
        h = F.normalize(h, p=2, dim=1)
        idx = torch.randperm(h.size()[0])[:h.size()[0]//5]
        h, label = h[idx], label[idx]

        loss = 0
        for i in range(h.size()[0]):
            idx = (label[i] == label).nonzero().squeeze()
            loss += (torch.exp((h[i] * h[idx]).sum(1)).sum() / torch.exp((h[i] * h).sum(1)).sum()) / idx.size()[0]
        return loss / h.size()[0]

    def forward(self,
                t_float_X,
                t_float_E,
                X_t_one_hot,
                Y,
                X_one_hot_2d,
                A_t,
                batch_src,
                batch_dst):
        """
        Parameters
        ----------
        t_float_X : torch.Tensor of shape (1)
            Sampled timestep divided by self.T_X.
        t_float_E : torch.Tensor of shape (1)
            Sampled timestep divided by self.T_E.
        X_t_one_hot : torch.Tensor of shape (|V|, 2 * F)
            One-hot encoding of the sampled node attributes.
        Y : torch.Tensor of shape (|V|)
            Categorical node labels.
        X_one_hot_2d : torch.Tensor of shape (|V|, 2 * F)
            Flattened one-hot encoding of the node attributes.
        A_t : dglsp.SparseMatrix
            Row-normalized sampled adjacency matrix.
        batch_src : torch.LongTensor of shape (B)
            Source node IDs for a batch of candidate edges (node pairs).
        batch_dst : torch.LongTensor of shape (B)
            Destination node IDs for a batch of candidate edges (node pairs).

        Returns
        -------
        logit_X : torch.Tensor of shape (|V|, F, 2)
            Predicted logits for the node attributes.
        logit_E : torch.Tensor of shape (B, 2)
            Predicted logits for the edge existence.
        """
        logit_X = self.pred_X(t_float_X,
                              X_t_one_hot,
                              Y)

        logit_E = self.pred_E(t_float_E,
                              X_one_hot_2d,
                              Y,
                              A_t,
                              batch_src,
                              batch_dst)
        
        h = self.pred_E.gnn_encoder(t_float_E,
                             X_one_hot_2d,
                             Y,
                             A_t)
        classification_loss = 0 #.1 * self.cont_loss(h, Y)
        return logit_X, logit_E, classification_loss
