# import torch
# import torch.nn as nn
# from torch_geometric.nn import MessagePassing
# from torch.nn.utils import spectral_norm

# class GCNConvLayer(MessagePassing):
#     def __init__(self, in_channels, out_channels):
#         super().__init__(aggr="add")
#         self.lin_l = spectral_norm(nn.Linear(in_channels, out_channels, bias=False))
#         self.lin_r = spectral_norm(nn.Linear(in_channels, out_channels, bias=False))

#     def forward(self, x, edge_index, edge_attr, deg):
#         out = self.propagate(edge_index, x=x, edge_attr=edge_attr, deg=deg)
#         out = self.lin_l(out)
#         return out

#     def message(self, x_j, edge_attr):
#         return x_j * edge_attr

#     def aggregate(self, inputs, index, deg):
#         """
#         TODO: add docstring.
#         """
#         out = super().aggregate(inputs, index)
#         deg = deg + 1e-7
#         return out / deg.view(-1, 1)


# class CorrectionNet(nn.Module):
#     def __init__(self, hidden_dim=8, n_layers=1):
#         super().__init__()
#         # self.enc = GCNConvLayer(1, hidden_dim)
#         self.enc = nn.Sequential(
#             spectral_norm(nn.Linear(1, hidden_dim//2)),
#             nn.GELU(),
#             spectral_norm(nn.Linear(hidden_dim//2, hidden_dim)),
#         )
#         self.layers = torch.nn.ModuleList([GCNConvLayer(hidden_dim, hidden_dim) for _ in range(n_layers)])
#         self.relu = nn.GELU()

#         self.dec = nn.Sequential(
#             spectral_norm(nn.Linear(hidden_dim, hidden_dim//2)),
#             nn.GELU(),
#             spectral_norm(nn.Linear(hidden_dim//2, 1)),
#         )

#     def forward(self, x, edge_index, edge_attr, deg,):
#         # h = self.enc(x, edge_index, edge_attr, deg)
#         # h = self.relu(self.enc(x))
#         h = self.enc(x)
#         for layer in self.layers:
#             h = layer(h, edge_index, edge_attr, deg)
#             # h = self.norm(h)
#             h = self.relu(h)
#         # out = self.dec(h, edge_index, edge_attr, deg)
#         out = self.dec(h)
#         return out


import torch
import torch.nn as nn
from torch_geometric.nn import MessagePassing
from torch.nn.utils import spectral_norm


class CorrectionNet(MessagePassing):
    """
    TODO: add docstring.
    """

    def __init__(self, hidden_dim=16):
        super().__init__(aggr="add")
        self.in_net = nn.Sequential(
            spectral_norm(nn.Linear(1, hidden_dim // 2)),
            nn.GELU(),
            spectral_norm(nn.Linear(hidden_dim // 2, hidden_dim)),
        )

        self.out_net = nn.Sequential(
            spectral_norm(nn.Linear(hidden_dim, hidden_dim // 2)),
            nn.GELU(),
            spectral_norm(nn.Linear(hidden_dim // 2, 1)),
        )

        self.lin_msg = spectral_norm(
            nn.Linear(hidden_dim, hidden_dim, bias=False)
        )
        self.lin_update = spectral_norm(
            nn.Linear(hidden_dim, hidden_dim, bias=False)
        )
        self.alpha = nn.Parameter(torch.tensor(0.0))
        self.beta = nn.Parameter(torch.tensor(0.0))

    def forward(self, x, edge_index, edge_attr, deg):
        """
        TODO: add docstring.
        """
        x = self.in_net(x)
        out = self.propagate(edge_index, x=x, edge_attr=edge_attr, deg=deg)
        return self.out_net(out)

    def message(self, x_j, edge_attr):
        """
        TODO: add docstring.
        """
        alpha = torch.sigmoid(self.alpha)
        msg = x_j * edge_attr
        msg = (1 - alpha) * msg + alpha * self.lin_msg(msg)
        return msg

    def update(self, aggr_out, x):
        """
        TODO: add docstring.
        """
        beta = torch.sigmoid(self.beta)
        return aggr_out * (1 - beta) + self.lin_msg(x) * beta

    def aggregate(self, inputs, index, deg):
        """
        TODO: add docstring.
        """
        out = super().aggregate(inputs, index)
        deg = deg + 1e-7
        return out / deg.view(-1, 1)