"""GEMM Block: 2-stage MAC pipeline (fetch + compute).

Stage 1 — Fetch (_fetch_stage):
  Collects DMA completion triggers (one per operand per tile).
  When all operands arrive, issues TCM read request (SPMem → MAC registers).

Stage 2 — Compute (_gemm_stage):
  Models MAC array computation time.
  Issues TCM write request (MAC result → SPMem).
  Routes output trigger to MathBlock, DmaWbBlock, or completion.

TCM access goes through TcmBlock for real BW serialization.
MAC compute time is cycle-accurate: ceil(Tm/mac_m) * ceil(Tk/mac_k) * ceil(Tn/mac_n).
"""
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 GemmDescriptor, Trigger


class GemmBlock:
    """2-stage MAC pipeline shared across all concurrent pipelines.

    Pipelines pre-load GemmDescriptors keyed by (pipeline_id, tile_id).
    Two SimPy processes run concurrently: _fetch_stage and _gemm_stage.
    """

    def __init__(
        self,
        env: simpy.Environment,
        trig_in: simpy.Store,
        fetch_to_gemm_trig: simpy.Store,
        to_math_trig: simpy.Store,
        to_dmaWB_trig: simpy.Store,
        completion_q: simpy.Store,
        *,
        tcm_port: simpy.Store,
        mac_m: int,
        mac_k: int,
        mac_n: int,
        bytes_per_element: int,
        clock_freq_ghz: float,
    ) -> None:
        self.env = env
        self.trig_in = trig_in
        self.fetch_to_gemm_trig = fetch_to_gemm_trig
        self.to_math_trig = to_math_trig
        self.to_dmaWB_trig = to_dmaWB_trig
        self.completion_q = completion_q
        self._tcm_port = tcm_port

        self._mac_m = mac_m
        self._mac_k = mac_k
        self._mac_n = mac_n
        self._bpe = bytes_per_element
        self._freq = clock_freq_ghz

        self._descriptor_table: dict[tuple[int, int], GemmDescriptor] = {}

        # Per-pipeline timing histograms
        self.t_tcm_load_per_tile: dict[int, list[float]] = {}
        self.t_compute_per_tile: dict[int, list[float]] = {}

    def load_descriptors(self, descs: dict[tuple, GemmDescriptor]) -> None:
        """Pre-load per-tile GEMM descriptors (cumulative across pipelines)."""
        self._descriptor_table.update(descs)

    def _compute_ns(self, desc: GemmDescriptor) -> float:
        """MAC array compute time for one tile (ns)."""
        cycles = ceil(desc.Tm / self._mac_m) * ceil(desc.Tk / self._mac_k) * ceil(desc.Tn / self._mac_n)
        return cycles / self._freq

    # -- Stage 1: Fetch (TCM → MAC load) --------------------------------------

    def _fetch_stage(self):
        """Collect DMA triggers per tile, issue TCM read for operand load."""
        pending: dict[tuple[int, int], list[Trigger]] = {}
        while True:
            trigger = yield self.trig_in.get()
            if trigger is None:
                yield self.fetch_to_gemm_trig.put(None)
                break

            key = (trigger.pipeline_id, trigger.tile_id)
            pending.setdefault(key, []).append(trigger)

            desc = self._descriptor_table.get(key)
            needed = desc.triggers_needed if desc else 2
            if len(pending[key]) < needed:
                continue

            del pending[key]

            # TCM load: read A and B tile data from SPMem → MAC registers
            if desc and desc.gemm_load:
                a_bytes = desc.Tm * desc.Tk * self._bpe
                b_bytes = desc.Tk * desc.Tn * self._bpe
                load_bytes = a_bytes + b_bytes

                t0 = self.env.now
                done = self.env.event()
                yield self._tcm_port.put(TcmRequest("read", load_bytes, done, tag="gemm_load"))
                yield done
                self.t_tcm_load_per_tile.setdefault(trigger.pipeline_id, []).append(
                    self.env.now - t0
                )

            yield self.fetch_to_gemm_trig.put(trigger)

    # -- Stage 2: Compute (MAC array) + Store (MAC → TCM) ---------------------

    def _gemm_stage(self):
        """MAC computation, then TCM store, then route to next block."""
        while True:
            trigger = yield self.fetch_to_gemm_trig.get()
            if trigger is None:
                break

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

            # MAC compute
            if desc and desc.gemm_compute:
                t_compute = self._compute_ns(desc)
                t0 = self.env.now
                if t_compute > 0:
                    yield self.env.timeout(t_compute)
                self.t_compute_per_tile.setdefault(trigger.pipeline_id, []).append(
                    self.env.now - t0
                )

            # Route output
            route = desc.next_block if desc else "MATH"
            out_trig = Trigger(
                tile_id=trigger.tile_id,
                pipeline_id=trigger.pipeline_id,
                source_block="GEMM",
            )

            if route == "MATH":
                yield self.to_math_trig.put(out_trig)
            elif route == "DMAWB":
                # TCM store before writeback
                if desc:
                    c_bytes = desc.Tm * desc.Tn * self._bpe
                    done = self.env.event()
                    yield self._tcm_port.put(TcmRequest("write", c_bytes, done, tag="gemm_store"))
                    yield done
                yield self.to_dmaWB_trig.put(out_trig)
            else:  # "DONE" — C stays in SPMem, no flush
                if desc:
                    c_bytes = desc.Tm * desc.Tn * self._bpe
                    done = self.env.event()
                    yield self._tcm_port.put(TcmRequest("write", c_bytes, done, tag="gemm_store"))
                    yield done
                yield self.completion_q.put(out_trig)
