"""Math Block: K-accumulation (GEMM helper) + element-wise ops (exp, log, etc.).

Two concurrent processing modes:

  1. K-accumulation (_run_k_accumulation):
     Receives triggers from GemmBlock after each K-tile compute.
     Issues TCM write for partial-result store.
     On final K-tile, routes to DMA_WB or completion.

  2. Element-wise ops (_run_element_wise):
     Receives triggers from DMA_IN after each tile read.
     Issues TCM read (load input), compute (SIMD), TCM write (store result).
     Routes output to DMA_WB for writeback.

TCM access goes through TcmBlock for real BW serialization.
SIMD compute time: ceil(num_elements / vector_width) / clock_freq.
"""
from __future__ import annotations

from math import ceil

import simpy

from kernbench.components.custom.pe_accel.blocks.tcm import TcmRequest
from kernbench.components.custom.pe_accel.types import MathDescriptor, MathOpDescriptor, Trigger


class MathBlock:
    """K-accumulation + element-wise math unit.

    Descriptor tables:
      - _accum_table:    MathDescriptor    (pipeline_id, tile_id) — for GEMM K-accumulation
      - _elemwise_table: MathOpDescriptor  (pipeline_id, tile_id) — for element-wise ops
    """

    def __init__(
        self,
        env: simpy.Environment,
        trig_in: simpy.Store,
        to_dmaWB_trig: simpy.Store,
        completion_q: simpy.Store,
        *,
        tcm_port: simpy.Store,
        bytes_per_element: int,
        clock_freq_ghz: float,
        vector_width: int = 256,
    ) -> None:
        self.env = env
        self.trig_in = trig_in              # from GemmBlock (K-accumulation)
        self.to_dmaWB_trig = to_dmaWB_trig
        self.completion_q = completion_q
        self._tcm_port = tcm_port

        self._bpe = bytes_per_element
        self._freq = clock_freq_ghz
        self._vector_width = vector_width

        # Descriptor tables
        self._accum_table: dict[tuple[int, int], MathDescriptor] = {}
        self._elemwise_table: dict[tuple[int, int], MathOpDescriptor] = {}

        # -- Timing histograms (per pipeline_id) --

        # K-accumulation
        self.t_tcm_store_per_tile: dict[int, list[float]] = {}

        # Element-wise ops
        self.t_math_op_load_per_tile: dict[int, list[float]] = {}
        self.t_math_op_compute_per_tile: dict[int, list[float]] = {}
        self.t_math_op_store_per_tile: dict[int, list[float]] = {}

    # -- Descriptor loading ----------------------------------------------------

    def load_descriptors(self, descs: dict[tuple, MathDescriptor]) -> None:
        """Pre-load K-accumulation descriptors (cumulative across pipelines)."""
        self._accum_table.update(descs)

    def load_math_op_descriptors(self, descs: dict[tuple, MathOpDescriptor]) -> None:
        """Pre-load element-wise op descriptors (cumulative across pipelines)."""
        self._elemwise_table.update(descs)

    # -- Mode 1: K-accumulation ------------------------------------------------

    def _run(self):
        """Backward-compat alias."""
        yield from self._run_k_accumulation()

    def _run_k_accumulation(self):
        """Receive GEMM output triggers, TCM store partial result, flush on last-K."""
        while True:
            trigger = yield self.trig_in.get()
            if trigger is None:
                break

            key = (trigger.pipeline_id, trigger.tile_id)
            desc = self._accum_table.get(key)

            # TCM store: write partial sum to SPMem
            if desc:
                c_bytes = desc.Tm * desc.Tn * self._bpe
                t0 = self.env.now
                done = self.env.event()
                yield self._tcm_port.put(TcmRequest("write", c_bytes, done, tag="k_accum_store"))
                yield done
                self.t_tcm_store_per_tile.setdefault(trigger.pipeline_id, []).append(
                    self.env.now - t0
                )

            if not desc or not desc.is_last_k:
                continue  # intermediate K-tile: store done, no flush yet

            out_trig = Trigger(
                tile_id=trigger.tile_id,
                pipeline_id=trigger.pipeline_id,
                source_block="MATH",
            )
            if desc.skip_dmaWB:
                yield self.completion_q.put(out_trig)
            else:
                yield self.to_dmaWB_trig.put(out_trig)

    # -- Mode 2: Element-wise ops ----------------------------------------------

    def _run_math_op(self, trig_q: simpy.Store):
        """Backward-compat alias."""
        yield from self._run_element_wise(trig_q)

    def _run_element_wise(self, trig_q: simpy.Store):
        """Receive DMA_IN triggers, apply element-wise op via TCM, route to DMA_WB.

        Per tile:
          1. TCM read  — load input tile from SPMem to SIMD
          2. Compute   — SIMD operation (exp/log/etc.)
          3. TCM write — store result from SIMD to SPMem
          4. Route to DMA_WB
        """
        while True:
            trigger = yield trig_q.get()
            if trigger is None:
                break

            key = (trigger.pipeline_id, trigger.tile_id)
            desc = self._elemwise_table.get(key)

            if desc:
                tile_bytes = desc.Tm * desc.Tn * self._bpe
                num_elements = desc.Tm * desc.Tn

                # 1. TCM read
                t0 = self.env.now
                done = self.env.event()
                yield self._tcm_port.put(TcmRequest("read", tile_bytes, done, tag="elemwise_load"))
                yield done
                self.t_math_op_load_per_tile.setdefault(trigger.pipeline_id, []).append(
                    self.env.now - t0
                )

                # 2. SIMD compute
                t0 = self.env.now
                compute_cycles = ceil(num_elements / self._vector_width)
                compute_ns = compute_cycles / self._freq
                if compute_ns > 0:
                    yield self.env.timeout(compute_ns)
                self.t_math_op_compute_per_tile.setdefault(trigger.pipeline_id, []).append(
                    self.env.now - t0
                )

                # 3. TCM write
                t0 = self.env.now
                done = self.env.event()
                yield self._tcm_port.put(TcmRequest("write", tile_bytes, done, tag="elemwise_store"))
                yield done
                self.t_math_op_store_per_tile.setdefault(trigger.pipeline_id, []).append(
                    self.env.now - t0
                )

            yield self.to_dmaWB_trig.put(Trigger(
                tile_id=trigger.tile_id,
                pipeline_id=trigger.pipeline_id,
                source_block="MATH_OP",
            ))
