"""Category types for categorial grammars.

Provides the free algebra of grammatical categories generated by
atoms and type constructors. This is the type system underlying
chart parsers for CCG, the Lambek calculus, multimodal type-logical
grammar, and any bespoke formalism assembled from biclosed monoidal
primitives.

Atomic categories
-----------------
``AtomicCategory("S")`` — a basic category like S, NP, N, PP.

Slash categories (internal homs)
--------------------------------
``SlashCategory(result, arg, "/")``  — right residual: result/arg
``SlashCategory(result, arg, "\\")`` — left residual: result\\arg

Product categories (tensor product)
------------------------------------
``ProductCategory(left, right)`` — tensor: left ⊗ right

Unit category (monoidal unit)
-----------------------------
``UnitCategory()`` — the monoidal unit I satisfying I ⊗ A ≅ A ≅ A ⊗ I.

Modal categories (unary type constructors)
------------------------------------------
``ModalCategory("◇", inner)`` — a modality applied to a category.
Supports arbitrary named modalities (◇, □, or user-defined).

Category systems
----------------
A ``CategorySystem`` manages a finite inventory of categories
with integer indices, supporting lookup and enumeration. This is
necessary because chart parsers operate over a finite set of
categories.

The ``from_generators`` factory method enumerates the free algebra
generated by atoms and a set of type constructors up to a given
depth, enabling fully parameterizable grammar construction.

Examples
--------
>>> cs = CategorySystem.from_atoms(["S", "NP", "N", "VP", "PP"])
>>> cs.add_slash(cs["S"], cs["NP"], "\\\\")  # S\\NP
>>> cs.add_slash(cs["S\\\\NP"], cs["NP"], "/")  # (S\\NP)/NP
>>> print(cs.size)
7

>>> # with modal operators and unit
>>> cs = CategorySystem.from_generators(
...     atoms=["S", "NP"],
...     constructors=["slash", "product", "unit", "diamond", "box"],
...     max_depth=1,
... )
"""

from __future__ import annotations

from dataclasses import dataclass
from collections.abc import Callable, Iterator


@dataclass(frozen=True, eq=True)
class AtomicCategory:
    """An atomic (primitive) grammatical category.

    Parameters
    ----------
    name : str
        The category name (e.g., "S", "NP", "N").
    """

    name: str

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


@dataclass(frozen=True, eq=True)
class SlashCategory:
    """A complex category formed by a slash operator.

    ``result / argument`` (forward): seeks argument on the right.
    ``result \\ argument`` (backward): seeks argument on the left.

    These are the internal homs of a biclosed monoidal category:
    ``/`` is the right residual and ``\\`` is the left residual.

    Parameters
    ----------
    result : Category
        The result category after combination.
    argument : Category
        The argument category sought.
    direction : str
        Slash direction: ``"/"`` (forward) or ``"\\\\"`` (backward).
    """

    result: Category
    argument: Category
    direction: str  # "/" or "\\"

    def __post_init__(self) -> None:
        if self.direction not in ("/", "\\"):
            raise ValueError(f"direction must be '/' or '\\\\', got {self.direction!r}")

    def __repr__(self) -> str:
        res = _cat_str(self.result)
        arg = _cat_str(self.argument)

        # parenthesize complex sub-categories
        if isinstance(self.result, SlashCategory):
            res = f"({res})"

        if isinstance(self.argument, SlashCategory):
            arg = f"({arg})"

        return f"{res}{self.direction}{arg}"


@dataclass(frozen=True, eq=True)
class ProductCategory:
    """A product category (tensor product) for type-logical grammar.

    Represents A ⊗ B in the Lambek calculus.

    Parameters
    ----------
    left : Category
        Left component.
    right : Category
        Right component.
    """

    left: Category
    right: Category

    def __repr__(self) -> str:
        l = _cat_str(self.left)
        r = _cat_str(self.right)

        if isinstance(self.left, (SlashCategory, ProductCategory)):
            l = f"({l})"

        if isinstance(self.right, (SlashCategory, ProductCategory)):
            r = f"({r})"

        return f"{l}⊗{r}"


@dataclass(frozen=True, eq=True)
class UnitCategory:
    """The monoidal unit category I.

    Satisfies the unit laws: I ⊗ A ≅ A ≅ A ⊗ I. Required for
    unit introduction and elimination rules.
    """

    def __repr__(self) -> str:
        return "I"


