import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv


class VGAEEncoder(nn.Module):
    """
    GraphSAGE-based variational encoder for VGAE.

    SAGEConv uses mean aggregation, which is degree-agnostic and more stable
    than GCNConv's degree-normalised aggregation on sparse knowledge graphs.
    Outputs (mu, log_std) for the reparameterisation trick.
    """

    def __init__(self, in_channels: int, hidden_channels: int, out_channels: int, dropout: float = 0.1):
        super().__init__()
        self.conv1 = SAGEConv(in_channels, hidden_channels)
        self.conv_mu = SAGEConv(hidden_channels, out_channels)
        self.conv_logstd = SAGEConv(hidden_channels, out_channels)
        self.dropout = dropout
        self._init_weights()

    def _init_weights(self):
        # Xavier uniform initialisation to control activation magnitude from epoch 0.
        for conv in [self.conv1, self.conv_mu, self.conv_logstd]:
            nn.init.xavier_uniform_(conv.lin_l.weight)
            nn.init.xavier_uniform_(conv.lin_r.weight)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=self.dropout, training=self.training)
        mu = self.conv_mu(x, edge_index)
        # Clamp both directions: prevents underflow (exp(-10)≈0) and overflow (exp(10)≈22k).
        log_std = self.conv_logstd(x, edge_index).clamp(min=-10, max=10)
        return mu, log_std


class VGAEModel(nn.Module):
    def __init__(self, in_channels: int, hidden_channels: int, out_channels: int, dropout: float = 0.1):
        super().__init__()
        self.encoder = VGAEEncoder(in_channels, hidden_channels, out_channels, dropout)

    def encode(self, x, edge_index):
        mu, log_std = self.encoder(x, edge_index)
        # At training time use the reparameterisation trick; at inference use the
        # deterministic mean mu for reproducible retrieval scores.
        if self.training:
            std = torch.exp(log_std)
            z = mu + torch.randn_like(std) * std
        else:
            z = mu
        self._mu = mu
        self._log_std = log_std
        return z

    def decode(self, z, edge_index, sigmoid=True):
        value = (z[edge_index[0]] * z[edge_index[1]]).sum(dim=1)
        return torch.sigmoid(value) if sigmoid else value

    def kl_loss(self):
        """KL( N(mu, std) || N(0, 1) ) — call after encode()."""
        return -0.5 * torch.mean(
            1 + 2 * self._log_std - self._mu.pow(2) - (2 * self._log_std).exp()
        )

    def forward(self, x, edge_index):
        return self.encode(x, edge_index)


def get_model(in_channels: int = 384, hidden_channels: int = 256, out_channels: int = 128) -> VGAEModel:
    return VGAEModel(in_channels, hidden_channels, out_channels)