rule_engine.py

"""
Rule-based inference engine with forward chaining and conflict resolution.
"""

from typing import List, Set, Dict, Any, Callable, Optional, Tuple
from collections import defaultdict
import copy


class Fact:
    """Represents a fact in the knowledge base."""
    
    def __init__(self, name: str, **attributes: Any) -> None:
        """
        Initialize a fact.
        
        Args:
            name: The name/type of the fact
            **attributes: Key-value pairs representing fact attributes
        """
        self.name = name
        self.attributes = attributes
    
    def __eq__(self, other: object) -> bool:
        if not isinstance(other, Fact):
            return False
        return self.name == other.name and self.attributes == other.attributes
    
    def __hash__(self) -> int:
        return hash((self.name, tuple(sorted(self.attributes.items()))))
    
    def __repr__(self) -> str:
        if not self.attributes:
            return f"Fact('{self.name}')"
        attrs = ', '.join(f"{k}={v}" for k, v in self.attributes.items())
        return f"Fact('{self.name}', {attrs})"


class Rule:
    """Represents an inference rule with conditions and conclusions."""
    
    def __init__(
        self, 
        name: str,
        conditions: List[Fact],
        conclusions: List[Fact],
        priority: int = 0
    ) -> None:
        """
        Initialize a rule.
        
        Args:
            name: Name of the rule
            conditions: List of facts that must be true for rule to fire
            conclusions: List of facts that become true when rule fires
            priority: Priority for conflict resolution (higher fires first)
        """
        self.name = name
        self.conditions = conditions
        self.conclusions = conclusions
        self.priority = priority
        self.fired = False
    
    def can_fire(self, facts: Set[Fact]) -> bool:
        """
        Check if rule can fire given current facts.
        
        Args:
            facts: Current set of known facts
            
        Returns:
            True if all conditions are met, False otherwise
        """
        if self.fired:
            return False
            
        for condition in self.conditions:
            if not self._fact_matches(condition, facts):
                return False
        return True
    
    def _fact_matches(self, condition: Fact, facts: Set[Fact]) -> bool:
        """Check if a condition fact matches any known fact."""
        for fact in facts:
            if fact.name == condition.name:
                # Check if all condition attributes match
                match = True
                for key, value in condition.attributes.items():
                    if key not in fact.attributes or fact.attributes[key] != value:
                        match = False
                        break
                if match:
                    return True
        return False
    
    def fire(self, facts: Set[Fact]) -> List[Fact]:
        """
        Fire the rule and return new facts.
        
        Args:
            facts: Current set of known facts
            
        Returns:
            List of new facts generated by this rule
        """
        if not self.can_fire(facts):
            return []
        
        self.fired = True
        new_facts = []
        for conclusion in self.conclusions:
            # Only add if not already present
            if conclusion not in facts:
                new_facts.append(conclusion)
        return new_facts
    
    def reset(self) -> None:
        """Reset the rule's fired status."""
        self.fired = False
    
    def __repr__(self) -> str:
        return f"Rule('{self.name}', priority={self.priority})"


class KnowledgeBase:
    """Manages facts and rules for inference."""
    
    def __init__(self) -> None:
        """Initialize an empty knowledge base."""
        self.facts: Set[Fact] = set()
        self.rules: List[Rule] = []
        self._rule_index: Dict[str, List[Rule]] = defaultdict(list)
    
    def add_fact(self, fact: Fact) -> None:
        """
        Add a fact to the knowledge base.
        
        Args:
            fact: The fact to add
        """
        self.facts.add(fact)
    
    def add_rule(self, rule: Rule) -> None:
        """
        Add a rule to the knowledge base.
        
        Args:
            rule: The rule to add
        """
        self.rules.append(rule)
        # Index rules by condition names for efficiency
        for condition in rule.conditions:
            self._rule_index[condition.name].append(rule)
    
    def get_applicable_rules(self) -> List[Rule]:
        """
        Get all rules that can currently fire.
        
        Returns:
            List of applicable rules sorted by priority
        """
        applicable = [rule for rule in self.rules if rule.can_fire(self.facts)]
        # Sort by priority (higher first) then by name for deterministic ordering
        applicable.sort(key=lambda r: (-r.priority, r.name))
        return applicable
    
    def forward_chain(self) -> Tuple[int, int]:
        """
        Perform forward chaining until no more rules can fire.
        
        Returns:
            Tuple of (iterations, facts_added) during inference
        """
        iterations = 0
        facts_added = 0
        
        while True:
            applicable_rules = self.get_applicable_rules()
            if not applicable_rules:
                break
            
            iterations += 1
            fired_in_iteration = 0
            
            for rule in applicable_rules:
                new_facts = rule.fire(self.facts)
                if new_facts:
                    fired_in_iteration += 1
                    for fact in new_facts:
                        self.facts.add(fact)
                    facts_added += len(new_facts)
            
            # If no rules fired in this iteration, we're done
            if fired_in_iteration == 0:
                break
        
        return iterations, facts_added
    
    def query(self, fact: Fact) -> bool:
        """
        Check if a fact can be derived from the knowledge base.
        
        Args:
            fact: The fact to query
            
        Returns:
            True if fact exists or can be derived, False otherwise
        """
        # First check if already known
        for known_fact in self.facts:
            if known_fact.name == fact.name:
                match = True
                for key, value in fact.attributes.items():
                    if key not in known_fact.attributes or known_fact.attributes[key] != value:
                        match = False
                        break
                if match:
                    return True
        return False
    
    def reset(self) -> None:
        """Reset all rules to unfired state."""
        for rule in self.rules:
            rule.reset()
    
    def __repr__(self) -> str:
        return f"KnowledgeBase({len(self.facts)} facts, {len(self.rules)} rules)"


