"""Compare scalar and bipolar-binary clocks with identical seeded writes."""
import argparse
import random
import torch
from beamboot.beam_probe import BeamAttentionProbe, calibrated_gain
from beamboot.binary_time_probe import BinaryTimeProbe


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--horizons', nargs='+', type=int, default=[128, 512, 1024])
    parser.add_argument('--seed', type=int, default=417)
    parser.add_argument('--dtype', choices=['float32', 'float16', 'bfloat16'], default='bfloat16')
    args = parser.parse_args()
    dtype = getattr(torch, args.dtype)
    for horizon in args.horizons:
        for name, mem in [('scalar', BeamAttentionProbe(horizon, dtype, calibrated_gain(horizon))),
                          ('binary', BinaryTimeProbe(horizon, dtype))]:
            rng = random.Random(args.seed)
            errors = 0
            for _ in range(horizon):
                a, v = rng.randrange(16), rng.randrange(16)
                mem.write(a, v)
                errors += mem.read(a) != v
            print(f'{name} {dtype} horizon={horizon} immediate_read_errors={errors}', flush=True)

if __name__ == '__main__': main()