"""Rule schemas: composable functors from CategorySystem to RuleSystem.

A ``RuleSchema`` is a functor that maps a ``CategorySystem`` to a
``RuleSystem``. Schemas are the abstract building blocks of grammar
formalisms: each schema defines a structural rule pattern, and
concrete rule instances are generated by instantiating the pattern
over a finite category inventory.

The two atomic schema types are:

- ``BinaryRuleSchema`` — defined by a match predicate on category
  pairs: ``match(left, right) -> result | None``.
- ``UnaryRuleSchema`` — defined by a match predicate on single
  categories: ``match(cat, system) -> list[result]``.

Schemas compose via ``|`` (union) and can carry default weights
via ``.weighted(w)``. Grammar formalisms are specific compositions::

    CCG = EVALUATION | HARMONIC_COMPOSITION | CROSSED_COMPOSITION
    LAMBEK = EVALUATION | ADJUNCTION_UNITS | TENSOR_INTRODUCTION | TENSOR_PROJECTION

The 13 bundled schemas (EVALUATION, HARMONIC_COMPOSITION, etc.)
are themselves composed from ~20 atomic schemas (ForwardApplication,
BackwardApplication, etc.), giving fine-grained control::

    # only forward application, no backward
    MY_GRAMMAR = ForwardApplication() | ForwardComposition()
"""

from __future__ import annotations

from abc import ABC, abstractmethod

from quivers.stochastic.categories import (
    AtomicCategory,
    Category,
    CategorySystem,
    ModalCategory,
    ProductCategory,
    SlashCategory,
    UnitCategory,
)
from quivers.stochastic._rule_system import RuleSystem


# ================================================================
# abstract base
# ================================================================


class RuleSchema(ABC):
    """A functor CategorySystem -> RuleSystem.

    Schemas are the abstract specification of structural rules.
    They compose via ``|`` (union) and ``+`` (alias for ``|``).
    """

    @abstractmethod
    def __call__(self, system: CategorySystem) -> RuleSystem:
        """Instantiate this schema over a category system.

        Parameters
        ----------
        system : CategorySystem
            The finite category inventory.

        Returns
        -------
        RuleSystem
            Concrete rules for this schema.
        """
        ...

    def __or__(self, other: RuleSchema) -> RuleSchema:
        """Union of two schemas."""
        return UnionSchema(self, other)

    def __add__(self, other: RuleSchema) -> RuleSchema:
        """Alias for union."""
        return self | other

    def weighted(self, weight: float) -> RuleSchema:
        """Attach a default weight to all rules from this schema."""
        return WeightedSchema(self, weight)


class UnionSchema(RuleSchema):
    """Union of two schemas (composes their rule systems via +)."""

    def __init__(self, left: RuleSchema, right: RuleSchema) -> None:
        self._left = left
        self._right = right

    def __call__(self, system: CategorySystem) -> RuleSystem:
        return self._left(system) + self._right(system)

    def __repr__(self) -> str:
        return f"({self._left!r} | {self._right!r})"


class WeightedSchema(RuleSchema):
    """Schema wrapper that attaches a default weight to all rules."""

    def __init__(self, inner: RuleSchema, weight: float) -> None:
        self._inner = inner
        self._weight = weight

    def __call__(self, system: CategorySystem) -> RuleSystem:
        rs = self._inner(system)
        binary_w = tuple(self._weight for _ in rs.binary_rules) or None
        unary_w = tuple(self._weight for _ in rs.unary_rules) or None

        return RuleSystem(
            binary_rules=rs.binary_rules,
            unary_rules=rs.unary_rules,
            n_categories=rs.n_categories,
            description=rs.description,
            binary_weights=binary_w,
            unary_weights=unary_w,
        )


# ================================================================
# atomic schema bases
# ================================================================


