"""Rerun the controlled grid from immutable laws and stored sample streams."""
from pathlib import Path
import argparse,hashlib,json,time
import torch
from reproduce import ROOT,load,verify
from unionmind.decision_sufficiency import study,benchmark

def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--device',default='cuda:0');parser.add_argument('--output',required=True)
    parser.add_argument('--primary-only',action='store_true')
    parser.add_argument('--benchmark',action='store_true')
    a=parser.parse_args();verify()
    assert a.device.startswith('cuda') and torch.cuda.is_available()
    out=Path(a.output).resolve();assert not out.exists(),'Use a new output directory.'
    (out/'confirmation').mkdir(parents=True)
    torch.cuda.set_device(a.device);torch.backends.cuda.matmul.allow_tf32=False
    config=json.loads((ROOT/'protocols/decision_sufficiency_protocol_01.json').read_text())
    config['execution']['device']=a.device
    if a.primary_only:config['estimation']['label_budgets']=[config['primary']['label_budget']]
    protocol=out/'protocol.json';protocol.write_text(json.dumps(config,indent=2)+'\n')
    archive=load('data/controlled/worlds',a.device);target=out/'confirmation/worlds.pt'
    torch.save(archive,target)
    manifest=json.loads((ROOT/'data/controlled/manifest.json').read_text())
    manifest['file_sha256']=hashlib.sha256(target.read_bytes()).hexdigest()
    (out/'confirmation/manifest.json').write_text(json.dumps(manifest,indent=2)+'\n')
    paths=[protocol,target,*list((ROOT/'unionmind/decision_sufficiency').glob('*.py'))]
    (out/'preconfirmation_lock.json').write_text(json.dumps(dict(
        purpose='Reproduction from archived observations, not new confirmation',
        timestamp_unix=time.time(),files={str(p):hashlib.sha256(p.read_bytes()).hexdigest() for p in paths}),indent=2)+'\n')
    for module in [study,benchmark]:module.ROOT=out;module.PROTOCOL=protocol
    study.evaluate('confirmation')
    if a.benchmark:benchmark.run()
    print('Controlled grid reconstructed. Original study results and criteria are unchanged.')
if __name__=='__main__':main()