import torch
import torch.nn as nn
from torch_geometric.nn import GINEConv, GPSConv


class EventGraphFM(nn.Module):
    def __init__(self, n_act: int, n_res: int, n_logs: int, d_model: int = 256,
                 n_layers: int = 6, n_heads: int = 8, act_dim: int = 64, res_dim: int = 32,
                 log_dim: int = 16, dropout: float = 0.1, edge_dim: int = 2, num_dim: int = 7):
        super().__init__()
        self.act_emb = nn.Embedding(n_act, act_dim, padding_idx=0)
        self.res_emb = nn.Embedding(n_res, res_dim, padding_idx=0)
        self.log_emb = nn.Embedding(n_logs, log_dim)
        in_dim = act_dim + res_dim + log_dim + num_dim + 1
        self.proj = nn.Linear(in_dim, d_model)
        self.layers = nn.ModuleList()
        for _ in range(n_layers):
            mlp = nn.Sequential(
                nn.Linear(d_model, 2 * d_model), nn.ReLU(), nn.Linear(2 * d_model, d_model)
            )
            conv = GINEConv(mlp, edge_dim=edge_dim)
            self.layers.append(GPSConv(d_model, conv, heads=n_heads, dropout=dropout, act="relu"))
        self.norm = nn.LayerNorm(d_model)
        self.act_out = nn.Linear(d_model, n_act)
        self.res_out = nn.Linear(d_model, n_res)

    def forward(self, gact, gres, num, log_id, edge_index, edge_attr, batch_vec, mask_flag):
        x = torch.cat([
            self.act_emb(gact),
            self.res_emb(gres),
            self.log_emb(log_id),
            num,
            mask_flag.unsqueeze(-1),
        ], dim=-1)
        h = self.proj(x)
        for layer in self.layers:
            h = layer(h, edge_index, batch=batch_vec, edge_attr=edge_attr)
        h = self.norm(h)
        return self.act_out(h), self.res_out(h)