"""Original flush-style forward/backward CPU schedule model; durations are synthetic and exclude communication and recomputation.
Run with Python 3.10+; standard library only. Writes results.json in the current directory.
"""
def wave_time(stage_times, microbatches):
    available = [0] * len(stage_times)
    for _ in range(microbatches):
        incoming = 0
        for stage, duration in enumerate(stage_times):
            start = max(available[stage], incoming)
            finish = start + duration
            available[stage] = finish
            incoming = finish
    return available[-1]

def schedule(stage_times, microbatches):
    forward = wave_time(stage_times, microbatches)
    backward = wave_time(list(reversed(stage_times)), microbatches)
    elapsed = forward + backward
    useful_work = 2 * microbatches * sum(stage_times)
    utilization = useful_work / (len(stage_times) * elapsed)
    assert 0 < utilization <= 1
    return elapsed, utilization

for label, times in (("balanced", [1, 1, 1, 1]),
                     ("imbalanced", [1, 1, 2, 1])):
    for count in (1, 4, 16):
        elapsed, utilization = schedule(times, count)
        print(f"{label} m={count}: time={elapsed}, "
              f"utilization={utilization:.3f}")
assert schedule([1, 1, 1, 1], 16)[0] == 38
assert schedule([1, 1, 2, 1], 16)[0] == 70

# Extended experiment: deterministic parameter sweep and invariants.
counts=[1,2,4,8,16,32,64,128]
sweep=[]
for overhead in [0.,0.1,0.5,1.0]:
    for count in counts:
        duration=32/count+overhead
        elapsed,utilization=schedule([duration]*4,count)
        expected=2*(count+3)*(32/count+overhead)
        assert abs(elapsed-expected)<1e-9
        sweep.append(dict(microbatches=count,overhead=overhead,duration_per_operation=duration,
                          elapsed=elapsed,utilization=utilization,compute_utilization=256/(4*elapsed)))
def events(times,count):
    available=[0.]*len(times); output=[]
    for micro in range(count):
        incoming=0.
        for stage,duration in enumerate(times):
            start=max(available[stage],incoming);finish=start+duration
            output.append(dict(stage=stage,microbatch=micro,start=start,finish=finish))
            available[stage]=finish;incoming=finish
    return output
data=dict(fixed_work_sweep=sweep, total_compute_per_stage_per_direction=32,
          balanced_forward=events([1,1,1,1],4),imbalanced_forward=events([1,1,2,1],4))
for overhead in [0.,0.1,0.5,1.]:
    best=min((r for r in sweep if r['overhead']==overhead),key=lambda r:r['elapsed'])
    print(f"per-operation overhead={overhead}: best tested microbatches={best['microbatches']}, elapsed={best['elapsed']:.3f}")

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