"""TLContext recipe descriptors (ADR-0065 D5).

A *recipe* is a named macro op (e.g. ``softmax_merge``) that TLContext
expands, at composite-emit time, into a sequence of flat ``OpSpec``s with
all addresses / sizes filled. The recipe table lives here, beside
``tl_context.py`` (the compiler analog). **PE_SCHEDULER never imports this
module** — it sees only the flat ops list (ADR-0065 D5 / D7 boundary).

Slot names used in ``EngineOp`` are either recipe operand names
(``RecipeDescriptor.operands``) or internal scratch slots produced by an
earlier op in ``engine_seq``. The scratch slot named ``"P"`` is, by
convention, the recipe's *primary output* — auto-bound into the head GEMM's
first matrix operand by the lowering pass.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Literal

PRIMARY_OUT_SLOT = "P"


@dataclass(frozen=True)
class PrimaryOutSpec:
    """Spec for a recipe's implicit primary output (slot ``"P"``)."""

    from_shape: str          # operand name to copy shape from
    from_dtype: str          # operand name to copy dtype from
    transform: str           # "identity" | "trans" (forward-compat)


@dataclass(frozen=True)
class EngineOp:
    """One micro-op in a recipe's ``engine_seq``.

    Slot fields (``src``, ``src_a``, ``src_b``, ``src_c``, ``dst``) name
    either ``RecipeDescriptor.operands`` or internal scratch slots.
    """

    engine: Literal["MATH", "GEMM", "DMA"]
    op_kind: str             # PE_MATH op_kind label (opaque to the engine)
    src: str | None = None
    src_a: str | None = None
    src_b: str | None = None
    src_c: str | None = None
    dst: str | None = None
    reduce_axis: int | None = None
    bcast_axis: int | None = None


@dataclass(frozen=True)
class RecipeDescriptor:
    operands: dict[str, str]                            # name → "R" | "RW"
    primary_out: PrimaryOutSpec | None                  # implicit output (slot "P")
    tile_alignment: Literal["single_shot", "tile_aligned"]
    internal_scratch_bytes_fn: Callable[..., int]
    engine_seq: tuple[EngineOp, ...]


RECIPE_DESCRIPTORS: dict[str, RecipeDescriptor] = {
    # Online-softmax merge (flash-attention accumulator update). Reads the
    # current score tile ``s`` and the running ``(m, l, O)``; emits the
    # rescaled probabilities ``P`` (slot "P") for the downstream P·V GEMM.
    "softmax_merge": RecipeDescriptor(
        operands={"s": "R", "m": "RW", "l": "RW", "O": "RW"},
        primary_out=PrimaryOutSpec(
            from_shape="s", from_dtype="s", transform="identity",
        ),
        tile_alignment="single_shot",
        internal_scratch_bytes_fn=lambda G, TILE, d, bpe: bpe * (
            G + G + G + G * TILE + G  # m_loc + m_new + corr + P + l_loc
        ),
        engine_seq=(
            EngineOp("MATH", "rmax",      src="s",                   dst="m_loc", reduce_axis=-1),
            EngineOp("MATH", "max_elem",  src_a="m",    src_b="m_loc", dst="m_new"),
            EngineOp("MATH", "exp_diff",  src_a="m",    src_b="m_new", dst="corr"),
            EngineOp("MATH", "exp_diff",  src_a="s",    src_b="m_new", dst="P", bcast_axis=0),
            EngineOp("MATH", "rsum",      src="P",                   dst="l_loc", reduce_axis=-1),
            EngineOp("MATH", "fma",       src_a="l",    src_b="corr", src_c="l_loc", dst="l"),
            EngineOp("MATH", "mul_bcast", src_a="O",    src_b="corr", dst="O", bcast_axis=1),
            EngineOp("MATH", "copy",      src="m_new",               dst="m"),
        ),
    ),
}
