"""Original scalar optimization example with invented timing; not a classifier experiment or Pollux scheduler reproduction.
Run with Python 3.10+; standard library only. Writes results.json in the current directory.
"""
# Identical examples make the averaged gradient independent of batch.
def run(batch):
    weight = 1.0
    learning_rate = 0.25
    steps = 0
    target = 0.0001
    step_seconds = 0.01 + batch / 10000
    while 0.5 * weight * weight > target:
        gradient = sum(weight for _ in range(batch)) / batch
        weight -= learning_rate * gradient
        steps += 1
        assert steps < 100
    return steps, steps * step_seconds, batch / step_seconds

results = {}
for batch in (64, 256, 1024):
    steps, seconds, throughput = run(batch)
    results[batch] = (steps, seconds, throughput)
    print(f"batch={batch}: steps={steps}, "
          f"modeled_seconds={seconds:.3f}, examples/s={throughput:.1f}")
small, large = results[64], results[1024]
assert small[0] == large[0] == 15
assert large[2] > small[2]
assert large[1] > small[1]
print("More examples/s did not mean less time to target.")

# Extended experiment: deterministic parameter sweep and invariants.
import math
batches = [16, 32, 64, 128, 256, 512, 1024, 2048]
noise_levels = [0.0, 0.2, 0.8]
sweep = []
for sigma in noise_levels:
    for batch in batches:
        eta, target = 0.25, 0.0001
        contraction = (1 - eta) ** 2
        variance_term = eta ** 2 * sigma ** 2 / batch
        floor = 0.5 * variance_term / (1 - contraction)
        second_moment, steps = 1.0, 0
        if floor < target:
            while 0.5 * second_moment > target:
                second_moment = contraction * second_moment + variance_term
                steps += 1
                assert steps < 10000
        else:
            steps = None
        step_time = 0.01 + batch / 10000
        sweep.append(dict(sigma=sigma, batch=batch, expected_loss_floor=floor,
                          updates=steps, modeled_seconds=None if steps is None else steps*step_time,
                          examples_per_second=batch/step_time))
assert [r['updates'] for r in sweep if r['sigma']==0] == [15]*len(batches)
assert all(r['updates'] is None for r in sweep if r['sigma']==0.8 and r['batch']<512)
data = dict(batch_sweep=sweep, learning_rate=0.25, target_expected_loss=0.0001,
            timing_model='seconds/update = 0.01 + batch / 10000',
            gradient_model='g = w + independent zero-mean noise; variance of batch mean = sigma^2 / batch')
for sigma in noise_levels:
    feasible=[r for r in sweep if r['sigma']==sigma and r['updates'] is not None]
    best=min(feasible,key=lambda r:r['modeled_seconds'])
    print(f"noise SD={sigma}: best tested batch={best['batch']}, expected-loss target in {best['modeled_seconds']:.4f} modeled s")

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