"""核心数据模型: Case 输入 / 实现方案结构体 / 评估结果.

约定:
- 所有"时延"内部统一用秒 (float), 输出 csv 时转 us;
- 所有"字节数"用 int (Byte);
- dtype 一律归一化为小写字符串, 如 "bf16"/"fp16"/"fp8"/"fp32".
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field, asdict
from typing import Optional


# ---------------------------------------------------------------------------
# dtype 工具
# ---------------------------------------------------------------------------

DTYPE_BYTES = {
    "fp32": 4,
    "f32": 4,
    "tf32": 4,
    "fp16": 2,
    "f16": 2,
    "bf16": 2,
    "fp8": 1,
    "fp8_e4m3": 1,
    "fp8_e5m2": 1,
    "int8": 1,
}

# Cube 累加器 (L0C) 中元素字节数: 16bit 输入 -> fp32 累加; fp8 输入 -> fp32 累加
L0C_DTYPE_BYTES = 4


def dtype_bytes(dtype: str) -> int:
    key = dtype.strip().lower()
    if key not in DTYPE_BYTES:
        raise ValueError(f"不支持的 dtype: {dtype!r}, 支持 {sorted(DTYPE_BYTES)}")
    return DTYPE_BYTES[key]


def ceil_div(a: int, b: int) -> int:
    return -(-a // b)


def align_up(x: int, align: int) -> int:
    return ceil_div(x, align) * align


def align_down(x: int, align: int) -> int:
    return (x // align) * align


# ---------------------------------------------------------------------------
# Case 输入
# ---------------------------------------------------------------------------

@dataclass
class BmmCase:
    """一个 batch_mat_mul_v3 case 的输入描述.

    对应算子接口:
        A: [batch_a, M, K]  (可带转置: trans_a=True 表示 [batch_a, K, M])
        B: [batch_b, K, N]  (可带转置: trans_b=True 表示 [batch_b, N, K])
        bias: [B, 1, N] 可选
        C: [batch_c, M, N], batch_c = broadcast(batch_a, batch_b)
    """

    case_id: str = ""
    batch_a: int = 1
    batch_b: int = 1
    m: int = 1
    n: int = 1
    k: int = 1
    dtype_a: str = "bf16"
    dtype_b: str = "bf16"
    dtype_c: str = "bf16"          # 输出 C 的 dtype (fp16/fp8 时 fixpipe 随路转换)
    trans_a: bool = False
    trans_b: bool = False
    has_bias: bool = False
    out_nd: bool = True            # 输出是否 ND 格式 (StreamK 要求 ND)
    deterministic_level: int = 0   # 确定性等级, >=2 禁用 StreamK

    # ---- 派生属性 ----
    @property
    def batch_c(self) -> int:
        return max(self.batch_a, self.batch_b)

    @property
    def dtype_in_bytes(self) -> int:
        # A/B 输入元素字节数 (要求 A/B 同 dtype, 不一致时取较大者并在校验中报 warning)
        return max(dtype_bytes(self.dtype_a), dtype_bytes(self.dtype_b))

    @property
    def dtype_out_bytes(self) -> int:
        return dtype_bytes(self.dtype_c)

    @property
    def flops(self) -> float:
        """总计算量 (乘加各计一次)."""
        return 2.0 * self.batch_c * self.m * self.n * self.k

    @property
    def input_bytes(self) -> float:
        """输入总数据量 (按广播前实际存储计)."""
        dt = self.dtype_in_bytes
        return self.batch_a * self.m * self.k * dt + self.batch_b * self.k * self.n * dt

    @property
    def output_bytes(self) -> float:
        return self.batch_c * self.m * self.n * self.dtype_out_bytes

    @property
    def ai_full(self) -> float:
        """case 固有全量算存比 AI_full = 2MNK / (MK + KN + MN) (单 batch)."""
        m, n, k = self.m, self.n, self.k
        denom = m * k + k * n + m * n
        return (2.0 * m * n * k / denom) if denom > 0 else 0.0

    def to_row(self) -> dict:
        d = asdict(self)
        return d


# ---------------------------------------------------------------------------
# 实现方案结构体 (标准结构体定义)
# ---------------------------------------------------------------------------

@dataclass
class ImplPlan:
    """BMM 实现方案 (理论分析输出 / 用户评估输入 共用的标准结构体).

    字段分四组:
    1) 分支与核间切分; 2) 核内 tiling; 3) 存储/Cache 策略; 4) 尾轮与流水策略.

    说明: 本结构体对齐 batch_mat_mul_v3 tiling 的概念层级
    (核间 grid 划分 -> singleCoreM/N/K -> L1 tile -> L0 tile),
    字段名采用理论文档符号, 便于与文档公式直接对照.
    """

    # --- 0) 基本信息 ---
    case_id: str = ""
    branch: str = ""               # 转Matmul / 特殊分支 / MergeBatch / IterBatch /
                                   # StreamK / ASW_Basic / ASW_Basic_降核
    npu: str = "Ascend950PR"
    op: str = "batch_mat_mul_v3"

    # --- 1) 核间切分 (grid 级) ---
    used_core_num: int = 0         # 实际使用 AIC 核数 (降核时 < C)
    split_b: int = 1               # 核间 B 维切分数
    m_cnt: int = 1                 # 核间 M 维切分数
    n_cnt: int = 1                 # 核间 N 维切分数
    grid_k: int = 1                # 核间 K 维切分数 (StreamK > 1)
    core_map: str = ""             # 核间分配策略, 如 "B->M->N线性映射+ASW滑窗蛇形(W=4)"

    # --- 2) 核内 tiling ---
    b_core: int = 0                # 每核 batch 数 (切 B 分支)
    merge_b0: int = 1              # MergeBatch 合并数 (IterBatch/其他 = 1)
    single_core_m: int = 0         # 每核输出 tile M
    single_core_n: int = 0         # 每核输出 tile N
    single_core_k: int = 0         # 每核 K 段长度 (核间切 K 时 < K)
    k_l1: int = 0                  # GM->L1 的 K 向粒度
    b_l1: int = 1                  # L1 内驻留 batch 数 (MergeBatch/IterBatch)
    l1_form: str = ""              # IterBatch L1 形态: a/b/c/d
    base_m: int = 0                # L0 级 tile
    base_n: int = 0
    base_k: int = 0

    # --- 3) 存储/Cache 策略 ---
    l2_policy_in: str = ""         # 输入 L2 策略: allocate(随路驻留) / non_allocate
    l2_policy_out: str = ""        # 输出: resident(驻留L2异步回写) / direct_gm(直写GM)
    swizzle_w: int = 0             # ASW 滑窗宽度 (0 = 不用)
    workspace_bytes: int = 0       # StreamK 中间结果 workspace (驻留 L2)

    # --- 4) 尾轮与流水策略 ---
    tail_strategy: str = ""        # A0 / A1a / A1b / 方案B (仅切 M/N 类分支)
    fixpipe_unitflag: bool = True  # fixpipe 开 unitflag 随路搬出
    out_dtype_bytes: int = 2       # fixpipe 写出元素字节数 (C 矩阵 dtype)

    # --- 备注 ---
    note: str = ""

    def to_row(self) -> dict:
        return asdict(self)

    @staticmethod
    def csv_fields() -> list:
        return list(ImplPlan.__dataclass_fields__.keys())

    @staticmethod
    def from_row(row: dict) -> "ImplPlan":
        """从 csv 行 (字符串字典) 恢复 ImplPlan."""
        kw = {}
        for name, f in ImplPlan.__dataclass_fields__.items():
            if name not in row or row[name] is None or str(row[name]).strip() == "":
                continue
            v = row[name]
            if f.type == "int":
                kw[name] = int(float(v))
            elif f.type == "bool":
                kw[name] = str(v).strip().lower() in ("1", "true", "yes", "y")
            else:
                kw[name] = v
        return ImplPlan(**kw)


# ---------------------------------------------------------------------------
# 评估结果
# ---------------------------------------------------------------------------

@dataclass
class HardwareTiming:
    """各硬件流水级时延 (秒) 与数据量明细."""

    # 搬入 (MTE2)
    gm_read_bytes: float = 0.0     # GM->L1 直读数据量 (不驻留/未命中 L2 的部分)
    l2_read_bytes: float = 0.0     # L2->L1 数据量 (驻留 L2 后重复读命中部分)
    t_mte2_gm: float = 0.0         # GM->L1 时延 (按 GM 带宽, 不累加 L2->L1)
    t_mte2_l2: float = 0.0         # L2->L1 时延 (按 L2 带宽)
    t_mte2: float = 0.0            # 搬入合计 = t_mte2_gm + t_mte2_l2 (两者发生在不同数据上)
    dma_cmd_count: float = 0.0     # GM->L1 DMA 命令次数 (T_cmd 分析用)
    t_dma_cmd: float = 0.0         # DMA 命令固定开销合计

    # 计算 (Cube MMAD)
    cube_flops: float = 0.0        # Cube 实际计算量 (MergeBatch 含冗余)
    t_mmad: float = 0.0

    # 搬出 (Fixpipe)
    fixpipe_bytes: float = 0.0     # 写出数据量 (按 C 矩阵 dtype / StreamK 临时矩阵按 4B)
    t_fixpipe: float = 0.0

    # 归约 (StreamK 专用)
    t_reduce: float = 0.0

    # 汇总
    t_steady: float = 0.0          # 稳态流水时延 = max(各级)
    t_drain: float = 0.0           # 流水排空暴露
    t_total: float = 0.0           # 端到端时延
    bottleneck: str = ""           # 瓶颈级: MTE2_GM / MTE2_L2 / MMAD / FIXPIPE / REDUCE

    def to_row(self, prefix: str = "") -> dict:
        return {prefix + k: v for k, v in asdict(self).items()}


@dataclass
class EvalResult:
    """单个 case 的完整评估输出 (csv 一行的内容)."""

    case: BmmCase = field(default_factory=BmmCase)
    plan: ImplPlan = field(default_factory=ImplPlan)
    timing: HardwareTiming = field(default_factory=HardwareTiming)
    feasible: bool = True          # 方案是否满足硬件约束
    violations: str = ""           # 违反的约束列表 (";" 分隔)
    bound_type: str = ""           # 计算Bound / 访存Bound / 写出Bound / 归约Bound
    advice: str = ""               # 瓶颈分析与优化建议

    def to_row(self) -> dict:
        row = {}
        row.update(self.case.to_row())
        row.update({"plan_" + k: v for k, v in self.plan.to_row().items()})
        row.update(self.timing.to_row())
        row.update({
            "feasible": self.feasible,
            "violations": self.violations,
            "bound_type": self.bound_type,
            "advice": self.advice,
        })
        return row