"""Tile schedule generators for GEMM and element-wise math operations.

Each generator produces a plan of tile commands with pre-computed addresses.
Pipelines use these plans to build descriptor tables and feed commands
to the shared hardware blocks.
"""
from __future__ import annotations

from math import ceil

from kernbench.components.custom.pe_accel.types import (
    CmdType,
    MathSchedulePlan,
    MathTileCommand,
    SchedulePlan,
    TileCommand,
)


def generate_gemm_tiles(
    M_tiles: int, K_tiles: int, N_tiles: int,
    M: int = 0, K: int = 0, N: int = 0,
    tile_m: int = 0, tile_k: int = 0, tile_n: int = 0,
    bytes_per_element: int = 2,
    A_addr: int = 0, B_addr: int = 0, C_addr: int = 0,
    pipeline_id: int = 0,
) -> SchedulePlan:
    """Generate GEMM tile commands in M (outer) -> N -> K (inner) order.

    Stamps is_last_k=True on the final K-tile per (m, n) pair.
    Emits one DMA_FLUSH per (m, n) pair after all K tiles.

    Per-tile addresses (row-major layout):
      A (M,K): A_addr + (m * tile_m * K + k * tile_k) * bpe
      B (K,N): B_addr + (k * tile_k * N + n * tile_n) * bpe
      C (M,N): C_addr + (m * tile_m * N + n * tile_n) * bpe
    """
    commands: list[TileCommand] = []
    cmd_id = 0
    tile_id = 0
    bpe = bytes_per_element

    for m in range(M_tiles):
        for n in range(N_tiles):
            c_tile_addr = C_addr + (m * tile_m * N + n * tile_n) * bpe

            for k in range(K_tiles):
                last_k = k == K_tiles - 1
                a_tile_addr = A_addr + (m * tile_m * K + k * tile_k) * bpe
                b_tile_addr = B_addr + (k * tile_k * N + n * tile_n) * bpe

                commands.append(TileCommand(
                    cmd_id=cmd_id, cmd_type=CmdType.DMA_LOAD,
                    tile_id=tile_id, m_idx=m, k_idx=k, n_idx=n,
                    is_last_k=last_k, pipeline_id=pipeline_id,
                    a_tile_addr=a_tile_addr, b_tile_addr=b_tile_addr,
                    c_tile_addr=c_tile_addr,
                ))
                cmd_id += 1

                commands.append(TileCommand(
                    cmd_id=cmd_id, cmd_type=CmdType.TENSOR_OP,
                    tile_id=tile_id, m_idx=m, k_idx=k, n_idx=n,
                    is_last_k=last_k, pipeline_id=pipeline_id,
                    a_tile_addr=a_tile_addr, b_tile_addr=b_tile_addr,
                    c_tile_addr=c_tile_addr,
                ))
                cmd_id += 1
                tile_id += 1

            # One flush per (m, n) pair after all K tiles
            commands.append(TileCommand(
                cmd_id=cmd_id, cmd_type=CmdType.DMA_FLUSH,
                tile_id=tile_id - 1, m_idx=m, k_idx=0, n_idx=n,
                pipeline_id=pipeline_id,
                c_tile_addr=c_tile_addr,
            ))
            cmd_id += 1

    return SchedulePlan(
        commands=commands, M_tiles=M_tiles, K_tiles=K_tiles, N_tiles=N_tiles
    )


def generate_math_tiles(
    M_tiles: int, N_tiles: int,
    M: int = 0, N: int = 0,
    tile_m: int = 0, tile_n: int = 0,
    bytes_per_element: int = 2,
    src_addr: int = 0, dst_addr: int = 0,
    pipeline_id: int = 0,
) -> MathSchedulePlan:
    """Generate element-wise math tile commands in row-major order.

    Per-tile addresses (row-major layout):
      src: src_addr + (m * tile_m * N + n * tile_n) * bpe
      dst: dst_addr + (m * tile_m * N + n * tile_n) * bpe
    """
    commands: list[MathTileCommand] = []
    cmd_id = 0
    tile_id = 0
    bpe = bytes_per_element

    for m in range(M_tiles):
        for n in range(N_tiles):
            offset = (m * tile_m * N + n * tile_n) * bpe
            commands.append(MathTileCommand(
                cmd_id=cmd_id,
                tile_id=tile_id,
                m_idx=m,
                n_idx=n,
                src_tile_addr=src_addr + offset,
                dst_tile_addr=dst_addr + offset,
                pipeline_id=pipeline_id,
            ))
            cmd_id += 1
            tile_id += 1

    return MathSchedulePlan(
        commands=commands, M_tiles=M_tiles, N_tiles=N_tiles
    )
