import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GATConv, LayerNorm


class CorporateActionAwareGNN(nn.Module):
    def __init__(
        self,
        num_features: int,
        hidden_channels: int = 64,
        num_heads: int = 8,
        dropout: float = 0.6,
    ):
        super().__init__()
        self.conv1 = GATConv(
            num_features,
            hidden_channels,
            heads=num_heads,
            dropout=dropout,
            concat=True,
        )
        self.norm1 = LayerNorm(hidden_channels * num_heads)
        self.conv2 = GATConv(
            hidden_channels * num_heads,
            hidden_channels,
            heads=num_heads,
            dropout=dropout,
            concat=True,
        )
        self.norm2 = LayerNorm(hidden_channels * num_heads)
        self.fc = nn.Linear(hidden_channels * num_heads, 1)
        self.dropout = nn.Dropout(dropout)

    def forward(self, data):
        x, edge_index = data.x, data.edge_index
        x = self.conv1(x, edge_index)
        x = self.norm1(x)
        x = F.elu(x)
        x = self.dropout(x)
        x = self.conv2(x, edge_index)
        x = self.norm2(x)
        x = F.elu(x)
        x = self.dropout(x)
        x = self.fc(x)
        return x