class BinaryRuleSchema(RuleSchema):
    """Schema for binary rules, defined by a match predicate.

    Subclasses implement ``match(left, right) -> result | None``.
    The schema iterates over all category pairs and collects matches.
    """

    @abstractmethod
    def match(
        self,
        left: Category,
        right: Category,
    ) -> Category | None:
        """Test whether (left, right) triggers this rule.

        Parameters
        ----------
        left : Category
            Left antecedent.
        right : Category
            Right antecedent.

        Returns
        -------
        Category or None
            Result category if the pair matches, else None.
        """
        ...

    @property
    def _description(self) -> str:
        return type(self).__name__

    def __call__(self, system: CategorySystem) -> RuleSystem:
        rules: list[tuple[int, int, int]] = []
        cats = list(system)

        for i, left in enumerate(cats):
            for j, right in enumerate(cats):
                result = self.match(left, right)

                if result is not None and result in system:
                    rules.append((system.index(result), i, j))

        return RuleSystem(
            binary_rules=tuple(rules),
            n_categories=system.size,
            description=self._description,
        )

    def __repr__(self) -> str:
        return self._description


class UnaryRuleSchema(RuleSchema):
    """Schema for unary rules, defined by a match predicate.

    Subclasses implement ``match(cat, system) -> list[result]``.
    The schema iterates over all categories and collects matches.
    """

    @abstractmethod
    def match(
        self,
        cat: Category,
        system: CategorySystem,
    ) -> list[Category]:
        """Find all valid results for a given input category.

        Parameters
        ----------
        cat : Category
            Input category.
        system : CategorySystem
            The full category inventory (needed for some schemas
            that must check existence of intermediate categories).

        Returns
        -------
        list[Category]
            Valid result categories.
        """
        ...

    @property
    def _description(self) -> str:
        return type(self).__name__

    def __call__(self, system: CategorySystem) -> RuleSystem:
        rules: list[tuple[int, int]] = []
        cats = list(system)

        for i, cat in enumerate(cats):
            for result in self.match(cat, system):
                if result in system:
                    rules.append((system.index(result), i))

        return RuleSystem(
            unary_rules=tuple(rules),
            n_categories=system.size,
            description=self._description,
        )

    def __repr__(self) -> str:
        return self._description


# ================================================================
# atomic binary schemas
# ================================================================


class ForwardApplication(BinaryRuleSchema):
    """C/B ⊗ B → C (right evaluation / forward application)."""

    @property
    def _description(self) -> str:
        return ">"

    def match(self, left, right):
        if (
            isinstance(left, SlashCategory)
            and left.direction == "/"
            and left.argument == right
        ):
            return left.result

        return None


class BackwardApplication(BinaryRuleSchema):
    """A ⊗ A\\C → C (left evaluation / backward application)."""

    @property
    def _description(self) -> str:
        return "<"

    def match(self, left, right):
        if (
            isinstance(right, SlashCategory)
            and right.direction == "\\"
            and right.argument == left
        ):
            return right.result

        return None


class ForwardComposition(BinaryRuleSchema):
    """C/B ⊗ B/A → C/A (right harmonic composition)."""

    @property
    def _description(self) -> str:
        return ">B"

    def match(self, left, right):
        if (
            isinstance(left, SlashCategory)
            and left.direction == "/"
            and isinstance(right, SlashCategory)
            and right.direction == "/"
            and left.argument == right.result
        ):
            return SlashCategory(left.result, right.argument, "/")

        return None


class BackwardComposition(BinaryRuleSchema):
    """A\\B ⊗ B\\C → A\\C (left harmonic composition)."""

    @property
    def _description(self) -> str:
        return "<B"

    def match(self, left, right):
        if (
            isinstance(left, SlashCategory)
            and left.direction == "\\"
            and isinstance(right, SlashCategory)
            and right.direction == "\\"
            and right.argument == left.result
        ):
            return SlashCategory(right.result, left.argument, "\\")

        return None


class ForwardCrossedComposition(BinaryRuleSchema):
    """C/B ⊗ B\\A → C\\A (right-left crossed composition)."""

    @property
    def _description(self) -> str:
        return ">Bx"

    def match(self, left, right):
        if (
            isinstance(left, SlashCategory)
            and left.direction == "/"
            and isinstance(right, SlashCategory)
            and right.direction == "\\"
            and left.argument == right.result
        ):
            return SlashCategory(left.result, right.argument, "\\")

        return None


