"""Original CPU latency-budget model. Prefill emits token one; transfer precedes token two. Uses maximum gap, not DistServe's measured goodput.
Run with Python 3.10+; standard library only. Writes results.json in the current directory.
"""
prefill_ms = 40.0
decode_ms = 10.0
first_token_limit_ms = 80.0
max_gap_limit_ms = 20.0
effective_bytes_per_second = 1024 ** 3

def evaluate(cache_mib):
    cache_bytes = cache_mib * 1024 ** 2
    transfer_ms = 1000 * cache_bytes / effective_bytes_per_second
    first_token_ms = prefill_ms
    handoff_gap_ms = transfer_ms + decode_ms
    steady_gap_ms = decode_ms
    compliant = (first_token_ms <= first_token_limit_ms
                 and handoff_gap_ms <= max_gap_limit_ms
                 and steady_gap_ms <= max_gap_limit_ms)
    return transfer_ms, handoff_gap_ms, compliant

for cache_mib in (4, 16):
    transfer, gap, compliant = evaluate(cache_mib)
    print(f"cache={cache_mib} MiB: transfer={transfer:.3f} ms, "
          f"handoff_gap={gap:.3f} ms, compliant={compliant}")
assert evaluate(4)[2] is True
assert evaluate(16)[2] is False
print(f"transfer_budget={max_gap_limit_ms - decode_ms:.1f} ms")

# Extended experiment: deterministic parameter sweep and invariants.
cache_sizes=[1,2,4,8,16,32,64];rates=[.25,.5,1.,2.,4.,8.]
sweep=[]
for rate in rates:
    for cache in cache_sizes:
        transfer=1000*cache/(1024*rate)
        gap=10+transfer
        sweep.append(dict(cache_mib=cache,rate_gib_per_s=rate,transfer_ms=transfer,
                          handoff_gap_ms=gap,passes_max_gap=gap<=20))
responses=[]
for outputs in [2,3,4,8,16,32,64,128]:
    transfer=15.625
    avg=10+transfer/(outputs-1)
    responses.append(dict(output_tokens=outputs,average_gap_ms=avg,max_gap_ms=25.625,
                          average_passes=avg<=20,maximum_passes=False))
assert responses[0]['average_gap_ms']==responses[0]['max_gap_ms']
assert responses[1]['average_passes']
data=dict(cache_bandwidth_sweep=sweep,response_length_sweep=responses,
          prefill_ms=40,decode_ms=10,first_token_limit_ms=80,max_gap_limit_ms=20)
print('At 1 GiB/s, maximum passing cache = 10.24 MiB under the 10 ms transfer budget.')
print('16 MiB handoff: average gap passes with >=3 total output tokens, but the maximum gap remains 25.625 ms.')

import json
from pathlib import Path
Path("results.json").write_text(json.dumps(data,indent=2,allow_nan=False)+"\n",encoding="utf-8")