@dataclass(frozen=True, eq=True)
class ModalCategory:
    """A modal category: a unary type constructor applied to a category.

    Represents modalities like ◇A (diamond) and □A (box) used in
    multimodal type-logical grammar (Moortgat 1997). Arbitrary
    named modalities are supported.

    Parameters
    ----------
    modality : str
        The modality name (e.g., "◇", "□", or a custom name).
    inner : Category
        The category under the modality.
    """

    modality: str
    inner: Category

    def __repr__(self) -> str:
        inner_str = _cat_str(self.inner)

        if isinstance(self.inner, (SlashCategory, ProductCategory)):
            inner_str = f"({inner_str})"

        return f"{self.modality}{inner_str}"


# union type for all category kinds
Category = (
    AtomicCategory | SlashCategory | ProductCategory | UnitCategory | ModalCategory
)

# all concrete category types for isinstance checks
_CATEGORY_TYPES = (
    AtomicCategory,
    SlashCategory,
    ProductCategory,
    UnitCategory,
    ModalCategory,
)


def _cat_str(cat: Category) -> str:
    """String representation of a category."""
    return repr(cat)


# ================================================================
# built-in constructor names for from_generators
# ================================================================

# each constructor is a callable: (system, current_cats) -> list[Category]
# they generate new categories from existing ones


def _slash_constructor(
    current: list[Category],
) -> list[Category]:
    """Generate all slash categories from current inventory."""
    new = []

    for result in current:
        for arg in current:
            new.append(SlashCategory(result, arg, "/"))
            new.append(SlashCategory(result, arg, "\\"))

    return new


def _product_constructor(
    current: list[Category],
) -> list[Category]:
    """Generate all product categories from current inventory."""
    new = []

    for left in current:
        for right in current:
            new.append(ProductCategory(left, right))

    return new


def _unit_constructor(
    current: list[Category],
) -> list[Category]:
    """Generate the unit category (only once)."""
    return [UnitCategory()]


def _modal_constructor(
    modality: str,
) -> Callable[[list[Category]], list[Category]]:
    """Create a modal constructor for a specific modality.

    Parameters
    ----------
    modality : str
        The modality name.

    Returns
    -------
    Callable[[list[Category]], list[Category]]
        A constructor function.
    """

    def constructor(current: list[Category]) -> list[Category]:
        return [ModalCategory(modality, c) for c in current]

    return constructor


# registry of built-in constructors
_BUILTIN_CONSTRUCTORS: dict[str, Callable[[list[Category]], list[Category]]] = {
    "slash": _slash_constructor,
    "product": _product_constructor,
    "unit": _unit_constructor,
    "diamond": _modal_constructor("◇"),
    "box": _modal_constructor("□"),
}


