dig.py

"""Directed Intent Graph primitives and the edge-aware models used by fastCEGS."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
from torch_geometric.data import Data
from torch_geometric.nn import SAGEConv


DEFAULT_POLICY = "default-forward"


def node_names(topology: Mapping) -> List[str]:
    return [n["name"] if isinstance(n, Mapping) else n for n in topology.get("nodes", [])]


def physical_edges(topology: Mapping) -> List[Tuple[str, str]]:
    result = []
    for edge in topology.get("edges", topology.get("links", [])):
        if isinstance(edge, str):
            source, target = edge.split("_", 1)
        else:
            source = edge["node1"]["name"]
            target = edge["node2"]["name"]
        if source != target and (source, target) not in result and (target, source) not in result:
            result.append((source, target))
    return result


def directed_edge(source: str, target: str, policy: str = DEFAULT_POLICY) -> Dict[str, str]:
    return {"source": source, "target": target, "policy": policy}


def normalize_dig(graph: Mapping, topology: Optional[Mapping] = None) -> Dict:
    """Normalize both the old CEGS intent graph and the fastCEGS DIG schema."""
    nodes = dict(graph.get("nodes", {}))
    raw_edges = graph.get("edges", [])
    edges: List[Dict[str, str]] = []
    if isinstance(raw_edges, Mapping):
        raw_edges = [dict(source=k.split("_", 1)[0], target=k.split("_", 1)[1], policy=v)
                     for k, v in raw_edges.items()]
    for edge in raw_edges:
        if isinstance(edge, str):
            source, target = edge.split("_", 1)
            edges.extend((directed_edge(source, target), directed_edge(target, source)))
        else:
            source = edge.get("source", edge.get("node1", {}).get("name"))
            target = edge.get("target", edge.get("node2", {}).get("name"))
            if source is None or target is None:
                continue
            policy = edge.get("policy", DEFAULT_POLICY)
            edges.append(directed_edge(source, target, policy))
            if "source" not in edge:  # physical topology-style edge
                edges.append(directed_edge(target, source, policy))
    if not edges and topology is not None:
        for source, target in physical_edges(topology):
            edges.extend((directed_edge(source, target), directed_edge(target, source)))
    # Parallel topology links must not distort edge aggregation or matching.
    unique = {(e["source"], e["target"]): e for e in edges}
    return {"nodes": nodes, "edges": list(unique.values())}


def outgoing(graph: Mapping, source: str) -> List[Dict[str, str]]:
    return [e for e in normalize_dig(graph)["edges"] if e["source"] == source]


def policy_lookup(graph: Mapping) -> Dict[Tuple[str, str], str]:
    return {(e["source"], e["target"]): e["policy"] for e in normalize_dig(graph)["edges"]}


def _encode(text_model, text: str) -> np.ndarray:
    value = text_model.encode(text)
    if hasattr(value, "detach"):
        value = value.detach().cpu().numpy()
    return np.asarray(value, dtype=np.float32).reshape(-1)


def create_dig_data(graph: Mapping, text_model) -> Data:
    """Encode Role and Policy text while retaining the direction of every edge."""
    graph = normalize_dig(graph)
    names = list(graph["nodes"])
    if not names:
        raise ValueError("A directed intent graph must contain at least one node")
    indices = {name: index for index, name in enumerate(names)}
    x = torch.tensor(np.stack([_encode(text_model, graph["nodes"][n]) for n in names]))
    pairs, features = [], []
    zero = np.zeros(x.shape[1], dtype=np.float32)
    for edge in graph["edges"]:
        if edge["source"] in indices and edge["target"] in indices:
            pairs.append([indices[edge["source"]], indices[edge["target"]]])
            features.append(_encode(text_model, edge.get("policy", DEFAULT_POLICY)))
    edge_index = (torch.tensor(pairs, dtype=torch.long).t().contiguous()
                  if pairs else torch.empty((2, 0), dtype=torch.long))
    edge_attr = (torch.tensor(np.stack(features)) if features
                 else torch.empty((0, len(zero)), dtype=torch.float32))
    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
    data.node_names = names
    return data


def unique_outgoing_edge_mean(x: torch.Tensor, edge_index: torch.Tensor,
                              edge_attr: torch.Tensor) -> torch.Tensor:
    """Equation 2 EdgeAgg: mean unique outgoing policy vectors per node."""
    result = torch.zeros_like(x)
    for index in range(x.shape[0]):
        mask = edge_index[0] == index
        if bool(mask.any()):
            result[index] = torch.unique(edge_attr[mask], dim=0).mean(dim=0)
    return result


class EdgeInitializer(nn.Module):
    def __init__(self, text_dim: int, hidden_size: int):
        super().__init__()
        self.mlp = nn.Sequential(nn.Linear(text_dim * 2, hidden_size), nn.ReLU(),
                                 nn.Linear(hidden_size, hidden_size))

    def forward(self, x, edge_index, edge_attr):
        edge_mean = unique_outgoing_edge_mean(x, edge_index, edge_attr)
        return self.mlp(torch.cat((x, edge_mean), dim=-1))


class EGNN(nn.Module):
    """Edge-augmented GraphSAGE from fastCEGS Equations 2, 10 and 11."""
    def __init__(self, text_dim=300, hidden_size=256, out_channels=128, num_layers=2):
        super().__init__()
        self.initializer = EdgeInitializer(text_dim, hidden_size)
        sizes = [hidden_size] + [hidden_size] * max(0, num_layers - 1) + [out_channels]
        self.convs = nn.ModuleList(SAGEConv(sizes[i], sizes[i + 1]) for i in range(num_layers))

    def forward(self, x, edge_index, edge_attr):
        x = self.initializer(x, edge_index, edge_attr)
        for index, conv in enumerate(self.convs):
            x = conv(x, edge_index)
            if index + 1 < len(self.convs):
                x = F.relu(x)
        return x


class EdgeBiasedAttention(nn.Module):
    """Dense multi-head attention with the dynamic edge gate in Equations 5-6."""
    def __init__(self, hidden_size: int, heads: int = 4):
        super().__init__()
        if hidden_size % heads:
            raise ValueError("hidden_size must be divisible by heads")
        self.heads, self.head_dim = heads, hidden_size // heads
        self.q = nn.Linear(hidden_size, hidden_size, bias=False)
        self.k = nn.Linear(hidden_size, hidden_size, bias=False)
        self.v = nn.Linear(hidden_size, hidden_size, bias=False)
        self.edge = nn.Linear(hidden_size, hidden_size, bias=False)
        self.gate = nn.Linear(hidden_size, heads)
        self.output = nn.Linear(hidden_size, hidden_size)

    def forward(self, x, edge_index, edge_features):
        n = x.shape[0]
        q = self.q(x).view(n, self.heads, self.head_dim).transpose(0, 1)
        k = self.k(x).view(n, self.heads, self.head_dim).transpose(0, 1)
        v = self.v(x).view(n, self.heads, self.head_dim).transpose(0, 1)
        scores = torch.einsum("hnd,hmd->hnm", q, k) / self.head_dim ** 0.5
        if edge_index.numel():
            projected = self.edge(edge_features).view(-1, self.heads, self.head_dim)
            gates = torch.sigmoid(self.gate(edge_features))
            for i, (source, target) in enumerate(edge_index.t().tolist()):
                scores[:, source, target] += gates[i] * (q[:, source] * projected[i]).sum(-1)
        weights = torch.softmax(scores, dim=-1)
        attended = torch.einsum("hnm,hmd->hnd", weights, v).transpose(0, 1).reshape(n, -1)
        return self.output(attended)


class EGFormer(nn.Module):
    """Graph-level encoder implementing fastCEGS Equations 2-7."""
    def __init__(self, text_dim=300, hidden_size=256, out_channels=128,
                 num_layers=2, heads=4):
        super().__init__()
        self.initializer = EdgeInitializer(text_dim, hidden_size)
        self.edge_projection = nn.Linear(text_dim, hidden_size)
        self.attention = nn.ModuleList(EdgeBiasedAttention(hidden_size, heads)
                                       for _ in range(num_layers))
        self.norm1 = nn.ModuleList(nn.LayerNorm(hidden_size) for _ in range(num_layers))
        self.ffn = nn.ModuleList(nn.Sequential(nn.Linear(hidden_size, hidden_size * 2), nn.ReLU(),
                                               nn.Linear(hidden_size * 2, hidden_size))
                                 for _ in range(num_layers))
        self.norm2 = nn.ModuleList(nn.LayerNorm(hidden_size) for _ in range(num_layers))
        self.pool_gate = nn.Linear(hidden_size, 1)
        self.output = nn.Linear(hidden_size, out_channels)

    def forward(self, x, edge_index, edge_attr):
        x = self.initializer(x, edge_index, edge_attr)
        edge_features = self.edge_projection(edge_attr)
        for attention, norm1, ffn, norm2 in zip(self.attention, self.norm1, self.ffn, self.norm2):
            x = norm1(x + attention(x, edge_index, edge_features))
            x = norm2(x + ffn(x))
        weights = torch.softmax(self.pool_gate(x).squeeze(-1), dim=0)
        return self.output((weights.unsqueeze(-1) * x).sum(dim=0))


def l1_similarity(left: torch.Tensor, right: torch.Tensor) -> float:
    return 1.0 / (1.0 + torch.linalg.vector_norm(left - right, ord=1).item())