"""DataExecutor: Phase 2 op_log-based data execution (ADR-0020 D6).

Executes GEMM/Math operations from the op_log using numpy.
Memory ops are skipped (already handled in Phase 1 via MemoryStore).
Same-timestamp independent ops can be batched for efficiency.
"""
from __future__ import annotations

from typing import Any

import numpy as np

from kernbench.sim_engine.memory_store import MemoryStore, _resolve_dtype
from kernbench.sim_engine.op_log import OpRecord


class DataExecutor:
    """Phase 2 executor: replay op_log with actual numpy computation.

    Args:
        op_log: list of OpRecords from Phase 1.
        store: MemoryStore snapshot from Phase 1 (contains tensor data).
    """

    def __init__(self, op_log: list[OpRecord], store: MemoryStore) -> None:
        self._op_log = op_log
        self.store = store
        # Per-slot time-ordered shape-keyed history. Populated on every
        # ipcq_copy WRITE; consulted on reads that find a shape-mismatched
        # value in MemoryStore (the slot was wrapped by a later inbound
        # before this read's Phase 2 turn). Required because Phase 1 cannot
        # snapshot math-output sources at outbound time (math executes only
        # in Phase 2), so token.data is None and slot wraps lose the recv-
        # time value. See test_attention_mesh_decode_diag (ADR-0059 mesh).
        self._slot_history: dict[tuple[str, int], list[tuple[float, Any]]] = {}

    def _resolve_read(
        self, space: str, addr: int,
        shape: tuple[int, ...] | None, dtype: str | None,
        t_at_or_before: float,
    ) -> Any:
        """Read (space, addr) with expected shape. On KeyError or shape
        mismatch in MemoryStore, fall back to ``_slot_history`` for the
        most recent shape-matching entry with t_write <= t_at_or_before.
        Returns None when no match is found."""
        try:
            return self.store.read(space, addr, shape=shape, dtype=dtype)
        except (KeyError, ValueError):
            pass
        hist = self._slot_history.get((space, addr))
        if hist is None:
            return None
        for t_w, val in reversed(hist):
            if t_w > t_at_or_before:
                continue
            if shape is None or getattr(val, "shape", None) == shape:
                return val
        return None

    # Ordering priority within the same t_start: memory copies must run
    # before math/gemm so that slot data is populated before a consumer
    # PE's math op reads it. With 0-ns PE_MATH overhead and tight SimPy
    # scheduling, ipcq_copy and math ops from different PEs can collide
    # at the exact same t_start.
    _KIND_ORDER = {"memory": 0, "gemm": 1, "math": 2, "unknown": 3}

    def run(self) -> None:
        """Execute all ops in op_log order.

        Primary sort: t_start (ascending).
        Secondary sort: op_kind priority — memory (ipcq_copy/dma_write)
        before gemm before math. This ensures IPCQ slot data arrives
        before a consumer PE's math op tries to read it, even when both
        share the same SimPy timestamp.
        """
        ops = sorted(
            self._op_log,
            key=lambda r: (r.t_start, self._KIND_ORDER.get(r.op_kind, 3)),
        )
        for op in ops:
            if op.op_kind != "memory" or op.op_name != "dma_read":
                self._execute_op(op)

    def _execute_op(self, op: OpRecord) -> None:
        if op.op_kind == "memory":
            self._execute_memory(op)
        elif op.op_kind == "gemm":
            self._execute_gemm(op)
        elif op.op_kind == "math":
            self._execute_math(op)

    def _execute_memory(self, op: OpRecord) -> None:
        """Replay memory copy ops in Phase 2 (ADR-0020 + ADR-0023).

        - dma_read: no-op (handle already references HBM source).
        - dma_write: copy (src_space, src_addr) → (dst_space, dst_addr).
          Required because Phase 2 may have just produced new data at the
          source addr (e.g. PE_MATH scratch output).
        - ipcq_copy: copy across PEs — sender's source → receiver's slot.
          Required because the source may be a Phase 2 math output, and
          a downstream math op on the receiver reads from the slot.

        Legacy entries without src/dst metadata are silently skipped.
        """
        p = op.params
        if op.op_name == "dma_write" or op.op_name == "ipcq_copy":
            src_space = p.get("src_space")
            src_addr = p.get("src_addr")
            dst_space = p.get("dst_space")
            dst_addr = p.get("dst_addr")
            if (src_space is None or src_addr is None
                    or dst_space is None or dst_addr is None):
                return
            # Prefer the Phase-1-time snapshot (captured at record_end /
            # outbound) so we don't read from a source that has since been
            # mutated by another op. Fall back to MemoryStore for sources
            # that had no Phase 1 data (e.g. math scratch outputs that
            # only get populated by Phase 2's math replay).
            data = p.get("snapshot")
            if data is None:
                data = self._resolve_read(
                    src_space, src_addr,
                    p.get("shape"), p.get("dtype"), op.t_start,
                )
                if data is None:
                    return
            self.store.write(dst_space, dst_addr, data)
            # Record this write in slot history so a later forwarded read
            # at src=dst_addr (a different ipcq_copy whose src is this slot)
            # can recover by shape even after the slot has been wrapped.
            if op.op_name == "ipcq_copy":
                self._slot_history.setdefault(
                    (dst_space, dst_addr), [],
                ).append((
                    op.t_start,
                    data.copy() if hasattr(data, "copy") else data,
                ))

    def _execute_gemm(self, op: OpRecord) -> None:
        """Execute GEMM: out = a @ b."""
        p = op.params
        if "src_a_addr" not in p:
            return  # composite record without full params
        default_space = p.get("addr_space", "tcm")
        # ADR-0027: per-operand + output spaces (fall back to single space
        # for legacy records without explicit space keys).
        src_a_space = p.get("src_a_space", default_space)
        src_b_space = p.get("src_b_space", default_space)
        dst_space = p.get("dst_space", default_space)
        dtype_in = p.get("dtype_in", "f16")
        dtype_out = p.get("dtype_out", dtype_in)

        # Phase 2 replay is best-effort: if an operand's data was never
        # produced (e.g. a torch.empty bench input, or a composite whose
        # numerics aren't exercised), skip the matmul rather than crash.
        a = self._resolve_read(src_a_space, p["src_a_addr"],
                               p.get("shape_a"), dtype_in, op.t_start)
        if a is None:
            try:
                a = self.store.read(src_a_space, p["src_a_addr"],
                                    shape=p.get("shape_a"), dtype=dtype_in)
            except (KeyError, ValueError):
                return
        b = self._resolve_read(src_b_space, p["src_b_addr"],
                               p.get("shape_b"), dtype_in, op.t_start)
        if b is None:
            try:
                b = self.store.read(src_b_space, p["src_b_addr"],
                                    shape=p.get("shape_b"), dtype=dtype_in)
            except (KeyError, ValueError):
                return

        # Compute in higher precision if specified
        dtype_acc = p.get("dtype_acc", "f32")
        a_f = a.astype(_resolve_dtype(dtype_acc))
        b_f = b.astype(_resolve_dtype(dtype_acc))
        result = np.matmul(a_f, b_f)

        # Accumulate into the existing output when requested (recipe O +=
        # P·V — the GEMM folds into the already-rescaled accumulator). The
        # composite is recorded at completion so the rescale has run first.
        if p.get("accumulate"):
            try:
                existing = self.store.read(dst_space, p["dst_addr"],
                                           shape=result.shape, dtype=dtype_out)
                result = existing.astype(_resolve_dtype(dtype_acc)) + result
            except (KeyError, ValueError):
                pass

        self.store.write(dst_space, p["dst_addr"],
                         result.astype(_resolve_dtype(dtype_out)))

    def _execute_math(self, op: OpRecord) -> None:
        """Execute math op: unary, binary, or reduction."""
        p = op.params
        math_op = p.get("op", op.op_name)
        dtype = p.get("dtype", "f32")
        input_addrs = p.get("input_addrs", [])
        input_shapes = p.get("input_shapes", [])
        # Per-input space/dtype (ADR-0023 CCL accumulation): math ops can
        # mix inputs from different MemoryStore spaces (e.g. acc in "hbm",
        # recv slot in "tcm"). Fall back to legacy single-space mode when
        # the per-input lists are absent.
        input_spaces = p.get("input_spaces") or [p.get("addr_space", "tcm")] * len(input_addrs)
        input_dtypes = p.get("input_dtypes") or [dtype] * len(input_addrs)
        # Per-input data snapshots (ADR-0020 D6): captured at op_log
        # record time. Phase 1 has correct values for slot/HBM addrs at
        # that moment, which lets Phase 2 sidestep the slot-wraparound
        # races where a later round overwrites a slot before this op
        # runs in t_start order.
        snapshots = p.get("input_snapshots") or [None] * len(input_addrs)
        dst_space = p.get("dst_space", p.get("addr_space", "tcm"))

        inputs = []
        for addr, shape, space, idtype, snap in zip(
            input_addrs, input_shapes, input_spaces, input_dtypes, snapshots
        ):
            if snap is not None:
                inputs.append(snap)
                continue
            resolved = self._resolve_read(space, addr, shape, idtype, op.t_start)
            if resolved is None:
                resolved = self.store.read(space, addr, shape=shape, dtype=idtype)
            inputs.append(resolved)

        result = _compute_math(math_op, inputs, p.get("axis"))
        if result is not None:
            self.store.write(dst_space, p["dst_addr"], result)

    def verify(self, expected: dict[tuple[str, int], np.ndarray],
               rtol: float = 1e-3, atol: float = 1e-3) -> dict[str, bool]:
        """Compare MemoryStore contents against expected tensors.

        Args:
            expected: {(space, addr): expected_ndarray}
            rtol, atol: tolerance for floating-point comparison.

        Returns:
            {key_str: passed} dict.
        """
        results = {}
        for (space, addr), exp in expected.items():
            key = f"{space}:0x{addr:x}"
            try:
                actual = self.store.read(space, addr)
                if np.issubdtype(actual.dtype, np.integer):
                    results[key] = bool(np.array_equal(actual, exp))
                else:
                    results[key] = bool(np.allclose(actual, exp, rtol=rtol, atol=atol))
            except KeyError:
                results[key] = False
        return results


