"""
norvig.logic – AIMA Chapters 7, 8, 9
=======================================
Uses SymPy as the symbolic logic engine.

    from ml_ai_library.norvig.logic import (
        # SymPy primitives re-exported
        Symbol, symbols,
        And, Or, Not, Implies, Equivalent,
        # KB wrappers
        PropKB,        # propositional knowledge base with TELL / ASK
        satisfiable,   # sympy.satisfiable (DPLL-based SAT)
        entails,       # KB |= query check
        # FOL helpers
        unify,         # Robinson unification (AIMA Fig 9.1)
        fol_bc_ask,    # backward chaining over definite clauses
    )

AIMA chapter map:
  PropKB / satisfiable / entails  → Ch 7 (propositional logic)
  unify / fol_bc_ask              → Ch 9 (first-order logic)
"""

from __future__ import annotations
from typing import Any, Dict, List, Optional

# SymPy symbolic logic engine ---------------------------------------------------
from sympy import Symbol, symbols                              # atomic propositions
from sympy.logic.boolalg import (
    And, Or, Not, Implies, Equivalent,
    to_cnf,                                                   # CNF conversion
)
from sympy.logic.inference import satisfiable                 # DPLL-based SAT solver


# ---------------------------------------------------------------------------
# Propositional Knowledge Base (AIMA §7.4)
# ---------------------------------------------------------------------------

class PropKB:
    """
    Propositional KB backed by SymPy (AIMA §7.4).

    tell(sentence) – add a SymPy Boolean expression
    ask(query)     – return True if KB ⊨ query (via satisfiability refutation)

    Example
    -------
    >>> P, Q = symbols('P Q')
    >>> kb = PropKB()
    >>> kb.tell(Implies(P, Q))
    >>> kb.tell(P)
    >>> kb.ask(Q)
    True
    """

    def __init__(self):
        self.sentences: List = []

    def tell(self, sentence):
        """Add a SymPy Boolean expression to the KB."""
        self.sentences.append(sentence)

    def ask(self, query) -> bool:
        """
        KB ⊨ query  iff  KB ∧ ¬query is unsatisfiable.
        Uses sympy.satisfiable (DPLL-based).
        """
        if not self.sentences:
            return False
        kb_conj = And(*self.sentences)
        return not satisfiable(And(kb_conj, Not(query)))

    def retract(self, sentence):
        """Remove a sentence from the KB."""
        self.sentences = [s for s in self.sentences if s != sentence]

    def __repr__(self):
        return f"PropKB({self.sentences})"


def entails(kb_sentences: List, query) -> bool:
    """
    Check KB ⊨ query using SymPy's satisfiable (AIMA §7.5).
    kb_sentences: list of SymPy Boolean expressions.
    """
    if not kb_sentences:
        return False
    kb_conj = And(*kb_sentences)
    return not satisfiable(And(kb_conj, Not(query)))


# ---------------------------------------------------------------------------
# FOL Unification (AIMA Fig 9.1) – implemented with plain Python dicts
# ---------------------------------------------------------------------------

def unify(x: Any, y: Any, theta: Optional[Dict] = None) -> Optional[Dict]:
    """
    Unification algorithm (AIMA Fig 9.1).
    Variables are lowercase strings; constants / functors are everything else.

    Returns substitution dict or None if unification fails.

    Example
    -------
    >>> unify(('Knows', 'John', 'x'), ('Knows', 'John', 'Jane'))
    {'x': 'Jane'}
    """
    if theta is None:
        theta = {}
    if theta is None:
        return None
    if x == y:
        return theta
    if _is_variable(x):
        return _unify_var(x, y, theta)
    if _is_variable(y):
        return _unify_var(y, x, theta)
    if isinstance(x, tuple) and isinstance(y, tuple) and len(x) == len(y):
        theta = unify(x[0], y[0], theta)
        if theta is None:
            return None
        return unify(x[1:], y[1:], theta)
    return None


def _is_variable(x) -> bool:
    return isinstance(x, str) and len(x) > 0 and x[0].islower()


def _unify_var(var: str, x: Any, theta: Dict) -> Optional[Dict]:
    if var in theta:
        return unify(theta[var], x, theta)
    if _is_variable(x) and x in theta:
        return unify(var, theta[x], theta)
    if _occurs_in(var, x, theta):          # occurs check
        return None
    return {**theta, var: x}


def _occurs_in(var, x, theta) -> bool:
    if var == x:
        return True
    if _is_variable(x) and x in theta:
        return _occurs_in(var, theta[x], theta)
    if isinstance(x, tuple):
        return any(_occurs_in(var, elem, theta) for elem in x)
    return False


# ---------------------------------------------------------------------------
# FOL Backward Chaining (AIMA Fig 9.6)
# ---------------------------------------------------------------------------

def fol_bc_ask(kb: List[tuple], query: tuple):
    """
    Backward-chaining FOL inference over definite clauses (AIMA Fig 9.6).

    kb    : list of (head, body_list) pairs
            e.g. [( ('Mortal','x'), [('Human','x')] )]
    query : goal atom tuple, e.g. ('Mortal','Socrates')

    Yields substitutions (dicts) satisfying the query.

    Example
    -------
    >>> kb = [
    ...   (('Mortal', 'x'), [('Human', 'x')]),
    ...   (('Human', 'Socrates'), []),
    ... ]
    >>> list(fol_bc_ask(kb, ('Mortal', 'Socrates')))
    [{'x': 'Socrates'}]
    """
    yield from _bc_or(kb, query, {})


def _bc_or(kb, goal, theta):
    for head, body in kb:
        theta2 = unify(head, goal, dict(theta))
        if theta2 is not None:
            yield from _bc_and(kb, body, theta2)


def _bc_and(kb, goals, theta):
    if not goals:
        yield theta
    else:
        first, rest = goals[0], goals[1:]
        for theta2 in _bc_or(kb, _subst(first, theta), theta):
            yield from _bc_and(kb, rest, theta2)


def _subst(atom, theta: Dict) -> tuple:
    return tuple(theta.get(a, a) if not isinstance(a, tuple) else _subst(a, theta)
                 for a in atom)
