"""Reproduce the 2026-09-10 serving scenarios in the blog.

Run with Python 3; uses only the standard library. Dollar costs and kappa are
author inputs. API prices and open-model parameter counts have linked sources.
"""
from pathlib import Path
import csv
import json
import math

DATA = Path(__file__).resolve().parents[1] / "data"
ASSUMPTIONS = json.loads((DATA / "model-assumptions.json").read_text())
PRICES = {
    "kimi_k3": {"input": 3, "cache_read": .3, "cache_write": 3, "output": 15},
    "gpt55": {"input": 5, "cache_read": .5, "cache_write": None, "output": 30},
    "gpt56_sol": {"input": 4, "cache_read": .4, "cache_write": 5, "output": 20},
    "gpt6_astra": {"input": 10, "cache_read": 1, "cache_write": 12.5, "output": 50},
}
SOURCES = {
    "kimi_prices": "https://platform.kimi.ai/",
    "openai_prices": "https://developers.openai.com/api/docs/pricing",
    "gpt55_model": "https://developers.openai.com/api/docs/models/gpt-5.5",
    "sol_model": "https://developers.openai.com/api/docs/models/gpt-5.6-sol",
    "cache_write_rules": "https://developers.openai.com/api/docs/guides/prompt-caching",
    "astra_long_context": "https://developers.openai.com/api/docs/models/gpt-6-astra",
    "kimi_parameters": "https://huggingface.co/moonshotai/Kimi-K3",
    "deepseek_parameters": "https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro",
    "nvl72_hardware": "https://www.nvidia.com/en-us/data-center/gb200-nvl72/",
    "hgx_hardware": "https://docs.nvidia.com/enterprise-reference-architectures/hgx-ai-factory/latest/components.html",
    "kimi_8k_benchmark": "https://zenn.dev/fixstars/articles/kimi-k3-benchmark",
}


def bill(model, input_tokens, output_tokens, cache_hit, miss_write_fraction=0):
    p = PRICES[model]
    long_context = model != "kimi_k3" and input_tokens > 272_000
    input_factor, output_factor = (2, 1.5) if long_context else (1, 1)
    if p["cache_write"] is None and miss_write_fraction:
        raise ValueError("This model has no separate cache-write rate in the source table")
    write_price = p["input"] if p["cache_write"] is None else p["cache_write"]
    miss_price = ((1 - miss_write_fraction) * p["input"]
                  + miss_write_fraction * write_price)
    return (input_tokens * ((1 - cache_hit) * miss_price
            + cache_hit * p["cache_read"]) * input_factor
            + output_tokens * p["output"] * output_factor) / 1e6


def threshold(hourly_cost, revenue_per_million_output, duty=.7, surplus=.3):
    return hourly_cost * 1e6 / (3600 * duty * (1 - surplus)
                              * revenue_per_million_output)