def _compute_math(op: str, inputs: list[np.ndarray], axis: int | None) -> np.ndarray | None:
    """Execute a math operation on numpy arrays."""
    if not inputs:
        return None

    x = inputs[0]

    # ADR-0063 §D3.1: copy is the vector-engine byte move used by
    # tl.copy_to to persist a scoped result to the persistent arena.
    # In data mode the bytes flow through as-is — identity op.
    if op == "copy":
        return x

    # Unary
    if op == "exp":
        return np.exp(x)
    if op == "log":
        return np.log(x)
    if op == "sqrt":
        return np.sqrt(x)
    if op == "abs":
        return np.abs(x)
    if op == "sigmoid":
        return 1.0 / (1.0 + np.exp(-x))
    if op == "cos":
        return np.cos(x)
    if op == "sin":
        return np.sin(x)

    # Reduction
    if op == "sum" or op == "rsum":
        return np.sum(x, axis=axis, keepdims=True)
    if op == "max" or op == "rmax":
        return np.max(x, axis=axis, keepdims=True)
    if op == "min":
        return np.min(x, axis=axis, keepdims=True)

    # Softmax (numerically stable)
    if op == "softmax":
        ax = axis if axis is not None else -1
        x_max = np.max(x, axis=ax, keepdims=True)
        e = np.exp(x - x_max)
        s = np.sum(e, axis=ax, keepdims=True)
        return e / s

    # Binary
    if len(inputs) >= 2:
        y = inputs[1]
        if op == "add":
            return x + y
        if op == "sub":
            return x - y
        if op == "mul":
            return x * y
        if op == "div":
            return x / y
        if op == "maximum" or op == "max_elem":
            return np.maximum(x, y)
        if op == "minimum":
            return np.minimum(x, y)
        # softmax_merge recipe ops (ADR-0065 D5). numpy broadcasting covers
        # the recipe's bcast_axis (e.g. (G,1) corr against (G,d) O).
        if op == "exp_diff":
            return np.exp(x - y)
        if op == "mul_bcast":
            return x * y

    # Ternary
    if len(inputs) >= 3:
        if op == "where":
            return np.where(inputs[0], inputs[1], inputs[2])
        if op == "fma":
            return inputs[0] * inputs[1] + inputs[2]
        if op == "clamp":
            return np.minimum(np.maximum(inputs[0], inputs[1]), inputs[2])

    return None
