"""
=============================================================================
CS349E · 第 16 讲 · 04_what_if.py —— 假设分析：SLO、质量要求、硬件价格、模型大小、流量一变，方案怎样变
=============================================================================

【这个脚本要回答的问题】
    部署方案不是一次性的答案，而是"在这组假设下"的答案。哪些假设一变，最优方案就变？
      演示 1：TPOT 的 SLO 从 100 ms 收紧到 15 ms（聊天）：成本和方案怎样变——延迟和成本的前沿。
      演示 2：质量要求：不许用 FP8（只能 BF16）要多花多少钱。
      演示 3：硬件：只有 H100；B200 涨价到 6 美元 / 小时；H200 降到 2.2 美元 / 小时。
      演示 4：模型大小：同样的聊天 SLO，8B / 70B / 405B。
      演示 5：流量：同样的聊天负载，5 到 500 请求 / 秒（只用解析模型）。

【关于数字】演示 1–4 的每个结果都是"解析模型挑前 4 个方案 → 模拟器检验 + 修复（至少 300 秒的到达）→ 选最便宜的"
    （planner.verified_best）；演示 5 流量太大、模拟太慢，只用解析模型。价格是假设值。
【运行】python3 04_what_if.py      （约 10 秒）
【依赖】numpy；本目录的 planner；../lecture01 的 console。
"""

import os
import sys
from dataclasses import replace

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from planner import (B200, H100, H200, LLAMA31_8B, LLAMA31_70B, LLAMA31_405B, PRICES, SCENARIOS, plan,  # noqa: E402
                     verified_best)
from console import print_header, print_observations, print_table  # noqa: E402

CHAT = SCENARIOS["chat"]


def show(rows, headers):
    print_table(headers, rows)


def cell(q, r):
    if q is None:
        return ["（不可行）", "—", "—", "—"]
    return [q.label(), str(q.gpus), f"{q.usd_per_mtok:.3f}", f"{r['tpot_p90'] * 1e3:.0f} ms"]


def demo1_tpot():
    print_header("演示 1：TPOT 的 SLO 收紧（聊天，TTFT ≤ 1 s，20 请求/s）")
    rows, res = [], []
    for tpot in (0.1, 0.05, 0.03, 0.02, 0.015, 0.01):
        q, r = verified_best(replace(CHAT, tpot=tpot))
        res.append((tpot, q))
        rows.append([f"{tpot * 1e3:.0f} ms"] + cell(q, r))
    show(rows, ["TPOT ≤", "方案", "GPU", "美元 / 百万输出 token", "模拟的 TPOT P90"])
    feas = [(t, q) for t, q in res if q is not None]
    first_change = next((t for t, q in feas if q.usd_per_mtok > feas[0][1].usd_per_mtok + 1e-9), None)
    last_t, last_q = feas[-1]
    obs = [
        f"TPOT 从 100 ms 收紧到 {feas[-2][0] * 1e3:.0f} ms，成本都是 {feas[0][1].usd_per_mtok:.3f} 美元/百万：在 20 请求/s 的流量下，"
        "\"凑整\"留下的余量把收紧的 SLO 吸收了（方案里每一步本来就只用了 SLO 的一部分，看模拟的 TPOT P90）；15 ms 时换成了张量并行 2，"
        "每一步更快（每块 GPU 只读一半的权重）。"
        if first_change == last_t else
        f"TPOT 收紧到 {first_change * 1e3:.0f} ms 时成本开始上涨。",
        f"收紧到 {last_t * 1e3:.0f} ms，最便宜的方案变成了 {last_q.label()}（{last_q.usd_per_mtok:.3f} 美元/百万，是宽松时的 "
        f"{last_q.usd_per_mtok / feas[0][1].usd_per_mtok:.1f} 倍）"
        + ("：这是这一讲的所有实验里解耦唯一一次胜出——decode 要用大的张量并行才够快，prefill 用小的更省，共置只能用同一种并行方式，两头顾不上"
           "（第 5 讲 DistServe 的论点）。" if last_q.mode == "disagg" else "。"),
        "延迟和成本的前沿就是这样一段一段换方案连起来的；投机解码（第 6 讲）还能进一步降低每个 token 的延迟，这个规划器没有把它算进去。",
    ]
    if any(q is None for _, q in res):
        obs.append("\"（不可行）\"表示这几种 GPU 和并行度里没有一种能让每一步快到这个程度。")
    print_observations(obs)


