"""Run a small, deterministic backend/precision audit on CPU.

This is an engineering smoke test, not a statistically general benchmark.
"""
import argparse
import json
import platform
import random
from pathlib import Path
import torch
import transformers
from beamboot.hf_binary_memory import HFBinaryMemory
from beamboot.causal_lm import CausalMachine
from beamboot.large_cpu import LargeCPU, LargeRuntime


def attention_case(dtype, backend, seed=417, horizon=128):
    m = HFBinaryMemory(horizon, dtype)
    m.layer.config._attn_implementation = backend
    rng = random.Random(seed)
    failures = []
    for i in range(horizon):
        address, value = rng.randrange(16), rng.randrange(16)
        m.write(address, value)
        actual = m.read(address)
        if actual != value:
            failures.append({'write': i+1, 'address': address, 'expected': value, 'actual': actual})
    return {'dtype': str(dtype), 'attention': backend, 'horizon': horizon,
            'seed': seed, 'errors': len(failures), 'failures': failures}


def machine_case(dtype, backend):
    program = [('LDI', 7), ('STORE', 3), ('LDI', 2), ('ADDM', 3), ('OUT', 0), ('HALT', 0)]
    # Instantiate in safe fp32, then deliberately override for the negative
    # control. The public runtime rejects BF16+eager because it is known bad.
    machine = CausalMachine(program)
    machine.model.to(dtype)
    machine.model.config._attn_implementation = backend
    machine._reset(128)
    ref = LargeRuntime(LargeCPU(program))
    failures = []
    for i in range(len(program)):
        actual, expected = machine.step(), ref.step()
        if actual != expected:
            failures.append({'step': i+1, 'expected_acc': expected.acc, 'actual_acc': actual.acc})
        if actual.halt: break
    return {'dtype': str(dtype), 'attention': backend, 'steps': i+1,
            'mismatches': len(failures), 'failures': failures}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', default='results/precision_smoke.json')
    args = parser.parse_args()
    cases = [(torch.float32, 'eager'), (torch.bfloat16, 'eager'), (torch.bfloat16, 'sdpa')]
    results = {'scope': 'CPU smoke test; constructed weights, not pretrained Beam; one deterministic seed',
               'python': platform.python_version(), 'torch': torch.__version__, 'transformers': transformers.__version__,
               'attention_read': [], 'machine': []}
    for dtype, backend in cases:
        results['attention_read'].append(attention_case(dtype, backend))
        print('attention', dtype, backend, results['attention_read'][-1]['errors'], flush=True)
        results['machine'].append(machine_case(dtype, backend))
        print('machine', dtype, backend, results['machine'][-1]['mismatches'], flush=True)
    path = Path(args.output)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(results, indent=2)+'\n')
    print('saved', path, flush=True)

if __name__ == '__main__': main()