"""Reproduce bounded Beam-style attention precision sweep on CPU."""
import argparse
import random
import torch
from beamboot.beam_probe import BeamAttentionProbe, calibrated_gain


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


if __name__ == '__main__': main()