class BackwardCrossedComposition(BinaryRuleSchema):
    """A/B ⊗ B\\C → A/C (left-right crossed composition)."""

    @property
    def _description(self) -> str:
        return "<Bx"

    def match(self, left, right):
        if (
            isinstance(left, SlashCategory)
            and left.direction == "/"
            and isinstance(right, SlashCategory)
            and right.direction == "\\"
            and right.argument == left.result
        ):
            return SlashCategory(right.result, left.argument, "/")

        return None


class CommutativeForwardApplication(BinaryRuleSchema):
    """B ⊗ C/B → C (commutative right evaluation)."""

    @property
    def _description(self) -> str:
        return "comm>"

    def match(self, left, right):
        if (
            isinstance(right, SlashCategory)
            and right.direction == "/"
            and right.argument == left
        ):
            return right.result

        return None


class CommutativeBackwardApplication(BinaryRuleSchema):
    """A\\C ⊗ A → C (commutative left evaluation)."""

    @property
    def _description(self) -> str:
        return "comm<"

    def match(self, left, right):
        if (
            isinstance(left, SlashCategory)
            and left.direction == "\\"
            and left.argument == right
        ):
            return left.result

        return None


class TensorIntroduction(BinaryRuleSchema):
    """A ⊗ B → A⊗B (product formation)."""

    @property
    def _description(self) -> str:
        return "⊗-intro"

    def match(self, left, right):
        return ProductCategory(left, right)


class LeftUnitElimination(BinaryRuleSchema):
    """I ⊗ A → A (left unitor)."""

    @property
    def _description(self) -> str:
        return "λ"

    def match(self, left, right):
        if isinstance(left, UnitCategory):
            return right

        return None


class RightUnitElimination(BinaryRuleSchema):
    """A ⊗ I → A (right unitor)."""

    @property
    def _description(self) -> str:
        return "ρ"

    def match(self, left, right):
        if isinstance(right, UnitCategory):
            return left

        return None


class ModalApplication(BinaryRuleSchema):
    """◇(C/B) ⊗ ◇B → ◇C (modal function application).

    Parameters
    ----------
    modality : str
        The modality symbol (e.g. "◇", "□").
    """

    def __init__(self, modality: str = "◇") -> None:
        self._modality = modality

    @property
    def _description(self) -> str:
        return f"{self._modality}-app"

    def match(self, left, right):
        if (
            not isinstance(left, ModalCategory)
            or left.modality != self._modality
            or not isinstance(right, ModalCategory)
            or right.modality != self._modality
        ):
            return None

        inner_left = left.inner
        inner_right = right.inner

        # ◇(C/B) ⊗ ◇B → ◇C
        if (
            isinstance(inner_left, SlashCategory)
            and inner_left.direction == "/"
            and inner_left.argument == inner_right
        ):
            return ModalCategory(self._modality, inner_left.result)

        # ◇A ⊗ ◇(A\C) → ◇C
        if (
            isinstance(inner_right, SlashCategory)
            and inner_right.direction == "\\"
            and inner_right.argument == inner_left
        ):
            return ModalCategory(self._modality, inner_right.result)

        return None


# ================================================================
# atomic unary schemas
# ================================================================


class RightLifting(UnaryRuleSchema):
    """A → C/(A\\C) (right adjunction unit / type raising)."""

    @property
    def _description(self) -> str:
        return "η-right"

    def match(self, cat, system):
        results = []

        for c in system:
            inner = SlashCategory(c, cat, "\\")

            if inner in system:
                raised = SlashCategory(c, inner, "/")

                if raised in system:
                    results.append(raised)

        return results


