"""Gated local-global graph attention, single-head MVP."""
import torch
from torch import nn


def phi(x):
    return torch.nn.functional.elu(x) + 1.0


def local_attention(x, edge_src, edge_dst, wq, wk, wv):
    q, k, v = x @ wq, x @ wk, x @ wv
    r = q.shape[-1]
    logits = (q[edge_src] * k[edge_dst]).sum(-1) / (r ** 0.5)
    n = x.shape[0]
    mx = torch.full((n,), -torch.inf, device=x.device, dtype=x.dtype)
    mx.scatter_reduce_(0, edge_src, logits, reduce="amax", include_self=True)
    weights_un = torch.exp(logits - mx[edge_src])
    den = torch.zeros(n, device=x.device, dtype=x.dtype)
    den.scatter_add_(0, edge_src, weights_un)
    weights = weights_un / den[edge_src].clamp_min(1e-12)
    out = torch.zeros((n, v.shape[-1]), device=x.device, dtype=x.dtype)
    out.index_add_(0, edge_src, weights[:, None] * v[edge_dst])
    return out


def global_linear_attention(x, wq, wk, wv, eps=1e-6):
    q, k, v = x @ wq, x @ wk, x @ wv
    pq, pk = phi(q), phi(k)
    # S = sum_j outer(phi(k_j), v_j), z = sum_j phi(k_j)
    s = pk.transpose(0, 1) @ v
    z = pk.sum(0)
    return (pq @ s) / ((pq @ z)[:, None] + eps)


class GatedLocalGlobalAttention(nn.Module):
    def __init__(self, d, r=None):
        super().__init__()
        r = r or d
        self.wq, self.wk, self.wv = nn.Linear(d, r, bias=False), nn.Linear(d, r, bias=False), nn.Linear(d, r, bias=False)
        self.out = nn.Linear(r, d)
        self.gate = nn.Linear(d, 1)

    def forward(self, x, edge_src, edge_dst, return_gate=False):
        l = local_attention(x, edge_src, edge_dst, self.wq.weight.T, self.wk.weight.T, self.wv.weight.T)
        g = global_linear_attention(x, self.wq.weight.T, self.wk.weight.T, self.wv.weight.T)
        gate = torch.sigmoid(self.gate(x))
        y = self.out(gate * l + (1.0 - gate) * g)
        return (y, gate.squeeze(-1)) if return_gate else y


def dense_attention(x, wq, wk, wv):
    q, k, v = x @ wq, x @ wk, x @ wv
    a = torch.softmax(q @ k.T / (q.shape[-1] ** 0.5), dim=-1)
    return a @ v