"""特殊分支: 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 通路需 B >= 2*AIV核数 且单 batch 输入输出能驻留 UB
            c2 = case.batch_c >= 2 * self.spec.aiv_num
            checks.append(ConditionCheck(
                "2_K=1的AIV触发: B >= 2*AIV核数 (开UB乒乓)",
                c2, f"B={case.batch_c} vs {2*self.spec.aiv_num}"))
        return checks

    # ------------------------------------------------------------------
    def make_plan(self, case: BmmCase) -> ImplPlan:
        s = self.spec
        if case.k == 0:
            sub = "K=0纯写值"
            note = "无任何计算, C=bias 或 0, 纯 AIV 写值; 按行均分到 AIV 核"
        else:
            sub = "K=1逐元素乘"
            note = ("退化为 C=A⊙B 无累加深度, Cube 16x16x16 粒度浪费 15/16; "
                    "走 AIV 通路 GM->UB->Mul->GM, UB 乒乓")
        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)",
            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,
        )