"""Strict ordered numerical gates; both complete CAL arms even if representation fails."""
import argparse
import subprocess
import time
import traceback
from context import ROOT, MODEL, EXP, AG, CFG, check, save, read, budget

SCIENCE = ROOT / 'models/male-cns-lif/.venv/bin/python'
ORACLE = ROOT / 'experiments/phase4ad-author-semantics/environment/.venv/bin/python'


def child(backend, arm, name, repeat=0):
    log = EXP / 'logs' / f'{backend}-{arm}-{name}-r{repeat}.log'
    command = [str(ORACLE if backend == 'oracle' else SCIENCE), '-B', str(MODEL / 'run_case.py'),
               backend, arm, name, '--repeat', str(repeat)]
    print('START', backend, arm, name, repeat, flush=True)
    started = time.perf_counter()
    with log.open('x', encoding='utf-8') as stream:
        result = subprocess.run(command, stdout=stream, stderr=subprocess.STDOUT, timeout=1800, check=False)
    print('FINISH', backend, arm, name, repeat, 'exit', result.returncode,
          'wall', round(time.perf_counter() - started, 2), flush=True)
    if result.returncode:
        raise RuntimeError(f'{backend} {arm}/{name} failed; preserve and inspect {log}')


def gate_case(arm, name, repeats):
    from compare_timing import compare, compare_gpu_repeats
    from calibration import control_identity
    child('gpu', arm, name)
    for repeat in range(1, repeats):
        child('gpu', arm, name, repeat)
    child('oracle', arm, name)
    result = compare(EXP / 'cases' / arm / name, EXP / 'gpu' / arm / (name + '-r0'),
                     EXP / 'oracle' / arm / name, EXP / 'audit' / f'compare-{arm}-{name}.json')
    assert result['pass_all'], f'Closed-loop exact comparison failed: {arm}/{name}'
    for repeat in range(1, repeats):
        report = compare_gpu_repeats(EXP / 'gpu' / arm / (name + '-r0'),
            EXP / 'gpu' / arm / f'{name}-r{repeat}', EXP / 'audit' / f'repeat-{arm}-{name}-r{repeat}.json')
        assert report['pass_all'], f'Exact repeat failed: {arm}/{name}/{repeat}'
    if arm == 'delay0':
        report = control_identity(EXP / 'gpu' / arm / (name + '-r0'), AG / 'gpu' / (name + '-r0'),
                                  EXP / 'cases' / arm / name)
        save(EXP / 'audit' / f'control-{name}.json', report)
        assert report['pass_all'], f'Zero-delay control changed relative to sealed AG: {name}'
    child('component', arm, name)
    return dict(arm=arm, case=name, pass_all=True, gpu_repeats=repeats)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('stage', choices=['fixtures', 'probes', 'calibration'])
    args = parser.parse_args()
    check(); budget()
    receipt = EXP / 'audit' / (args.stage + '-gate.json')
    assert not receipt.exists()
    assert read(EXP / 'origin-audit/summary.json')['pass_all']
    assert read(EXP / 'audit/shift-gate.json')['pass_all']
    arms = ['delay' + str(delay) for delay in CFG['delays_steps']]
    try:
        if args.stage == 'fixtures':
            rows = [gate_case(arm, name, CFG['fixture_repeats']) for arm in arms for name in CFG['fixtures']]
            save(receipt, dict(pass_all=True, cases=rows, scope='Both arms; fixed real-edge induced fixtures only'))
        elif args.stage == 'probes':
            assert read(EXP / 'audit/fixtures-gate.json')['pass_all']
            rows = [gate_case(arm, name, 1) for arm in arms for name in CFG['calibration']['full_oracle_probes']]
            save(receipt, dict(pass_all=True, cases=rows, scope='Both arms; same two original first-CAL fullgraph cases'))
        else:
            assert read(EXP / 'audit/fixtures-gate.json')['pass_all']
            assert read(EXP / 'audit/probes-gate.json')['pass_all']
            from calibration import summarize, control_identity
            from compare_timing import compare_gpu_repeats
            records = read(EXP / 'protocol/calibration-cases.json')
            assert len(records) == 60
            for arm in arms:
                for record in records:
                    if record['name'] not in CFG['calibration']['full_oracle_probes']:
                        child('gpu', arm, record['name'])
                    name = record['name']
                    if arm == 'delay0':
                        control = control_identity(EXP / 'gpu' / arm / (name + '-r0'),
                            AG / 'gpu' / (name + '-r0'), EXP / 'cases' / arm / name)
                        save(EXP / 'audit' / f'cal-control-{name}-r0.json', control)
                        assert control['pass_all'], 'Zero-delay CAL control diverged: ' + name
                    if record['seed'] == CFG['calibration']['seeds'][0]:
                        child('gpu', arm, record['name'], 1)
                        repeat = compare_gpu_repeats(EXP / 'gpu' / arm / (name + '-r0'),
                            EXP / 'gpu' / arm / (name + '-r1'), EXP / 'audit' / f'cal-repeat-{arm}-{name}.json')
                        assert repeat['pass_all'], 'CAL exact repeat diverged: ' + arm + '/' + name
                        if arm == 'delay0':
                            control = control_identity(EXP / 'gpu' / arm / (name + '-r1'),
                                AG / 'gpu' / (name + '-r1'), EXP / 'cases' / arm / name)
                            save(EXP / 'audit' / f'cal-control-{name}-r1.json', control)
                            assert control['pass_all'], 'Zero-delay CAL repeat control diverged: ' + name
            result = summarize()
            save(receipt, dict(pass_all=result['execution_gate_pass'],
                representation_pass_all=result['pass_all'], scope='Complete both fixed CAL arms; not heldout or learning',
                summary='calibration/summary.json'))
            assert result['execution_gate_pass'], 'CAL numerical/control/repeat gate failed'
    except Exception as error:
        save(EXP / 'audit' / ('stopped-' + args.stage + '.json'), dict(stage=args.stage,
            type=type(error).__name__, message=str(error), traceback=traceback.format_exc(),
            later_stages_not_permitted=True, partial_outputs_preserved=True, parameters_changed=False))
        raise
    budget()


if __name__ == '__main__':
    main()