class LeftLifting(UnaryRuleSchema):
    """A → (C/A)\\C (left adjunction unit / type raising)."""

    @property
    def _description(self) -> str:
        return "η-left"

    def match(self, cat, system):
        results = []

        for c in system:
            inner = SlashCategory(c, cat, "/")

            if inner in system:
                raised = SlashCategory(c, inner, "\\")

                if raised in system:
                    results.append(raised)

        return results


class LeftProjection(UnaryRuleSchema):
    """A⊗B → A (left tensor projection)."""

    @property
    def _description(self) -> str:
        return "π₁"

    def match(self, cat, system):
        if isinstance(cat, ProductCategory) and cat.left in system:
            return [cat.left]

        return []


class RightProjection(UnaryRuleSchema):
    """A⊗B → B (right tensor projection)."""

    @property
    def _description(self) -> str:
        return "π₂"

    def match(self, cat, system):
        if isinstance(cat, ProductCategory) and cat.right in system:
            return [cat.right]

        return []


class UnitCoercion(UnaryRuleSchema):
    """A → I (unit elimination / weakening)."""

    @property
    def _description(self) -> str:
        return "unit-elim"

    def match(self, cat: Category, system: CategorySystem) -> list[Category]:
        unit = UnitCategory()

        if unit in system:
            return [unit]

        return []


class ModalInjection(UnaryRuleSchema):
    """A → ◇A (modal introduction).

    Parameters
    ----------
    modality : str
        The modality symbol (e.g. "◇", "□").
    """

    def __init__(self, modality: str = "◇") -> None:
        self._modality = modality

    @property
    def _description(self) -> str:
        return f"{self._modality}-intro"

    def match(self, cat: Category, system: CategorySystem) -> list[Category]:
        modal = ModalCategory(self._modality, cat)

        if modal in system:
            return [modal]

        return []


class ModalProjection(UnaryRuleSchema):
    """◇A → A (modal elimination).

    Parameters
    ----------
    modality : str
        The modality symbol (e.g. "◇", "□").
    """

    def __init__(self, modality: str = "◇") -> None:
        self._modality = modality

    @property
    def _description(self) -> str:
        return f"{self._modality}-elim"

    def match(self, cat, system):
        if (
            isinstance(cat, ModalCategory)
            and cat.modality == self._modality
            and cat.inner in system
        ):
            return [cat.inner]

        return []


# ================================================================
# generalized composition (special: recursive match)
# ================================================================


class GeneralizedComposition(RuleSchema):
    """Generalized composition through nested slashes (B^n).

    C/B ⊗ (B|₁Z₁)|₂Z₂...|ₙZₙ → (C|₁Z₁)|₂Z₂...|ₙZₙ

    This is a special schema because the match recurses into
    nested slash structure. At depth 1 it is ordinary harmonic
    composition; this schema only generates rules for depth >= 2.

    Parameters
    ----------
    max_depth : int
        Maximum nesting depth to traverse.
    """

    def __init__(self, max_depth: int = 2) -> None:
        self._max_depth = max_depth

    def __call__(self, system: CategorySystem) -> RuleSystem:
        rules: list[tuple[int, int, int]] = []
        cats = list(system)

        for i, left in enumerate(cats):
            if not isinstance(left, SlashCategory) or left.direction != "/":
                continue

            for j, right in enumerate(cats):
                self._collect(system, rules, left, right, i, j, depth=2)

        return RuleSystem(
            binary_rules=tuple(rules),
            n_categories=system.size,
            description=f"gen-comp(depth={self._max_depth})",
        )

    def _collect(self, system, rules, left, right, li, ri, depth):
        if depth > self._max_depth:
            return

        if not isinstance(right, SlashCategory):
            return

        if right.result == left.argument:
            result = SlashCategory(left.result, right.argument, right.direction)

            if result in system:
                rules.append((system.index(result), li, ri))

        if isinstance(right.result, SlashCategory):
            self._collect(system, rules, left, right.result, li, ri, depth + 1)

    def __repr__(self) -> str:
        return f"GeneralizedComposition(max_depth={self._max_depth})"


# ================================================================
# pattern-based schemas (from DSL rule declarations)
# ================================================================


