#!/usr/bin/env python3
"""model_matrix.py — 模型选型与容量推演（无需 API key，本地可跑）

内置「截至 2026-08-30」的官方模型参数与限流基线，演示架构师报价/选型的三步：
1. 给定负载画像（输入/输出 token、QPS、缓存命中率），按模型算月成本
2. 检查是否撞 ITPM/OTPM/RPM（Start 档基线）
3. 输出选型建议（延迟档位 + 5 折 Batch 场景）

数据来源（全部访问于 2026-08-30，易变事实使用前请重验）:
  - 模型表/价格: https://platform.claude.com/docs/en/models/overview
  - 限流: https://platform.claude.com/docs/en/api/rate-limits
  - 缓存计价: https://platform.claude.com/docs/en/build-with-claude/prompt-caching
运行: python3 model_matrix.py
"""

from __future__ import annotations

# (model, $/MTok in, $/MTok out, ctx, max_out, latency tier, [RPM, ITPM, OTPM] @ Start tier)
MODELS = {
    "claude-fable-5":   dict(inp=10.0, out=50.0, ctx=1_000_000, latency="Slower",   start=[1_000, 500_000, 100_000]),
    "claude-opus-5":    dict(inp=5.0,  out=25.0, ctx=1_000_000, latency="Moderate", start=[1_000, 2_000_000, 400_000]),
    "claude-sonnet-5":  dict(inp=2.0,  out=10.0, ctx=1_000_000, latency="Fast",     start=[1_000, 2_000_000, 400_000]),
    "claude-haiku-4-5": dict(inp=1.0,  out=5.0,  ctx=200_000,   latency="Fastest",  start=[1_000, 2_000_000, 400_000]),
}

# 缓存计价系数（官方: 写 1.25x / 读 0.1x，5m TTL）
CACHE_WRITE_MULT, CACHE_READ_MULT = 1.25, 0.10


def monthly_cost(model: str, qps: float, in_tok: int, out_tok: int,
                 cache_hit_rate: float = 0.0, batch: bool = False) -> float:
    """月成本估算（USD）。in_tok 含缓存部分；cache_hit_rate 为可命中缓存的前缀占比。"""
    m = MODELS[model]
    cached = in_tok * cache_hit_rate
    fresh = in_tok - cached
    per_call = (
        fresh * m["inp"] / 1e6
        + cached * m["inp"] * CACHE_READ_MULT / 1e6   # 命中读按 0.1x
        + out_tok * m["out"] / 1e6
    )
    if batch:
        per_call *= 0.5                                # Batch API 全线 5 折
    return per_call * qps * 86400 * 30


def check_rate_limits(model: str, qps: float, in_tok: int, out_tok: int,
                      cache_hit_rate: float = 0.0) -> list[str]:
    """Start 档限流检查。注意: 多数模型 cache_read 不计入 ITPM（cache-aware ITPM）。"""
    rpm, itpm, otpm = MODELS[model]["start"]
    cached = in_tok * cache_hit_rate
    issues = []
    if qps * 60 > rpm:
        issues.append(f"RPM: 需 {qps*60:,.0f} > 上限 {rpm:,}")
    # ITPM 只计未命中输入（fresh）+ 缓存写入；简化：把 cached 当读不计入
    if (in_tok - cached) * qps * 60 > itpm:
        issues.append(f"ITPM: 需 {(in_tok-cached)*qps*60:,.0f} > 上限 {itpm:,}（缓存读取未计入）")
    if out_tok * qps * 60 > otpm:
        issues.append(f"OTPM: 需 {out_tok*qps*60:,.0f} > 上限 {otpm:,}")
    return issues


def main() -> None:
    # 负载画像：企业知识库助手（长 system+检索上下文，中等输出）
    profile = dict(qps=4, in_tok=12_000, out_tok=1_500, cache_hit_rate=0.8)

    print("=== 负载画像 ===")
    print(f"  QPS={profile['qps']}  输入={profile['in_tok']:,} tok/次"
          f"（缓存命中 {profile['cache_hit_rate']:.0%}）  输出={profile['out_tok']:,} tok/次")
    print()
    print(f"{'模型':<20} {'延迟档':<9} {'月成本(USD)':>14} {'Start档限流':>6}  {'Batch月成本':>12}")
    print("-" * 72)
    for name, m in MODELS.items():
        cost = monthly_cost(name, batch=False, **profile)
        batch_cost = monthly_cost(name, batch=True, **profile)
        issues = check_rate_limits(name, **profile)
        flag = "❌ " + "; ".join(issues) if issues else "✅"
        print(f"{name:<20} {m['latency']:<9} {cost:>14,.0f} {flag:>6}  {batch_cost:>12,.0f}")

    print()
    print("=== 结论（本负载画像下） ===")
    haiku = monthly_cost("claude-haiku-4-5", batch=False, **profile)
    fable = monthly_cost("claude-fable-5", batch=False, **profile)
    print(f"  · Fable 5 是 Haiku 4.5 的 {fable/haiku:.0f}x 成本——先在低档模型调好 prompt/检索再升档")
    print(f"  · 80% 缓存命中 + 0.1x 读价 ≈ 输入侧省 "
          f"{(1-profile['cache_hit_rate']*0.9)*100:.0f}%；且缓存读取不占 ITPM（多数模型）")
    print("  · 若无需即时响应，Batch 再砍 50%（代价：小时级延迟）")

    # 数据完整性自检（与官方表核对的关键锚点）
    assert MODELS["claude-haiku-4-5"]["ctx"] == 200_000
    assert all(m["ctx"] == 1_000_000 for n, m in MODELS.items() if n != "claude-haiku-4-5")
    assert CACHE_WRITE_MULT == 1.25 and CACHE_READ_MULT == 0.10
    print("\n锚点自检通过：Haiku 4.5=200K 窗口、其余 1M、缓存 1.25x/0.1x（截至 2026-08-30）")


if __name__ == "__main__":
    main()