class CategorySystem:
    """A finite inventory of grammatical categories with integer indices.

    Manages the mapping between ``Category`` objects and their
    integer indices, which are used by the chart parsers.

    Parameters
    ----------
    categories : list of Category
        The initial category inventory.
    """

    def __init__(self, categories: list[Category] | None = None) -> None:
        self._cats: list[Category] = []
        self._index: dict[Category, int] = {}

        if categories:
            for cat in categories:
                self.add(cat)

    @classmethod
    def from_atoms(cls, names: list[str]) -> CategorySystem:
        """Create a category system from atomic category names.

        Parameters
        ----------
        names : list of str
            Names of atomic categories (e.g., ["S", "NP", "N"]).

        Returns
        -------
        CategorySystem
            A new system with the given atomic categories.
        """
        cats: list[Category] = [AtomicCategory(name) for name in names]
        return cls(cats)

    @classmethod
    def from_atoms_and_slash_depth(
        cls,
        names: list[str],
        max_depth: int = 1,
    ) -> CategorySystem:
        """Create a category system with atomic categories and all
        slash combinations up to a given depth.

        This auto-enumerates the category space, which is necessary
        for chart parsing over a finite category set.

        Parameters
        ----------
        names : list of str
            Atomic category names.
        max_depth : int
            Maximum nesting depth for slash categories.

        Returns
        -------
        CategorySystem
            A system with all categories up to the given depth.
        """
        system = cls.from_atoms(names)

        if max_depth < 1:
            return system

        # iteratively build slash categories up to max_depth
        for _ in range(max_depth):
            current = list(system._cats)

            for result in current:
                for arg in current:
                    system.add(SlashCategory(result, arg, "/"))
                    system.add(SlashCategory(result, arg, "\\"))

        return system

    @classmethod
    def from_generators(
        cls,
        atoms: list[str],
        constructors: list[str] | None = None,
        max_depth: int = 1,
    ) -> CategorySystem:
        """Create a category system from atoms and type constructors.

        Enumerates the free algebra generated by the given atoms and
        constructors up to ``max_depth`` layers of nesting. This is
        the fully general factory for building category systems for
        any grammar formalism.

        Built-in constructor names:

        - ``"slash"`` — generates all A/B and A\\B
        - ``"product"`` — generates all A⊗B
        - ``"unit"`` — adds the monoidal unit I
        - ``"diamond"`` — generates ◇A for all A
        - ``"box"`` — generates □A for all A

        Parameters
        ----------
        atoms : list of str
            Atomic category names.
        constructors : list of str or None
            Constructor names. If None, defaults to ``["slash"]``.
        max_depth : int
            Number of constructor application layers.

        Returns
        -------
        CategorySystem
            A system with all generated categories.

        Raises
        ------
        ValueError
            If an unknown constructor name is given.
        """
        if constructors is None:
            constructors = ["slash"]

        # resolve constructor functions
        fns = []

        for name in constructors:
            if name not in _BUILTIN_CONSTRUCTORS:
                raise ValueError(
                    f"unknown constructor {name!r}; available: "
                    f"{sorted(_BUILTIN_CONSTRUCTORS.keys())}"
                )

            fns.append(_BUILTIN_CONSTRUCTORS[name])

        system = cls.from_atoms(atoms)

        for _ in range(max_depth):
            current = list(system._cats)

            for fn in fns:
                for cat in fn(current):
                    system.add(cat)

        return system

    def add(self, cat: Category) -> int:
        """Add a category to the system.

        Parameters
        ----------
        cat : Category
            The category to add.

        Returns
        -------
        int
            The index of the category (existing or newly assigned).
        """
        if cat in self._index:
            return self._index[cat]

        idx = len(self._cats)
        self._cats.append(cat)
        self._index[cat] = idx
        return idx

    def add_slash(
        self,
        result: Category,
        argument: Category,
        direction: str,
    ) -> int:
        """Add a slash category and return its index.

        Parameters
        ----------
        result : Category
            The result category.
        argument : Category
            The argument category.
        direction : str
            ``"/"`` or ``"\\\\"``

        Returns
        -------
        int
            Index of the new or existing slash category.
        """
        cat = SlashCategory(result, argument, direction)
        return self.add(cat)

    @property
    def size(self) -> int:
        """Number of categories in the system."""
        return len(self._cats)

    def __len__(self) -> int:
        return len(self._cats)

    def __getitem__(self, key: str | int | Category) -> Category | int:
        """Look up by name (str), index (int), or category object.

        Parameters
        ----------
        key : str or int or Category
            If str: returns the AtomicCategory with that name.
            If int: returns the Category at that index.
            If Category: returns the integer index.

        Returns
        -------
        Category or int
            The looked-up value.

        Raises
        ------
        KeyError
            If the key is not found.
        """
        if isinstance(key, str):
            cat = AtomicCategory(key)

            if cat not in self._index:
                raise KeyError(f"category {key!r} not in system")

            return cat

        if isinstance(key, int):
            if key < 0 or key >= len(self._cats):
                raise KeyError(f"index {key} out of range")

            return self._cats[key]

        if isinstance(key, _CATEGORY_TYPES):
            if key not in self._index:
                raise KeyError(f"category {key!r} not in system")

            return self._index[key]

        raise TypeError(f"invalid key type: {type(key).__name__}")

    def index(self, cat: Category) -> int:
        """Get the integer index of a category.

        Parameters
        ----------
        cat : Category
            The category to look up.

        Returns
        -------
        int
            The index.

        Raises
        ------
        KeyError
            If the category is not in the system.
        """
        if cat not in self._index:
            raise KeyError(f"category {cat!r} not in system")

        return self._index[cat]

    def __contains__(self, cat: Category) -> bool:
        return cat in self._index

    def __iter__(self) -> Iterator[Category]:
        return iter(self._cats)

    def __repr__(self) -> str:
        cats_str = ", ".join(repr(c) for c in self._cats[:10])

        if len(self._cats) > 10:
            cats_str += f", ... ({len(self._cats)} total)"

        return f"CategorySystem([{cats_str}])"