def demo2_quality():
    print_header("演示 2：质量要求——不许用 FP8")
    rows = []
    for key in ("chat", "agent"):
        sc = SCENARIOS[key]
        q8, _ = verified_best(sc)
        q16, _ = verified_best(replace(sc, precisions=("bf16",)))
        rows.append([sc.name, q8.label(), f"{q8.usd_per_mtok:.3f}", q16.label(), f"{q16.usd_per_mtok:.3f}",
                     f"{q16.usd_per_mtok / q8.usd_per_mtok:.2f}×"])
    show(rows, ["场景", "允许 FP8：方案", "成本", "只能 BF16：方案", "成本", "贵了"])
    print_observations([
        "不许量化要多付的钱，就是第 3 讲里\"精度和质量\"那笔账的另一面：如果 FP8 在你的任务上的质量损失可以接受（第 3 讲用 KL、"
        "困惑度、下游任务来衡量），这部分钱就能省下来。规划器把质量当成一个约束，而不是一个可以优化的目标——评估质量要另外做。",
    ])


def demo3_hardware():
    print_header("演示 3：硬件与价格（聊天）")
    rows = []
    cases = [("默认价格（H100 2.0、H200 2.6、B200 4.0）", {}, None),
             ("只有 H100", {}, (H100,)),
             ("B200 涨到 6.0", {B200.name: 6.0}, None),
             ("H200 降到 2.2", {H200.name: 2.2}, None)]
    res = []
    for name, over, chips in cases:
        prices = dict(PRICES)
        prices.update(over)
        kw = {"prices": prices}
        if chips:
            kw["chips"] = chips
        q, r = verified_best(CHAT, **kw)
        res.append(q)
        rows.append([name] + cell(q, r)[:3])
    show(rows, ["假设", "方案", "GPU", "美元 / 百万输出 token"])
    d, h, b6, h2 = res
    tie = abs(d.usd_per_mtok - h.usd_per_mtok) < 1e-9
    print_observations([
        (f"默认价格下，2 块 B200 和 4 块 H100 的成本恰好一样（{d.usd_per_mtok:.3f} 美元/百万）：2 × 4.0 = 4 × 2.0 美元/小时——"
         "在 20 请求/s 的流量下，整数台数让两者打平。" if tie else
         f"默认价格下 {d.label()} 最便宜（{d.usd_per_mtok:.3f}），只有 H100 时 {h.usd_per_mtok:.3f}。")
        + f"B200 涨到 6.0 之后，最便宜的变成 {b6.label()}；H200 降到 2.2 时是 {h2.label()}。"
        "规划器的输出永远是\"在这组价格、这个流量下\"的答案。",
        "02 脚本演示 3 按\"每块 GPU 的效率\"比较时 B200 明显最好；流量小的时候，整数台的粒度可以把这个优势抹平——"
        "这也是\"更大的芯片\"对小服务不一定划算的原因之一。",
    ])


def demo4_model():
    print_header("演示 4：模型大小（聊天的 SLO 和负载）")
    rows = []
    for m in (LLAMA31_8B, LLAMA31_70B, LLAMA31_405B):
        q, r = verified_best(replace(CHAT, model=m))
        rows.append([m.name, f"{m.n_params() / 1e9:.0f} B"] + cell(q, r)[:3])
    show(rows, ["模型", "参数", "方案", "GPU", "美元 / 百万输出 token"])
    print_observations([
        "成本大致随模型大小增长（decode 每一步都要读全部权重）：同样的 SLO 和流量，405B 比 70B 贵好几倍，70B 比 8B 贵好几倍。"
        "这也是为什么\"用多大的模型\"本身就是推理效率里最重要的一个决定——蒸馏、MoE（第 3 讲）、投机解码的草稿模型都在用小模型换成本。",
        "405B 的方案必须用张量并行（单卡放不下），通信成了每一步里不可忽略的一项（第 4 讲）。",
    ])


def demo5_scale():
    print_header("演示 5：流量（聊天的负载和 SLO，只用解析模型）")
    rows = []
    for rate in (5, 20, 100, 500):
        sc = replace(CHAT, rate=float(rate))
        p = plan(sc, tail_aware=True)[0]
        eff = p.gpus * p.price * (sc.rate / p.capacity()) / (sc.rate * sc.output * 3600) * 1e6
        rows.append([f"{rate}", p.label(), str(p.gpus), f"{p.usd_per_mtok:.3f}", f"{eff:.3f}", f"{p.capacity() / rate:.2f}"])
    show(rows, ["流量（请求/s）", "方案", "GPU", "按整数台（美元/百万）", "按效率（美元/百万）", "容量 / 流量"])
    print_observations([
        "流量小的时候，整数台的浪费很大（5 请求/s 时，一台机器的容量远大于需要）；流量大了，按整数台的成本逼近\"每块 GPU 的效率\"，"
        "选择也更接近\"效率最高的配置\"。小服务的成本问题常常不是\"选哪种 GPU\"，而是\"怎样和别的服务共享 GPU\"。",
    ])


if __name__ == "__main__":
    demo1_tpot()
    demo2_quality()
    demo3_hardware()
    demo4_model()
    demo5_scale()
    assert LLAMA31_70B is CHAT.model