rule_mining.py

from __future__ import annotations
from collections import defaultdict
from dataclasses import dataclass
from typing import Dict, List, Tuple, Set

from .kg import KnowledgeGraph


@dataclass(frozen=True)
class Rule:
    # body: (r1(x,y), r2(y,z)) -> head: r3(x,z)
    r1: str
    r2: str
    head: str

    def __str__(self) -> str:
        return f"{self.r1}(x,y) ∧ {self.r2}(y,z) → {self.head}(x,z)"


@dataclass(frozen=True)
class UnaryRule:
    # body: r1(x,y) -> head: r2(x,y) (subrelation), or symmetry: r(x,y) -> r(y,x)
    body: str
    head: str
    symmetry: bool = False

    def __str__(self) -> str:
        if self.symmetry:
            return f"{self.body}(x,y) → {self.head}(y,x)"
        return f"{self.body}(x,y) → {self.head}(x,y)"


@dataclass
class RuleStat:
    support: int
    body_count: int

    @property
    def confidence(self) -> float:
        return 0.0 if self.body_count == 0 else self.support / self.body_count


def mine_binary_path_rules(kg: KnowledgeGraph, min_support: int = 1, min_conf: float = 0.2) -> Dict[Rule, RuleStat]:
    # Optimized counting using adjacency index: out_index[r][h] = {t}
    out_index: Dict[str, Dict[str, Set[str]]] = defaultdict(lambda: defaultdict(set))
    for r, pairs in kg.facts.items():
        for h, t in pairs:
            out_index[r][h].add(t)

    # Count body occurrences (unique x,z per r1,r2) and head occurrences
    aggregate_body: Dict[Tuple[str, str, str], Set[Tuple[str, str]]] = defaultdict(set)  # (r1, r2, head) -> {(x,z)}
    aggregate_support: Dict[Tuple[str, str, str], Set[Tuple[str, str]]] = defaultdict(set)

    rels = list(kg.facts.keys())
    for r1, pairs1 in kg.facts.items():
        for x, y in pairs1:
            # For each possible r2, traverse y -> z
            for r2 in rels:
                zs = out_index[r2].get(y, ())
                if not zs:
                    continue
                for z in zs:
                    # Record body for all potential heads
                    for head in rels:
                        aggregate_body[(r1, r2, head)].add((x, z))
                        if (x, z) in kg.facts[head]:
                            aggregate_support[(r1, r2, head)].add((x, z))

    # Aggregate support over distinct (x,z) bindings
    stats: Dict[Rule, RuleStat] = {}
    for (r1, r2, head), body_pairs in aggregate_body.items():
        body_c = len(body_pairs)
        supp = len(aggregate_support.get((r1, r2, head), set()))
        if body_c == 0:
            continue
        conf = supp / body_c
        if supp >= min_support and conf >= min_conf:
            stats[Rule(r1, r2, head)] = RuleStat(support=supp, body_count=body_c)
    return stats


def mine_unary_rules(kg: KnowledgeGraph, min_support: int = 1, min_conf: float = 0.3) -> Dict[UnaryRule, RuleStat]:
    stats: Dict[UnaryRule, RuleStat] = {}
    # Symmetry candidates: if many pairs (x,y) also have (y,x)
    for r, pairs in kg.facts.items():
        body_c = len(pairs)
        supp = sum(1 for (x, y) in pairs if (y, x) in pairs)
        rule = UnaryRule(body=r, head=r, symmetry=True)
        if supp >= min_support and (body_c == 0 or supp / body_c >= min_conf):
            stats[rule] = RuleStat(support=supp, body_count=body_c)

    # Subrelation candidates: if r1 pairs are mostly included in r2
    rels = list(kg.facts.keys())
    for i, r1 in enumerate(rels):
        for r2 in rels:
            if r1 == r2:
                continue
            pairs1 = kg.facts[r1]
            pairs2 = kg.facts[r2]
            if not pairs1:
                continue
            supp = sum(1 for p in pairs1 if p in pairs2)
            body_c = len(pairs1)
            if supp >= min_support and supp / body_c >= min_conf:
                stats[UnaryRule(body=r1, head=r2, symmetry=False)] = RuleStat(support=supp, body_count=body_c)

    return stats


def apply_unary(rule: UnaryRule, h: str, t: str) -> Tuple[str, str]:
    # Given a pair (h,t) for body relation, return the implied pair for head relation
    if rule.symmetry:
        return (t, h)
    return (h, t)