def main() -> None:
    """Demo of the inference engine with a simple medical diagnosis system."""
    # Create knowledge base
    kb = KnowledgeBase()
    
    # Add initial facts
    kb.add_fact(Fact("patient", fever=True))
    kb.add_fact(Fact("patient", headache=True))
    kb.add_fact(Fact("patient", age="adult"))
    
    # Add rules
    kb.add_rule(Rule(
        "flu_diagnosis",
        [Fact("patient", fever=True), Fact("patient", headache=True)],
        [Fact("diagnosis", disease="flu")],
        priority=10
    ))
    
    kb.add_rule(Rule(
        "migraine_diagnosis",
        [Fact("patient", headache=True), Fact("patient", age="adult")],
        [Fact("diagnosis", disease="migraine")],
        priority=5
    ))
    
    kb.add_rule(Rule(
        "fever_treatment",
        [Fact("diagnosis", disease="flu")],
        [Fact("treatment", medication="antipyretic")],
        priority=15
    ))
    
    kb.add_rule(Rule(
        "headache_treatment",
        [Fact("diagnosis", disease="migraine")],
        [Fact("treatment", medication="analgesic")],
        priority=15
    ))
    
    print("Initial state:")
    print(f"Facts: {len(kb.facts)}")
    print(f"Rules: {len(kb.rules)}")
    print()
    
    # Perform inference
    iterations, facts_added = kb.forward_chain()
    
    print(f"Forward chaining completed in {iterations} iterations")
    print(f"Added {facts_added} new facts")
    print()
    
    # Check results
    print("Derived facts:")
    for fact in kb.facts:
        print(f"  {fact}")
    print()
    
    # Query specific facts
    print("Queries:")
    flu_query = kb.query(Fact("diagnosis", disease="flu"))
    migraine_query = kb.query(Fact("diagnosis", disease="migraine"))
    treatment_query = kb.query(Fact("treatment", medication="antipyretic"))
    
    print(f"Flu diagnosis: {flu_query}")
    print(f"Migraine diagnosis: {migraine_query}")
    print(f"Antipyretic treatment: {treatment_query}")
    
    # Test conflict resolution with another example
    print("\n" + "="*50)
    print("Testing conflict resolution:")
    
    kb2 = KnowledgeBase()
    
    # Facts
    kb2.add_fact(Fact("animal", has_fur=True))
    kb2.add_fact(Fact("animal", says="meow"))
    kb2.add_fact(Fact("animal", has_tail=True))
    
    # Rules with different priorities
    kb2.add_rule(Rule(
        "cat_id",
        [Fact("animal", has_fur=True), Fact("animal", says="meow")],
        [Fact("classification", species="cat")],
        priority=5
    ))
    
    kb2.add_rule(Rule(
        "mammal_id",
        [Fact("animal", has_fur=True)],
        [Fact("classification", species="mammal")],
        priority=1  # Lower priority
    ))
    
    # This rule should fire first due to higher priority
    kb2.add_rule(Rule(
        "pet_id",
        [Fact("animal", says="meow"), Fact("animal", has_tail=True)],
        [Fact("classification", category="pet")],
        priority=10
    ))
    
    iterations2, facts_added2 = kb2.forward_chain()
    print(f"Completed in {iterations2} iterations, added {facts_added2} facts")
    
    print("Final facts:")
    for fact in kb2.facts:
        print(f"  {fact}")


if __name__ == "__main__":
    main()