"""Structural PE_CPU dispatch cost model (ADR-0064 Revision 2, with D8
single-op-cmd fast-path amendment).

Replaces the Rev1 per-op-type calibration table (``cpu_issue_cost.py``,
removed) with a structural formula derived from each command's
``logical_bytes`` (ADR-0064 D2):

    dispatch_cycles(cmd) = FIXED(cmd_type) + cmd.logical_bytes * R

The per-command ``FIXED`` term models the queue-tail update / MMIO-class
RTT / completion-event registration; the byte term ``R`` models queue-write
bandwidth. The primary signal the model exposes is **command-count
reduction** (FIXED-dominated); the byte term is a secondary refinement
(ADR-0064 Context).

ADR-0064 D8 (single-op-cmd fast-path): every single-op command — the DMA
descriptors (``DmaReadCmd`` / ``DmaWriteCmd``) and the single-op compute
dispatches (``GemmCmd`` / ``MathCmd`` / ``CopyCmd``) — carries a
separate, lighter FIXED than the generic control-path cost. The default
(``fixed_per_single_op_cmd_cycles = 8``) reflects descriptor-ring-push /
single-instruction-dispatch patterns common in modern accelerators
(NVIDIA TMA ~1 ISA cycle, AMD AQL packet write ~5–15 cycles, generic
Synopsys/Xilinx DMA IP ~5–20 cycles). The general control-path cost
(``fixed_per_cmd_cycles = 40``) applies only to ``CompositeCmd`` — the
one command that needs scheduler-side plan generation, per-tile RW-hazard
tracking, and completion-handle wiring.

Knobs are cycle-domain only; cycle→ns uses the PE node's ``clock_freq_ghz``
(ADR-0064 D3). Defaults anchor a typical 1-OpSpec GEMM composite (≈54 bytes)
at ≈43 ns on a 16 B/cycle on-die descriptor queue, while a 64 B single-op
descriptor lands at ≈12 ns.

A composite whose ``logical_bytes`` exceeds ``max_composite_logical_bytes``
is rejected at emit time with a ``ValueError`` (ADR-0064 D7, revised: hard
cap, no auto-segmentation).
"""
from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class PeCostModel:
    """Cycle-domain dispatch cost knobs (ADR-0064 D3/D4/D8)."""

    fixed_per_cmd_cycles: float = 40.0
    byte_cycles_recip: float = 0.0625          # = 16 bytes/cycle
    max_composite_logical_bytes: int = 1024     # D7 hard cap
    fixed_per_single_op_cmd_cycles: float = 8.0  # D8 single-op-cmd fast-path

    def dispatch_cycles(self, cmd_type: type, logical_bytes: int) -> float:
        """FIXED(cmd_type) + logical_bytes × R (ADR-0064 D1/D8).

        ``CompositeCmd`` uses the general control-path
        ``fixed_per_cmd_cycles``; every single-op command (DmaRead/
        DmaWrite/Gemm/Math/Copy) uses the lighter
        ``fixed_per_single_op_cmd_cycles``.
        """
        # Local import avoids a hard module-load cycle between common.pe_cost_model
        # and common.pe_commands; only CompositeCmd is needed to discriminate
        # the control path from the single-op fast path.
        from kernbench.common.pe_commands import CompositeCmd
        if cmd_type is CompositeCmd:
            fixed = self.fixed_per_cmd_cycles
        else:
            fixed = self.fixed_per_single_op_cmd_cycles
        return fixed + logical_bytes * self.byte_cycles_recip


DEFAULT_PE_COST_MODEL = PeCostModel()


def from_node_attrs(attrs: dict) -> PeCostModel:
    """Build a PeCostModel from a PE node's ``pe_cost_model:`` attrs block
    (ADR-0064 D4). Missing keys fall back to defaults."""
    block = attrs.get("pe_cost_model", {}) or {}
    d = DEFAULT_PE_COST_MODEL
    return PeCostModel(
        fixed_per_cmd_cycles=float(
            block.get("fixed_per_cmd_cycles", d.fixed_per_cmd_cycles)),
        byte_cycles_recip=float(
            block.get("byte_cycles_recip", d.byte_cycles_recip)),
        max_composite_logical_bytes=int(
            block.get("max_composite_logical_bytes",
                      d.max_composite_logical_bytes)),
        fixed_per_single_op_cmd_cycles=float(
            block.get("fixed_per_single_op_cmd_cycles",
                      d.fixed_per_single_op_cmd_cycles)),
    )
