"""时延评估引擎: 按 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 BmmCase, HardwareTiming


def output_to_l2(case: BmmCase, spec: NpuSpec = ASCEND950PR,
                 workspace_bytes: float = 0.0) -> bool:
    """Fixpipe 输出落点决策 (issue#30, 设计文档 docs/05 §4.2 R4).

    to_l2 (输出写 L2 写口 5.2TB/s, GM 写流量 = 0, 异步回写不占算子时延)
      ⟺  整 case 输入 V_in + 输出 V_out [+ StreamK workspace] ≤ L2
    否则输出**直写 GM**: 输入优先驻留 L2 (输入存在重复读), 输出计入 GM
    读写共享总线 (与读累加, issue#23).
    """
    return (case.input_bytes + case.output_bytes + workspace_bytes) <= spec.l2_bytes


@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,
              dtype_a=None, dtype_b=None) -> float:
    """Cube 计算时延: 单核计算量 / 单核 dtype 感知算力 (issue#28)."""
    rate = spec.q_cube(dtype_a, dtype_b)
    return flops_per_core / rate 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, out_to_gm: bool = False) -> float:
    """StreamK 单 tile 归约时延 (v0.98 §七).

    部分和 dtype = L0C dtype (4B, 防精度丢失), 驻留 L2, AIV 归约 (fp32 求和,
    AIV 按 fp32 通量, 不随输入 dtype 变 — issue#28):
      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(fp32)
      写回         tile x outB / W (最终输出落点, issue#30):
                   整 case 可驻留 (S_A) 时 W = W_L2; 否则直写 GM (W = W_GM)
    """
    b4 = 4
    w_l2 = spec.bw_l2
    w_out = spec.bw_gm if out_to_gm else w_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_out
    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,
                    reduce_serial: bool = True,
                    fixpipe_to_gm: bool = True) -> HardwareTiming:
    """汇总各级时延, 判定瓶颈.

    带宽端口口径 (issue#23, 用户澄清 + KB):
    - GM 1.6TB/s 为**读写共享总线**: MTE2 的 GM 读与 Fixpipe 直写 GM 并发时无法
      拆分读写占用, 时延累加 (读+写)/1.6TB/s, 并入 MTE2 搬移链;
    - L2 带宽读写各自独享 5.2TB/s: L2 重复读段(读口) 与 Fixpipe→L2 写(写口)
      互不竞争; MTE2 引擎顺序服务 GM/L2 装载, 两段相加 (官方 T≈HBM/1.6+L2/5.2);
    - Fixpipe→L2 (resident) 写出独立为 FIXPIPE 级, 不进 GM 总线.

    归约计账约定 (issue#9): REDUCE 默认**串行追加** (reduce_serial=True, 归约不可
    掩盖, 体现在 t_drain 中), 不进稳态 max(); 仅当调用方显式声明归约可流水掩盖
    (reduce_serial=False) 时才进稳态 max(). 避免"既取最大又串行追加"的双倍计账.
    """
    t_mte2 = t_mte2_gm + t_mte2_l2 + t_dma_cmd
    fix_gm = t_fixpipe if fixpipe_to_gm else 0.0   # Fixpipe 直写 GM 的时延
    fix_l2 = 0.0 if fixpipe_to_gm else t_fixpipe   # Fixpipe→L2 (5.2TB/s 写口)
    # MTE2 搬移链 = GM 总线(读写共享, 读+直写累加) + L2 重复读段 + DMA 命令开销
    # (同一 MTE2 引擎顺序服务 GM/L2 两类装载; GM 写由 Fixpipe 并发发起, 共享 GM 总线)
    mte2_chain = t_mte2_gm + t_dma_cmd + fix_gm + t_mte2_l2
    stages = {
        "MTE2": mte2_chain,
        "MMAD": t_mmad,
        "FIXPIPE": fix_l2,   # 仅 Fixpipe→L2 (5.2 写口, 与 L2 读口互不竞争)
    }
    if not reduce_serial:
        stages["REDUCE"] = t_reduce
    bottleneck = max(stages, key=stages.get)
    t_steady = max(stages.values())
    # 归约串行追加时, 若它是全链路最大项则瓶颈标注为 REDUCE (但时延只计一次)
    if reduce_serial and t_reduce > t_steady:
        bottleneck = "REDUCE"
    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": "访存Bound(GM读写共享+L2重复读)",
        "FIXPIPE": "写出Bound(L2写口)",
        "REDUCE": "归约Bound",
    }.get(bottleneck, "")