def _match_pattern(
    pattern,
    category: Category,
    bindings: dict[str, Category],
    variables: frozenset[str],
) -> bool:
    """Match a category pattern against a concrete category.

    Attempts to unify ``pattern`` with ``category``, extending
    ``bindings`` with variable assignments. Returns True on
    success, False on failure. Mutates ``bindings`` in place.

    Parameters
    ----------
    pattern
        A category pattern AST node (CatPatternName, CatPatternSlash,
        or CatPatternProduct).
    category : Category
        The concrete category to match against.
    bindings : dict
        Current variable bindings (mutated on success).
    variables : frozenset
        The set of variable names (from the rule's parameter list).
    """
    from quivers.dsl.ast_nodes import CatPatternName, CatPatternSlash, CatPatternProduct

    if isinstance(pattern, CatPatternName):
        name = pattern.name

        if name in variables:
            # pattern variable: bind or check consistency
            if name in bindings:
                return bindings[name] == category

            bindings[name] = category
            return True

        else:
            # category atom constant: must match exactly
            return isinstance(category, AtomicCategory) and category.name == name

    elif isinstance(pattern, CatPatternSlash):
        if not isinstance(category, SlashCategory):
            return False

        if pattern.direction != category.direction:
            return False

        return _match_pattern(
            pattern.result, category.result, bindings, variables
        ) and _match_pattern(pattern.argument, category.argument, bindings, variables)

    elif isinstance(pattern, CatPatternProduct):
        if not isinstance(category, ProductCategory):
            return False

        return _match_pattern(
            pattern.left, category.left, bindings, variables
        ) and _match_pattern(pattern.right, category.right, bindings, variables)

    return False


def _instantiate_pattern(
    pattern,
    bindings: dict[str, Category],
    variables: frozenset[str],
) -> Category | None:
    """Instantiate a category pattern with bound variables.

    Parameters
    ----------
    pattern
        A category pattern AST node.
    bindings : dict
        Variable bindings from matching.
    variables : frozenset
        The set of variable names.

    Returns
    -------
    Category or None
        The concrete category, or None if a variable is unbound.
    """
    from quivers.dsl.ast_nodes import CatPatternName, CatPatternSlash, CatPatternProduct

    if isinstance(pattern, CatPatternName):
        if pattern.name in variables:
            return bindings.get(pattern.name)

        return AtomicCategory(pattern.name)

    elif isinstance(pattern, CatPatternSlash):
        result = _instantiate_pattern(pattern.result, bindings, variables)
        argument = _instantiate_pattern(pattern.argument, bindings, variables)

        if result is None or argument is None:
            return None

        return SlashCategory(result, argument, pattern.direction)

    elif isinstance(pattern, CatPatternProduct):
        left = _instantiate_pattern(pattern.left, bindings, variables)
        right = _instantiate_pattern(pattern.right, bindings, variables)

        if left is None or right is None:
            return None

        return ProductCategory(left, right)

    return None


class PatternBinarySchema(BinaryRuleSchema):
    """Binary rule schema from a DSL rule declaration.

    Matches category pairs against two premise patterns and
    instantiates the conclusion pattern with bound variables.

    Parameters
    ----------
    left_pattern
        Category pattern for the left antecedent.
    right_pattern
        Category pattern for the right antecedent.
    conclusion_pattern
        Category pattern for the consequent.
    variables : frozenset of str
        Universally quantified pattern variables.
    name : str
        The rule's name (for description).
    """

    def __init__(
        self,
        left_pattern,
        right_pattern,
        conclusion_pattern,
        variables: frozenset[str],
        name: str = "",
    ) -> None:
        self._left = left_pattern
        self._right = right_pattern
        self._conclusion = conclusion_pattern
        self._variables = variables
        self._name = name

    @property
    def _description(self) -> str:
        return self._name or "pattern-binary"

    def match(self, left: Category, right: Category) -> Category | None:
        bindings: dict[str, Category] = {}

        if not _match_pattern(self._left, left, bindings, self._variables):
            return None

        if not _match_pattern(self._right, right, bindings, self._variables):
            return None

        return _instantiate_pattern(
            self._conclusion,
            bindings,
            self._variables,
        )


