"""Data-mode numeric tests for the softmax_merge recipe MATH ops (ADR-0065 N2).

N2 makes the recipe's 8 MATH ops compute in data mode: `_compute_math` gains
rmax/rsum/max_elem/exp_diff/mul_bcast, the recipe MATH stages carry operand
addresses (`_math_stage`), and op_log `record_end` promotes an address-
carrying MATH stage to op_kind="math". The MATH ops fully compute the online-
softmax update of `m` and `l` (and rescale `O`); the `+P·V` of `O` is the
GEMM-accumulate step (N3).
"""
from __future__ import annotations

import numpy as np

from kernbench.common.pe_commands import CompositeCmd, TensorHandle
from kernbench.components.builtin.tiling import generate_plan_from_ops
from kernbench.sim_engine.data_executor import DataExecutor, _compute_math
from kernbench.sim_engine.memory_store import MemoryStore
from kernbench.sim_engine.op_log import OpRecord
from kernbench.triton_emu.tl_context import TLContext

G, TILE, D = 8, 64, 128


def test_compute_math_recipe_ops():
    x = np.random.randn(4, 6).astype(np.float32)
    y = np.random.randn(4, 1).astype(np.float32)
    assert np.allclose(_compute_math("rmax", [x], -1), x.max(-1, keepdims=True))
    assert np.allclose(_compute_math("rsum", [x], -1), x.sum(-1, keepdims=True))
    assert np.allclose(_compute_math("max_elem", [y, x.max(-1, keepdims=True)], None),
                       np.maximum(y, x.max(-1, keepdims=True)))
    assert np.allclose(_compute_math("exp_diff", [x, y], None), np.exp(x - y))
    assert np.allclose(_compute_math("mul_bcast", [x, y], None), x * y)


def _tcm(addr, shape):
    return TensorHandle(id=f"h{addr:x}", addr=addr, shape=shape, dtype="f16",
                        nbytes=2 * int(np.prod(shape)), space="tcm")


def test_softmax_merge_computes_m_l_O_rescale():
    """Run the recipe's prologue MATH ops through the DataExecutor and check
    m, l (fully updated) and O (rescaled by corr) match a numpy reference."""
    rng = np.random.default_rng(0)
    s = rng.standard_normal((G, TILE)).astype(np.float16)
    m0 = rng.standard_normal((G, 1)).astype(np.float16)
    l0 = np.abs(rng.standard_normal((G, 1))).astype(np.float16)
    O0 = rng.standard_normal((G, D)).astype(np.float16)

    sh = _tcm(0x1000, (G, TILE))
    mh = _tcm(0x2000, (G, 1))
    lh = _tcm(0x3000, (G, 1))
    Oh = _tcm(0x4000, (G, D))
    V = TensorHandle(id="V", addr=0x5000, shape=(TILE, D), dtype="f16",
                     nbytes=2 * TILE * D, space="hbm")

    tl = TLContext(pe_id=0, num_programs=1, scratch_base=0x100000,
                   scratch_size=1 << 20)
    tl.composite(
        prologue=[{"op": "softmax_merge", "s": sh, "m": mh, "l": lh, "O": Oh}],
        op="gemm", b=V, out=Oh,
    )
    cmd = [c for c in tl.commands if isinstance(c, CompositeCmd)][-1]
    plan = generate_plan_from_ops(cmd.ops, tile_m=32, tile_k=32, tile_n=32,
                                  bytes_per_element=2,
                                  pe_prefix="sip0.cube0.pe0")

    store = MemoryStore()
    store.write("tcm", sh.addr, s)
    store.write("tcm", mh.addr, m0)
    store.write("tcm", lh.addr, l0)
    store.write("tcm", Oh.addr, O0)

    records = [
        OpRecord(t_start=float(i), t_end=float(i), component_id="pe_math",
                 op_kind="math", op_name=st.params["op"], params=dict(st.params))
        for i, st in enumerate(plan.prologue_stages)
    ]
    DataExecutor(records, store).run()

    # numpy reference (online-softmax merge of a single new tile into m/l/O).
    sf = s.astype(np.float32)
    m_loc = sf.max(-1, keepdims=True)
    m_new = np.maximum(m0.astype(np.float32), m_loc)
    corr = np.exp(m0.astype(np.float32) - m_new)
    P = np.exp(sf - m_new)
    l_new = l0.astype(np.float32) * corr + P.sum(-1, keepdims=True)
    O_resc = O0.astype(np.float32) * corr

    m_got = store.read("tcm", mh.addr, shape=(G, 1), dtype="f16").astype(np.float32)
    l_got = store.read("tcm", lh.addr, shape=(G, 1), dtype="f16").astype(np.float32)
    O_got = store.read("tcm", Oh.addr, shape=(G, D), dtype="f16").astype(np.float32)

    assert np.allclose(m_got, m_new, atol=2e-2), f"m: {m_got.ravel()[:3]} vs {m_new.ravel()[:3]}"
    assert np.allclose(l_got, l_new, rtol=5e-2, atol=5e-2), "l mismatch"
    assert np.allclose(O_got, O_resc, rtol=5e-2, atol=5e-2), "O rescale mismatch"