"""方案评估器: 对 (case, plan) 做硬件约束校验 + 时延评估 + 瓶颈分析.

评估模式入口: 用户自带实现方案 (ImplPlan), 软件评估其在 NPU 上的
各级硬件时延、流水情况与瓶颈, 并做可行性校验.
"""

from __future__ import annotations

from .hardware import NpuSpec, ASCEND950PR
from .models import BmmCase, ImplPlan, EvalResult
from .timing import bound_type_of
from .branches.merge_batch import MergeBatchBranch
from .branches.iter_batch import IterBatchBranch
from .branches.to_matmul import ToMatmulBranch
from .branches.special import SpecialBranch
from .branches.stream_k import StreamKBranch
from .branches.asw_basic import AswBasicBranch

_BRANCH_EVAL = {
    "MergeBatch": MergeBatchBranch,
    "IterBatch": IterBatchBranch,
    "转Matmul": ToMatmulBranch,
    "特殊分支": SpecialBranch,
    "StreamK": StreamKBranch,
    "ASW_Basic": AswBasicBranch,
    "ASW_Basic_降核": AswBasicBranch,
}


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

    def evaluate(self, case: BmmCase, plan: ImplPlan) -> EvalResult:
        res = EvalResult(case=case, plan=plan)

        # 1) 硬件约束校验
        violations = self._check_constraints(case, plan)
        res.feasible = not violations
        res.violations = "; ".join(violations)

        # 2) 时延评估 (按方案分支调用对应模型)
        branch_cls = _BRANCH_EVAL.get(plan.branch)
        if branch_cls is None:
            res.advice = (f"分支 {plan.branch!r} 的评估模型待后续迭代; "
                          f"当前支持: {sorted(_BRANCH_EVAL)}")
            return res
        timing = branch_cls(self.spec).evaluate(case, plan)
        res.timing = timing
        res.bound_type = bound_type_of(timing.bottleneck)

        # 3) 瓶颈分析建议
        res.advice = self._advice(case, plan, res)
        return res

    # ------------------------------------------------------------------
    def _check_constraints(self, case: BmmCase, plan: ImplPlan) -> list:
        """约束校验: 委托给 constraints.py 单一约束源 (与生成同源, issue#5/#6)."""
        from .constraints import check_plan_constraints
        return check_plan_constraints(case, plan, self.spec)

    # ------------------------------------------------------------------
    def _advice(self, case: BmmCase, plan: ImplPlan, res: EvalResult) -> str:
        t = res.timing
        tips = []
        if not res.feasible:
            tips.append("方案违反硬件约束, 需先修正: " + res.violations)
        bn = t.bottleneck
        if bn == "MTE2_GM":
            tips.append("瓶颈在 GM 搬入: 可考虑增大 tile 提升 dValue/单核搬移量, "
                        "或利用 L2 驻留吸收重复读 (MergeBatch/ASW swizzle 方向)")
        elif bn == "MTE2_L2":
            tips.append("瓶颈在 L2 重复读: 优化核间分配/swizzle 窗口压低活跃工作集")
        elif bn == "MMAD":
            tips.append("瓶颈在 Cube 计算: 已接近理论算力上限, 检查是否有冗余计算 "
                        "(MergeBatch 交叉项) 可消除")
        elif bn == "FIXPIPE":
            tips.append("瓶颈在 Fixpipe 写出: 检查输出 dtype (fp16/fp8 可减半写出量), "
                        "或评估输出驻留 L2 异步回写策略")
        if plan.branch == "MergeBatch" and plan.k_l1 < case.k:
            tips.append("警告: MergeBatch 处于 L1 绑定情形 (k_L1<K), 理论证明其恒劣于 "
                        "IterBatch (v1.1 §4.4), 建议改用 IterBatch")
        return " | ".join(tips) if tips else "方案合理, 流水均衡"