"""特殊分支: K=0 (纯写值) / K=1 (逐元素乘), 走 AIV 向量通路.

理论依据: 《BMM算子优化分析 v0.98》§九 + docs/02_分支理论/04_特殊分支.md.

K 维是 Cube 存在的意义 (累加深度). K=0/1 时 Cube 的 16x16x16 粒度浪费,
走 AIV 向量通路 (GM->UB->算->GM) 优于 Cube 通路. 与切分正交的前置判断.
"""

from __future__ import annotations

from ..hardware import NpuSpec, ASCEND950PR
from ..models import BmmCase, ImplPlan, HardwareTiming
from ..timing import assemble_timing
from .base import Branch, ConditionCheck


class SpecialBranch(Branch):
    name = "特殊分支"

    def __init__(self, spec: NpuSpec = ASCEND950PR):
        super().__init__(spec)

    def check_conditions(self, case: BmmCase) -> list:
        c1 = case.k <= 1
        checks = [ConditionCheck("1_K<=1 (Cube 无用)", c1, f"K={case.k}")]
        if case.k == 1:
            # K=1 的 AIV 通路恒可用 (issue#12/#17): B>=2*AIV 开 UB 乒乓; B<128 退化为
            # AIV 单缓冲 (无乒乓, 逐 batch 串行搬入), 不再是无方案空洞.
            b = case.batch_c
            pingpong = b >= 2 * self.spec.aiv_num
            mode = "UB乒乓" if pingpong else "AIV单缓冲(逐batch串行, B<2*AIV)"
            checks.append(ConditionCheck(
                "2_K=1的AIV通路: 恒可用 (B>=128 开UB乒乓, 否则单缓冲)",
                True, f"B={b}, 模式={mode}"))
        return checks

    # ------------------------------------------------------------------
    def make_plan(self, case: BmmCase) -> ImplPlan:
        s = self.spec
        if case.k == 0:
            sub = "K=0纯写值"
            mode = ""
            note = "无任何计算, C=bias 或 0, 纯 AIV 写值; 按行均分到 AIV 核"
        else:
            sub = "K=1逐元素乘"
            pingpong = case.batch_c >= 2 * s.aiv_num
            mode = "UB乒乓" if pingpong else "AIV单缓冲"
            note = (f"退化为 C=A⊙B 无累加深度, Cube 16x16x16 粒度浪费 15/16; "
                    f"走 AIV 通路 GM->UB->Mul->GM, {mode} "
                    f"({'B>=2*AIV 双batch乒乓流水' if pingpong else 'B<2*AIV 逐batch单缓冲串行'})")
        return ImplPlan(
            case_id=case.case_id, branch=self.name,
            used_core_num=s.aiv_num,  # 用 AIV 核
            split_b=1, m_cnt=1, n_cnt=1, grid_k=1,
            core_map="AIV 核间按行均分 (无 Cube tile 概念)",
            b_core=0, merge_b0=1,
            single_core_m=0, single_core_n=0, single_core_k=case.k,
            k_l1=0, b_l1=1,
            l1_form="UB驻留(AIV)" if case.k == 0 else "UB驻留(AIV) " + mode,
            base_m=0, base_n=0, base_k=0,
            l2_policy_in="allocate", l2_policy_out="direct_gm",
            swizzle_w=0, workspace_bytes=0,
            tail_strategy="不涉及(AIV逐元素)",
            fixpipe_unitflag=False,
            out_dtype_bytes=case.dtype_out_bytes,
            note=f"{sub}: {note}",
        )

    # ------------------------------------------------------------------
    def evaluate(self, case: BmmCase, plan: ImplPlan) -> HardwareTiming:
        """AIV 通路时延: 瓶颈在搬移 (AIV 算力远剩)."""
        s = self.spec
        b = case.batch_c
        m, n, k = case.m, case.n, case.k
        dt = case.dtype_in_bytes
        out_b = case.dtype_out_bytes

        if k == 0:
            # 纯写值: 仅写出
            t_in, t_compute, t_out = 0.0, 0.0, b * m * n * out_b / s.bw_gm
            in_bytes, cube_flops = 0.0, 0.0
        else:
            # 逐元素乘: 搬入 A+B, 搬出 C, AIV 算力远剩
            in_bytes = b * (m * k + k * n) * dt
            out_bytes = b * m * n * out_b
            t_in = in_bytes / s.bw_gm
            t_out = out_bytes / s.bw_gm
            # AIV 求积吞吐 (近似按 Q_AIV)
            t_compute = b * m * n / s.q_aiv
            cube_flops = float(b * m * n)

        t_total = max(t_in, t_out, t_compute)
        return assemble_timing(
            t_mte2_gm=t_in, t_mte2_l2=0.0, t_dma_cmd=0.0,
            t_mmad=t_compute, t_fixpipe=t_out, t_reduce=0.0, t_drain=0.0,
            gm_read_bytes=in_bytes, l2_read_bytes=0.0, dma_cmd_count=0.0,
            cube_flops=cube_flops,
            fixpipe_bytes=b * m * n * out_b,
        )