"""时延评估引擎: 按 NPU 硬件流水模型计算各级时延.

模型依据 (v0.98 §3.1 + 用户补充说明):

    T_total = max(T_MMAD, T_MTE2, T_MTE1, T_Fixpipe [, T_Reduce]) + T_drain

搬入 (MTE2) 的两条路径:
  1. GM->L1 直读: 数据量按 GM 带宽计时.
     算子软件可配置是否随路驻留 L2 (Cache Hint allocate);
     无论是否驻留, 该次搬入时延都按 GM 带宽计算, 不累加后续 L2->L1.
     驻留 L2 的意义只在于: 后续重复读可以走路径 2.
  2. L2->L1: 数据已驻留 L2 后的再次/重复读取, 按 L2 带宽 (5.2TB/s) 计时.

Cube (MMAD): 按单核算力 Q16 x 使用核数计时. MergeBatch 的冗余计算计入 cube_flops.

Fixpipe 搬出: 写出 dtype = C 矩阵 dtype (fp16/fp8 时随路转换, 数据量按 C dtype 计);
  StreamK 的中间部分和为防精度丢失按 L0C dtype (4B) 计, 且需再次读入 AIV 归约.

流水掩盖: 双缓冲下稳态时延取各级最大; T_drain 为末级排空暴露, 由各分支模型给出.
"""

from __future__ import annotations

from dataclasses import dataclass

from .hardware import NpuSpec, ASCEND950PR
from .models import HardwareTiming


@dataclass
class MoveInPlan:
    """一次稳态搬入的描述 (单核视角, 单位 Byte)."""
    gm_bytes: float = 0.0    # GM->L1 直读量
    l2_bytes: float = 0.0    # L2->L1 命中量
    dma_cmds: float = 0.0    # GM->L1 DMA 命令次数


def eval_mte2(move: MoveInPlan, spec: NpuSpec = ASCEND950PR,
              active_cores: int | None = None) -> tuple[float, float, float, float]:
    """MTE2 搬入时延.

    返回 (t_mte2_gm, t_mte2_l2, t_mte2_total, t_dma_cmd).

    带宽模型: 每核独立 DMA 引擎, 单核带宽份额 = 聚合带宽/C (尾轮文档 §2.4).
    active_cores < C 时聚合带宽 = active_cores x 单核份额 (尾轮不加速).
    """
    cores = active_cores or spec.aic_num
    bw_gm_eff = spec.bw_pc * cores if cores < spec.aic_num else spec.bw_gm
    bw_l2_eff = spec.bw_l2_pc * cores if cores < spec.aic_num else spec.bw_l2

    # move 为单核数据量; 聚合时延 = 单核量 / 单核带宽份额 (等价 总量/聚合带宽)
    t_gm = move.gm_bytes / spec.bw_pc if move.gm_bytes > 0 else 0.0
    t_l2 = move.l2_bytes / spec.bw_l2_pc if move.l2_bytes > 0 else 0.0
    t_cmd = move.dma_cmds * spec.t_cmd
    # GM 直读与 L2 命中是不同数据的两段搬入, 时延相加; DMA 开销叠加在 GM 路径上
    return t_gm, t_l2, t_gm + t_l2 + t_cmd, t_cmd


def eval_mmad(flops_per_core: float, spec: NpuSpec = ASCEND950PR) -> float:
    """Cube 计算时延: 单核计算量 / 单核算力."""
    return flops_per_core / spec.q16 if flops_per_core > 0 else 0.0


def eval_fixpipe(bytes_per_core: float, to_l2: bool,
                 spec: NpuSpec = ASCEND950PR) -> float:
    """Fixpipe 搬出时延.

    to_l2=True: 写出驻留 L2 (5.2TB/s 写口, 异步回写 GM 不占算子内时延);
    to_l2=False: 直写 GM, 与读共享 1.6TB/s 总线.
    """
    if bytes_per_core <= 0:
        return 0.0
    bw = spec.bw_l2_pc if to_l2 else spec.bw_pc
    return bytes_per_core / bw


def eval_streamk_reduce(tile_elems: float, grid_k: int, out_dtype_bytes: int,
                        spec: NpuSpec = ASCEND950PR) -> float:
    """StreamK 单 tile 归约时延 (v0.98 §七).

    部分和 dtype = L0C dtype (4B, 防精度丢失), 驻留 L2, AIV 归约:
      AIC 写部分和 grid_k x tile x 4B / W_L2
      AIV 读回     grid_k x tile x 4B / W_L2
      AIV 求和     grid_k x tile / Q_AIV
      写回         tile x outB / W_L2
    """
    b4 = 4
    w_l2 = spec.bw_l2
    t_write_partial = grid_k * tile_elems * b4 / w_l2
    t_read_back = grid_k * tile_elems * b4 / w_l2
    t_sum = grid_k * tile_elems / spec.q_aiv
    t_write_out = tile_elems * out_dtype_bytes / w_l2
    return t_write_partial + t_read_back + t_sum + t_write_out


def assemble_timing(t_mte2_gm: float, t_mte2_l2: float, t_dma_cmd: float,
                    t_mmad: float, t_fixpipe: float, t_reduce: float,
                    t_drain: float,
                    gm_read_bytes: float, l2_read_bytes: float,
                    dma_cmd_count: float, cube_flops: float,
                    fixpipe_bytes: float) -> HardwareTiming:
    """汇总各级时延, 判定瓶颈."""
    t_mte2 = t_mte2_gm + t_mte2_l2 + t_dma_cmd
    stages = {
        "MTE2_GM": t_mte2_gm + t_dma_cmd,
        "MTE2_L2": t_mte2_l2,
        "MMAD": t_mmad,
        "FIXPIPE": t_fixpipe,
        "REDUCE": t_reduce,
    }
    bottleneck = max(stages, key=stages.get)
    t_steady = max(stages.values())
    return HardwareTiming(
        gm_read_bytes=gm_read_bytes, l2_read_bytes=l2_read_bytes,
        t_mte2_gm=t_mte2_gm, t_mte2_l2=t_mte2_l2, t_mte2=t_mte2,
        dma_cmd_count=dma_cmd_count, t_dma_cmd=t_dma_cmd,
        cube_flops=cube_flops, t_mmad=t_mmad,
        fixpipe_bytes=fixpipe_bytes, t_fixpipe=t_fixpipe,
        t_reduce=t_reduce,
        t_steady=t_steady, t_drain=t_drain, t_total=t_steady + t_drain,
        bottleneck=bottleneck,
    )


def bound_type_of(bottleneck: str) -> str:
    return {
        "MMAD": "计算Bound",
        "MTE2_GM": "访存Bound(GM)",
        "MTE2_L2": "访存Bound(L2)",
        "FIXPIPE": "写出Bound",
        "REDUCE": "归约Bound",
    }.get(bottleneck, "")