ComputerAssignments/CA10_knowledge_graphs/src/reasoning/neural_reasoning.py

"""
Neural Reasoning for Knowledge Graphs
Implements multi-hop reasoning, path-based reasoning, and neural-symbolic integration
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Tuple, Optional, Set
import numpy as np
from collections import defaultdict, deque
import sys
import os

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
from core.knowledge_graph import KnowledgeGraph, Triple


class MultiHopReasoning(nn.Module):
    """
    Multi-hop reasoning module for knowledge graphs
    Performs reasoning over multiple hops to infer new facts
    """

    def __init__(
        self,
        num_entities: int,
        num_relations: int,
        embedding_dim: int,
        max_hops: int = 3,
        dropout: float = 0.3,
    ):
        super(MultiHopReasoning, self).__init__()

        self.num_entities = num_entities
        self.num_relations = num_relations
        self.embedding_dim = embedding_dim
        self.max_hops = max_hops

        # Entity and relation embeddings
        self.entity_embeddings = nn.Embedding(num_entities, embedding_dim)
        self.relation_embeddings = nn.Embedding(num_relations, embedding_dim)

        # Hop-specific transformations
        self.hop_transforms = nn.ModuleList(
            [nn.Linear(embedding_dim, embedding_dim) for _ in range(max_hops)]
        )

        # Attention mechanism for path aggregation
        self.path_attention = nn.MultiheadAttention(
            embedding_dim, num_heads=4, dropout=dropout
        )

        # Path composition
        self.path_composer = nn.LSTM(
            embedding_dim,
            embedding_dim,
            num_layers=2,
            batch_first=True,
            dropout=dropout,
        )

        # Final scoring
        self.scorer = nn.Sequential(
            nn.Linear(embedding_dim * 2, embedding_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embedding_dim, 1),
        )

        self._init_weights()

    def _init_weights(self):
        nn.init.xavier_uniform_(self.entity_embeddings.weight)
        nn.init.xavier_uniform_(self.relation_embeddings.weight)

    def forward(self, head_idx, relation_path, tail_idx):
        """
        Args:
            head_idx: Starting entity indices (batch_size,)
            relation_path: Sequence of relations (batch_size, path_length)
            tail_idx: Target entity indices (batch_size,)
        """
        batch_size = head_idx.size(0)

        # Get initial embeddings
        current_emb = self.entity_embeddings(head_idx)
        tail_emb = self.entity_embeddings(tail_idx)

        # Traverse the path
        path_embeddings = [current_emb]

        for hop in range(relation_path.size(1)):
            rel_emb = self.relation_embeddings(relation_path[:, hop])

            # Apply hop-specific transformation
            if hop < self.max_hops:
                rel_emb = self.hop_transforms[hop](rel_emb)

            # Compose with current embedding
            current_emb = current_emb + rel_emb
            path_embeddings.append(current_emb)

        # Stack path embeddings
        path_stack = torch.stack(path_embeddings, dim=1)  # (batch, path_len+1, dim)

        # Apply attention over path
        attended_path, _ = self.path_attention(path_stack, path_stack, path_stack)

        # Apply LSTM for path composition
        lstm_out, _ = self.path_composer(attended_path)
        final_path_emb = lstm_out[:, -1, :]  # Take last output

        # Combine with tail and score
        combined = torch.cat([final_path_emb, tail_emb], dim=-1)
        scores = self.scorer(combined).squeeze(-1)

        return scores

    def find_paths(
        self,
        kg: KnowledgeGraph,
        start_entity: str,
        end_entity: str,
        max_length: int = 3,
    ) -> List[List[str]]:
        """
        Find all paths between two entities using BFS

        Args:
            kg: Knowledge graph
            start_entity: Starting entity ID
            end_entity: Target entity ID
            max_length: Maximum path length

        Returns:
            List of paths (each path is a list of relation IDs)
        """
        if start_entity not in kg.entities or end_entity not in kg.entities:
            return []

        paths = []
        queue = deque([(start_entity, [])])
        visited = set()

        while queue:
            current, path = queue.popleft()

            if len(path) > max_length:
                continue

            if current == end_entity and len(path) > 0:
                paths.append(path)
                continue

            if (current, tuple(path)) in visited:
                continue
            visited.add((current, tuple(path)))

            # Explore neighbors
            for relation, neighbors in kg.entity_adjacency[current].items():
                if relation.startswith("inv_"):
                    continue
                for neighbor in neighbors:
                    new_path = path + [relation]
                    queue.append((neighbor, new_path))

        return paths


class PathRankingReasoning(nn.Module):
    """
    Path Ranking Algorithm for Knowledge Graph Reasoning
    Learns to rank paths based on their relevance for link prediction
    """

    def __init__(
        self,
        num_relations: int,
        embedding_dim: int,
        max_path_length: int = 4,
        dropout: float = 0.3,
    ):
        super(PathRankingReasoning, self).__init__()

        self.num_relations = num_relations
        self.embedding_dim = embedding_dim
        self.max_path_length = max_path_length

        # Relation embeddings
        self.relation_embeddings = nn.Embedding(num_relations, embedding_dim)

        # Path encoder (Bi-LSTM)
        self.path_encoder = nn.LSTM(
            embedding_dim,
            embedding_dim // 2,
            num_layers=2,
            batch_first=True,
            bidirectional=True,
            dropout=dropout,
        )

        # Path ranker
        self.path_ranker = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embedding_dim // 2, 1),
            nn.Sigmoid(),
        )

        self._init_weights()

    def _init_weights(self):
        nn.init.xavier_uniform_(self.relation_embeddings.weight)

    def forward(self, paths):
        """
        Args:
            paths: List of paths, each path is a tensor of relation indices (batch_size, max_path_len)
        """
        # Get relation embeddings for paths
        path_embs = self.relation_embeddings(paths)  # (batch, path_len, dim)

        # Encode paths
        encoded_paths, _ = self.path_encoder(path_embs)  # (batch, path_len, dim)

        # Pool path representations (mean pooling)
        pooled_paths = encoded_paths.mean(dim=1)  # (batch, dim)

        # Rank paths
        path_scores = self.path_ranker(pooled_paths).squeeze(-1)  # (batch,)

        return path_scores, pooled_paths


class NeuralSymbolicReasoning(nn.Module):
    """
    Neural-Symbolic Integration for Knowledge Graph Reasoning
    Combines neural learning with symbolic logic rules
    """

    def __init__(
        self,
        num_entities: int,
        num_relations: int,
        embedding_dim: int,
        num_rules: int = 50,
        dropout: float = 0.3,
    ):
        super(NeuralSymbolicReasoning, self).__init__()

        self.num_entities = num_entities
        self.num_relations = num_relations
        self.embedding_dim = embedding_dim
        self.num_rules = num_rules

        # Entity and relation embeddings
        self.entity_embeddings = nn.Embedding(num_entities, embedding_dim)
        self.relation_embeddings = nn.Embedding(num_relations, embedding_dim)

        # Rule embeddings (learnable logic rules)
        self.rule_embeddings = nn.Parameter(torch.randn(num_rules, embedding_dim))

        # Rule attention
        self.rule_attention = nn.MultiheadAttention(
            embedding_dim, num_heads=4, dropout=dropout
        )

        # Logic operators (AND, OR, NOT)
        self.logic_and = nn.Sequential(
            nn.Linear(embedding_dim * 2, embedding_dim), nn.Tanh()
        )
        self.logic_or = nn.Sequential(
            nn.Linear(embedding_dim * 2, embedding_dim), nn.Tanh()
        )
        self.logic_not = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim), nn.Tanh()
        )

        # Rule applicator
        self.rule_applicator = nn.Sequential(
            nn.Linear(embedding_dim * 3, embedding_dim * 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embedding_dim * 2, embedding_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embedding_dim, 1),
        )

        self._init_weights()

    def _init_weights(self):
        nn.init.xavier_uniform_(self.entity_embeddings.weight)
        nn.init.xavier_uniform_(self.relation_embeddings.weight)
        nn.init.xavier_uniform_(self.rule_embeddings)

    def forward(self, head_idx, relation_idx, tail_idx):
        """
        Args:
            head_idx: Head entity indices (batch_size,)
            relation_idx: Relation indices (batch_size,)
            tail_idx: Tail entity indices (batch_size,)
        """
        batch_size = head_idx.size(0)

        # Get embeddings
        head_emb = self.entity_embeddings(head_idx)
        rel_emb = self.relation_embeddings(relation_idx)
        tail_emb = self.entity_embeddings(tail_idx)

        # Prepare for attention
        triple_emb = torch.stack(
            [head_emb, rel_emb, tail_emb], dim=1
        )  # (batch, 3, dim)
        rules = self.rule_embeddings.unsqueeze(0).expand(
            batch_size, -1, -1
        )  # (batch, num_rules, dim)

        # Apply rule attention
        attended_rules, attention_weights = self.rule_attention(
            triple_emb, rules, rules, average_attn_weights=True
        )

        # Aggregate attended rules
        aggregated_rules = attended_rules.mean(dim=1)  # (batch, dim)

        # Combine with triple information
        combined = torch.cat([head_emb, rel_emb, tail_emb], dim=-1)

        # Apply rules
        scores = self.rule_applicator(combined).squeeze(-1)

        return scores, attention_weights

    def apply_logic_operator(self, emb1, emb2, operator="and"):
        """Apply logic operators to embeddings"""
        if operator == "and":
            return self.logic_and(torch.cat([emb1, emb2], dim=-1))
        elif operator == "or":
            return self.logic_or(torch.cat([emb1, emb2], dim=-1))
        elif operator == "not":
            return self.logic_not(emb1)
        else:
            raise ValueError(f"Unknown operator: {operator}")


class ReasoningChain(nn.Module):
    """
    Chain-of-Thought Reasoning for Knowledge Graphs
    Implements step-by-step reasoning with intermediate steps
    """

    def __init__(
        self,
        num_entities: int,
        num_relations: int,
        embedding_dim: int,
        max_reasoning_steps: int = 5,
        dropout: float = 0.3,
    ):
        super(ReasoningChain, self).__init__()

        self.num_entities = num_entities
        self.num_relations = num_relations
        self.embedding_dim = embedding_dim
        self.max_reasoning_steps = max_reasoning_steps

        # Entity and relation embeddings
        self.entity_embeddings = nn.Embedding(num_entities, embedding_dim)
        self.relation_embeddings = nn.Embedding(num_relations, embedding_dim)

        # Reasoning step encoder
        self.step_encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=embedding_dim,
                nhead=8,
                dim_feedforward=embedding_dim * 4,
                dropout=dropout,
                batch_first=True,
            ),
            num_layers=3,
        )

        # Step predictor
        self.step_predictor = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embedding_dim, num_relations),
        )

        # Confidence estimator
        self.confidence_estimator = nn.Sequential(
            nn.Linear(embedding_dim, embedding_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embedding_dim // 2, 1),
            nn.Sigmoid(),
        )

        self._init_weights()

    def _init_weights(self):
        nn.init.xavier_uniform_(self.entity_embeddings.weight)
        nn.init.xavier_uniform_(self.relation_embeddings.weight)

    def forward(self, start_entity, target_relation, target_entity):
        """
        Args:
            start_entity: Starting entity indices (batch_size,)
            target_relation: Target relation indices (batch_size,)
            target_entity: Target entity indices (batch_size,)
        """
        batch_size = start_entity.size(0)

        # Get initial embeddings
        current_emb = self.entity_embeddings(start_entity)
        target_rel_emb = self.relation_embeddings(target_relation)
        target_ent_emb = self.entity_embeddings(target_entity)

        # Reasoning chain
        reasoning_steps = [current_emb]
        predicted_relations = []
        confidences = []

        for step in range(self.max_reasoning_steps):
            # Stack reasoning history
            step_sequence = torch.stack(reasoning_steps, dim=1)  # (batch, step+1, dim)

            # Encode reasoning steps
            encoded_steps = self.step_encoder(step_sequence)

            # Get current state
            current_state = encoded_steps[:, -1, :]  # (batch, dim)

            # Predict next relation
            next_relation_logits = self.step_predictor(current_state)
            predicted_relations.append(next_relation_logits)

            # Estimate confidence
            confidence = self.confidence_estimator(current_state)
            confidences.append(confidence)

            # Update current embedding
            next_rel_idx = next_relation_logits.argmax(dim=-1)
            next_rel_emb = self.relation_embeddings(next_rel_idx)
            current_emb = current_emb + next_rel_emb

            reasoning_steps.append(current_emb)

        # Final score
        final_emb = reasoning_steps[-1]
        score = F.cosine_similarity(final_emb, target_ent_emb, dim=-1)

        return {
            "score": score,
            "reasoning_steps": reasoning_steps,
            "predicted_relations": predicted_relations,
            "confidences": confidences,
        }


class ExplainableReasoning(nn.Module):
    """
    Explainable Reasoning Module
    Provides interpretable explanations for predictions
    """

    def __init__(
        self,
        num_entities: int,
        num_relations: int,
        embedding_dim: int,
        dropout: float = 0.3,
    ):
        super(ExplainableReasoning, self).__init__()

        self.num_entities = num_entities
        self.num_relations = num_relations
        self.embedding_dim = embedding_dim

        # Entity and relation embeddings
        self.entity_embeddings = nn.Embedding(num_entities, embedding_dim)
        self.relation_embeddings = nn.Embedding(num_relations, embedding_dim)

        # Attention for explanation
        self.explanation_attention = nn.MultiheadAttention(
            embedding_dim, num_heads=4, dropout=dropout
        )

        # Evidence aggregator
        self.evidence_aggregator = nn.Sequential(
            nn.Linear(embedding_dim * 3, embedding_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(embedding_dim, embedding_dim // 2),
            nn.ReLU(),
            nn.Linear(embedding_dim // 2, 1),
        )

        self._init_weights()

    def _init_weights(self):
        nn.init.xavier_uniform_(self.entity_embeddings.weight)
        nn.init.xavier_uniform_(self.relation_embeddings.weight)

    def forward(self, head_idx, relation_idx, tail_idx, evidence_triples=None):
        """
        Args:
            head_idx: Head entity indices (batch_size,)
            relation_idx: Relation indices (batch_size,)
            tail_idx: Tail entity indices (batch_size,)
            evidence_triples: Supporting evidence triples (batch_size, num_evidence, 3)
        """
        # Get query triple embeddings
        head_emb = self.entity_embeddings(head_idx)
        rel_emb = self.relation_embeddings(relation_idx)
        tail_emb = self.entity_embeddings(tail_idx)

        query_emb = torch.stack([head_emb, rel_emb, tail_emb], dim=1)  # (batch, 3, dim)

        if evidence_triples is not None:
            batch_size, num_evidence, _ = evidence_triples.shape

            # Get evidence embeddings
            evidence_heads = self.entity_embeddings(evidence_triples[:, :, 0])
            evidence_rels = self.relation_embeddings(evidence_triples[:, :, 1])
            evidence_tails = self.entity_embeddings(evidence_triples[:, :, 2])

            evidence_emb = torch.stack(
                [evidence_heads, evidence_rels, evidence_tails], dim=2
            )
            evidence_emb = evidence_emb.view(batch_size, num_evidence * 3, -1)

            # Apply attention to find relevant evidence
            attended_evidence, attention_weights = self.explanation_attention(
                query_emb, evidence_emb, evidence_emb, average_attn_weights=False
            )

            # Aggregate evidence
            combined = torch.cat(
                [head_emb, rel_emb, tail_emb], dim=-1
            )  # (batch, dim*3)
            score = self.evidence_aggregator(combined).squeeze(-1)

            return score, attention_weights
        else:
            # No evidence, just score the triple
            combined = torch.cat([head_emb, rel_emb, tail_emb], dim=-1)
            score = self.evidence_aggregator(combined).squeeze(-1)

            return score, None

    def explain_prediction(
        self, kg: KnowledgeGraph, head: str, relation: str, tail: str, top_k: int = 5
    ) -> List[Tuple[Triple, float]]:
        """
        Generate explanation for a prediction

        Args:
            kg: Knowledge graph
            head: Head entity ID
            relation: Relation ID
            tail: Tail entity ID
            top_k: Number of supporting triples to return

        Returns:
            List of (triple, relevance_score) tuples
        """
        # Find supporting triples (simplified version)
        supporting_triples = []

        # Find triples involving the head entity
        head_triples = kg.query_triples(subject=head)
        for triple in head_triples[:top_k]:
            supporting_triples.append((triple, 1.0))  # Simplified relevance score

        # Find triples involving the tail entity
        tail_triples = kg.query_triples(object=tail)
        for triple in tail_triples[:top_k]:
            if len(supporting_triples) < top_k:
                supporting_triples.append((triple, 0.8))

        return supporting_triples[:top_k]