class PatternUnarySchema(UnaryRuleSchema):
    """Unary rule schema from a DSL rule declaration.

    Matches a single category against a premise pattern and
    instantiates the conclusion pattern with bound variables.

    Parameters
    ----------
    premise_pattern
        Category pattern for the antecedent.
    conclusion_pattern
        Category pattern for the consequent.
    variables : frozenset of str
        Universally quantified pattern variables.
    name : str
        The rule's name (for description).
    """

    def __init__(
        self,
        premise_pattern,
        conclusion_pattern,
        variables: frozenset[str],
        name: str = "",
    ) -> None:
        self._premise = premise_pattern
        self._conclusion = conclusion_pattern
        self._variables = variables
        self._name = name

    @property
    def _description(self) -> str:
        return self._name or "pattern-unary"

    def match(
        self,
        cat: Category,
        system: CategorySystem,
    ) -> list[Category]:
        bindings: dict[str, Category] = {}

        if not _match_pattern(self._premise, cat, bindings, self._variables):
            return []

        result = _instantiate_pattern(
            self._conclusion,
            bindings,
            self._variables,
        )

        if result is None:
            return []

        return [result]


# ================================================================
# bundled schemas (the 13 traditional primitives)
# ================================================================

EVALUATION = ForwardApplication() | BackwardApplication()
HARMONIC_COMPOSITION = ForwardComposition() | BackwardComposition()
CROSSED_COMPOSITION = ForwardCrossedComposition() | BackwardCrossedComposition()
COMMUTATIVE_EVALUATION = (
    CommutativeForwardApplication() | CommutativeBackwardApplication()
)
ADJUNCTION_UNITS = RightLifting() | LeftLifting()
TENSOR_INTRODUCTION = TensorIntroduction()
TENSOR_PROJECTION = LeftProjection() | RightProjection()
UNIT_INTRODUCTION = LeftUnitElimination() | RightUnitElimination()
UNIT_ELIMINATION = UnitCoercion()
MODAL_INTRODUCTION = ModalInjection()
MODAL_ELIMINATION = ModalProjection()
MODAL_APPLICATION = ModalApplication()


def generalized_composition(max_depth: int = 2) -> RuleSchema:
    """Create a generalized composition schema.

    Parameters
    ----------
    max_depth : int
        Maximum nesting depth.

    Returns
    -------
    RuleSchema
        The generalized composition schema.
    """
    return GeneralizedComposition(max_depth)


# ================================================================
# grammar presets (specific compositions of bundled schemas)
# ================================================================

CCG = EVALUATION | HARMONIC_COMPOSITION | CROSSED_COMPOSITION
LAMBEK = EVALUATION | ADJUNCTION_UNITS | TENSOR_INTRODUCTION | TENSOR_PROJECTION
NL = EVALUATION
LP = EVALUATION | COMMUTATIVE_EVALUATION | ADJUNCTION_UNITS


# ================================================================
# schema registry for DSL name resolution
# ================================================================

SCHEMA_REGISTRY: dict[str, RuleSchema] = {
    "evaluation": EVALUATION,
    "harmonic_composition": HARMONIC_COMPOSITION,
    "crossed_composition": CROSSED_COMPOSITION,
    "generalized_composition": GeneralizedComposition(),
    "adjunction_units": ADJUNCTION_UNITS,
    "tensor_introduction": TENSOR_INTRODUCTION,
    "tensor_projection": TENSOR_PROJECTION,
    "commutative_evaluation": COMMUTATIVE_EVALUATION,
    "unit_introduction": UNIT_INTRODUCTION,
    "unit_elimination": UNIT_ELIMINATION,
    "modal_introduction": MODAL_INTRODUCTION,
    "modal_elimination": MODAL_ELIMINATION,
    "modal_application": MODAL_APPLICATION,
}