def save_csv(name, rows):
    with (DATA / name).open("w", encoding="utf-8-sig", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def main():
    bills = []
    for input_tokens in (100_000, 300_000):
        for hit in (.8, .85, .9):
            for model in PRICES:
                for write_fraction in ((0,) if model in ("kimi_k3", "gpt55") else (0, 1)):
                    v = bill(model, input_tokens, 1000, hit, write_fraction)
                    bills.append(dict(model=model, input_tokens=input_tokens,
                        output_tokens=1000, cache_hit=hit,
                        miss_write_fraction=write_fraction, bill_usd=v,
                        revenue_usd_per_million_output=v * 1000))
    save_csv("agent-serving-bills.csv", bills)

    thresholds = []
    for r in bills:
        if r["input_tokens"] != 100_000 or r["model"] not in ("kimi_k3", "gpt6_astra"):
            continue
        for cost in ((80,) if r["model"] == "kimi_k3" else (360, 720)):
            for surplus in (0, .3):
                q = threshold(cost, r["revenue_usd_per_million_output"], surplus=surplus)
                thresholds.append(dict(**r, assumed_hourly_cost_usd=cost,
                    assumed_duty=.7, target_resource_surplus_fraction=surplus,
                    required_active_output_tokens_per_second=q))
                revenue = q * 3600 * .7 / 1e6 * r["revenue_usd_per_million_output"]
                assert math.isclose(revenue * (1 - surplus), cost)
    save_csv("agent-serving-thresholds.csv", thresholds)

    kappa_rows = []
    for write_fraction in (0, 1):
        price_ratio = bill("gpt6_astra", 100_000, 1000, .85, write_fraction) / bill("kimi_k3", 100_000, 1000, .85)
        for kappa in (.75, 1, 1.25):
            kappa_rows.append(dict(astra_miss_write_fraction=write_fraction,
                cache_hit=.85, input_tokens=100_000, output_tokens=1000,
                reference_active_parameters_b=104, revenue_ratio=price_ratio,
                assumed_kappa=kappa, implied_astra_active_parameters_b=104 * price_ratio * kappa))
    save_csv("astra-active-calibration.csv", kappa_rows)

    generation_rows = []
    for model in PRICES:
        v = bill(model, 100_000, 1000, .85)
        ratio = v / bill("kimi_k3", 100_000, 1000, .85)
        generation_rows.append(dict(model=model, input_tokens=100_000, output_tokens=1000,
            cache_hit=.85, miss_write_fraction=0, bill_usd=v,
            revenue_ratio_vs_kimi=ratio, assumed_kappa=1,
            price_normalized_active_parameters_b=104*ratio,
            parameter_status="reported" if model == "kimi_k3" else "conditional_price_normalization"))
    save_csv("generation-price-calibration.csv", generation_rows)

    model_rows = []
    project_flops = (ASSUMPTIONS['gpu_count'] * ASSUMPTIONS['gpu_dense_bf16_flops_s']
        * ASSUMPTIONS['project_days'] * 86400 * ASSUMPTIONS['effective_fraction'])
    pretrain_share = ASSUMPTIONS['main_pretraining_fraction']
    for active_b in (300, 350, 400):
        for ratio in (2800 / 104, 30, 1600 / 49):
            total_t = active_b * ratio / 1000
            model_rows.append(dict(active_parameters_b=active_b,
                total_active_ratio=ratio, total_parameters_t=total_t,
                raw_4bit_weight_tb=total_t / 2,
                nvl72_hbm_after_raw_4bit_weights_tb=13.4 - total_t / 2,
                project_flops=project_flops, pretraining_share=pretrain_share,
                cumulative_training_tokens_t=pretrain_share * project_flops / (6 * active_b * 1e9) / 1e12))
    active_b = ASSUMPTIONS['representative_active_parameters_b']
    tokens_t = pretrain_share * project_flops / (6 * active_b * 1e9) / 1e12
    total_t = tokens_t / ASSUMPTIONS['tokens_per_total_parameter']
    raw_tb = total_t * ASSUMPTIONS['serving_weight_bits'] / 8
    central = dict(active_parameters_b=active_b, total_active_ratio=total_t*1000/active_b,
        total_parameters_t=total_t, raw_4bit_weight_tb=raw_tb,
        nvl72_hbm_after_raw_4bit_weights_tb=ASSUMPTIONS['nvl72_hbm_tb']-raw_tb,
        project_flops=project_flops, pretraining_share=pretrain_share,
        cumulative_training_tokens_t=tokens_t)
    model_rows.append(central)
    save_csv("astra-model-scenarios.csv", model_rows)

    old_v = bill("kimi_k3", 8000, 1000, .9) * 1000
    new_v = bill("kimi_k3", 100_000, 1000, .9) * 1000
    summary = dict(as_of="2026-09-10", sources=SOURCES, standard_prices_usd_per_million=PRICES,
        openai_long_context_threshold_input_tokens=272000,
        assumptions={"kappa_definition": "(Astra/Kimi service-resource cost shares) times (Kimi/Astra cost per active parameter under matched workload); local linear approximation",
            "kappa_is_measured": False,
            "hourly_costs_are_author_scenarios": True,
            "100k_workload_astra_throughput_is_measured": False,
            "same_numerical_token_counts_assumed": True,
            "effective_duty": .7, "target_resource_surplus_fraction": .3},
        request_bills=bills, throughput_thresholds=thresholds,
        active_parameter_sensitivity=kappa_rows, model_scenarios=model_rows,
        generation_price_calibration=generation_rows,
        agent_vs_8_to_1_at_90pct_cache={"old_revenue_per_million_output": old_v,
            "new_revenue_per_million_output": new_v, "revenue_ratio": new_v / old_v,
            "old_fresh_forward_tokens_per_output": 1 + 8 * .1,
            "new_fresh_forward_tokens_per_output": 1 + 100 * .1,
            "fresh_forward_work_ratio": (1 + 100 * .1) / (1 + 8 * .1)},
        central_model_scenario=central,
        preferred_author_scenario={"active_parameters_b": [300, 400],
            "total_active_ratio": [2800 / 104, 1600 / 49],
            "architecture_only_total_parameters_t": [300 * (2800 / 104) / 1000, 400 * (1600 / 49) / 1000],
            "joint_training_constraint_file": "training-constraints.json"},
        memory_examples=[dict(total_parameters_t=t, raw_4bit_weight_tb=t / 2,
            nvl72_hbm_after_raw_4bit_weights_tb=13.4 - t / 2) for t in (8, 10, 12)])
    (DATA / "serving-calibration.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

    assert math.isclose(old_v, 19.56)
    assert math.isclose(new_v, 72)
    assert math.isclose(bill("gpt6_astra", 100_000, 1000, .85), .285)
    assert math.isclose(bill("gpt6_astra", 100_000, 1000, .85, 1), .3225)
    assert math.isclose(bill("gpt6_astra", 300_000, 1000, .85), 1.485)
    assert math.isclose(bill("gpt6_astra", 272_000, 1000, .85), .6892)
    for r in bills:
        if r["model"] == "gpt6_astra":
            sol_bill = bill("gpt56_sol", r["input_tokens"], 1000, r["cache_hit"], r["miss_write_fraction"])
            assert math.isclose(r["bill_usd"] / sol_bill, 2.5)
    assert math.isclose(central["cumulative_training_tokens_t"] / central["total_parameters_t"], ASSUMPTIONS["tokens_per_total_parameter"])
    assert math.isclose(central["cumulative_training_tokens_t"], 296.2285714285714)
    assert math.isclose(bill("gpt55", 100_000, 1000, .85), .1475)
    print(json.dumps({"central": central, "kappa": kappa_rows,
        "thresholds_85pct_target": [r for r in thresholds if r["cache_hit"] == .85 and r["target_resource_surplus_fraction"] == .3],
        "rows": {"bills": len(bills), "thresholds": len(thresholds), "kappa": len(kappa_rows), "models": len(model_rows)}}, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()
