"""Composable rule systems for chart parsing.

A ``RuleSystem`` encapsulates the structural inference rules that govern
a grammar -- binary combination rules and unary type-change rules -- as
index tensors suitable for CKY chart parsing. Rule systems are the
compositional unit: they can be merged (``+``), inspected, and passed
to any ``ChartParser``.

The primitive rule schemas live in ``quivers.stochastic.schema``,
which provides composable functors ``CategorySystem -> RuleSystem``.
The convenience functions here (``ccg_rules``, ``lambek_rules``)
instantiate schema presets over a given category system.

Categorical perspective
-----------------------
A rule system is a presentation of the arrows in a free category
generated by a set of basic combinators. Merging two rule systems
takes the coproduct of their generating sets.
"""

from __future__ import annotations

from quivers.stochastic._rule_system import RuleSystem
from quivers.stochastic.schema import (
    EVALUATION,
    HARMONIC_COMPOSITION,
    CROSSED_COMPOSITION,
    ADJUNCTION_UNITS,
    COMMUTATIVE_EVALUATION,
    TENSOR_INTRODUCTION,
    TENSOR_PROJECTION,
    GeneralizedComposition,
)
from quivers.stochastic.categories import CategorySystem

# re-export RuleSystem so existing imports keep working
__all__ = ["RuleSystem"]


# ================================================================
# CCG: a specific combination of biclosed monoidal primitives
# ================================================================


def ccg_rules(
    system: CategorySystem,
    *,
    enable_composition: bool = True,
    enable_crossed_composition: bool = True,
    enable_type_raising: bool = False,
    generalized_composition_depth: int = 1,
) -> RuleSystem:
    """Build a CCG rule system from a category inventory.

    Parameters
    ----------
    system : CategorySystem
        The finite category inventory.
    enable_composition : bool
        Enable harmonic composition (>B, <B).
    enable_crossed_composition : bool
        Enable crossed composition (>Bx, <Bx).
    enable_type_raising : bool
        Enable type raising (>T, <T).
    generalized_composition_depth : int
        Maximum depth for generalized composition (B^n).

    Returns
    -------
    RuleSystem
        The CCG rule system.
    """
    schema = EVALUATION

    if enable_composition:
        schema = schema | HARMONIC_COMPOSITION

    if enable_crossed_composition:
        schema = schema | CROSSED_COMPOSITION

    if enable_type_raising:
        schema = schema | ADJUNCTION_UNITS

    if generalized_composition_depth > 1 and enable_composition:
        schema = schema | GeneralizedComposition(
            max_depth=generalized_composition_depth,
        )

    rs = schema(system)

    return RuleSystem(
        binary_rules=rs.binary_rules,
        unary_rules=rs.unary_rules,
        n_categories=rs.n_categories,
        description="CCG",
    )


# ================================================================
# Lambek calculus: another combination of biclosed primitives
# ================================================================


def lambek_rules(
    system: CategorySystem,
    *,
    associative: bool = True,
    commutative: bool = False,
    enable_lifting: bool = True,
    enable_product: bool = True,
) -> RuleSystem:
    """Build a Lambek calculus rule system from a category inventory.

    Parameters
    ----------
    system : CategorySystem
        The finite category inventory.
    associative : bool
        Use associative Lambek calculus L (default True).
    commutative : bool
        Use commutative product / LP calculus (default False).
    enable_lifting : bool
        Enable Lambek lifting rules (adjunction units).
    enable_product : bool
        Enable product introduction / projection.

    Returns
    -------
    RuleSystem
        The Lambek calculus rule system.
    """
    schema = EVALUATION

    if commutative:
        schema = schema | COMMUTATIVE_EVALUATION

    if enable_lifting:
        schema = schema | ADJUNCTION_UNITS

    if enable_product:
        schema = schema | TENSOR_INTRODUCTION | TENSOR_PROJECTION

    rs = schema(system)

    variant = "L"

    if not associative:
        variant = "NL"

    if commutative:
        variant = "LP"

    return RuleSystem(
        binary_rules=rs.binary_rules,
        unary_rules=rs.unary_rules,
        n_categories=rs.n_categories,
        description=f"Lambek({variant})",
    )


# ================================================================
# Custom / manual rules
# ================================================================


def custom_rules(
    binary: list[tuple[int, int, int]] | None = None,
    unary: list[tuple[int, int]] | None = None,
    n_categories: int = 0,
    description: str = "custom",
) -> RuleSystem:
    """Build a rule system from explicit rule triples.

    Parameters
    ----------
    binary : list of (int, int, int) or None
        Binary rules as (result, left, right) triples.
    unary : list of (int, int) or None
        Unary rules as (result, input) pairs.
    n_categories : int
        Total number of categories.
    description : str
        Label for the rule system.

    Returns
    -------
    RuleSystem
        The custom rule system.
    """
    return RuleSystem(
        binary_rules=tuple(binary or []),
        unary_rules=tuple(unary or []),
        n_categories=n_categories,
        description=description,
    )
