"""CPU candidate/record audit for the compact contextual RL evidence bundle.

Re-extraction covers fixed raw samples only. Does not replay neural training.
"""
import argparse
from collections import defaultdict
import hashlib
import json
from pathlib import Path
import numpy as np


def read(path):
    return json.loads(path.read_text())


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def table(path):
    data = read(path)
    return [dict(zip(data['columns'], row, strict=True)) for row in data['rows']]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--bundle', type=Path, default=Path(__file__).resolve().parents[1])
    ap.add_argument('--data', type=Path, required=True, help='Supplied baseline code/data directory')
    ap.add_argument('--output', type=Path, required=True, help='Fresh directory outside the bundle')
    args = ap.parse_args()
    root = args.bundle.resolve()
    out = args.output.resolve()
    assert not out.is_relative_to(root), 'Do not alter the review bundle.'
    assert not out.exists()
    manifest = read(root / 'MANIFEST.json')
    for name, record in manifest.items():
        p = root / name
        assert p.resolve().is_relative_to(root) and p.is_file() and not p.is_symlink()
        assert p.stat().st_size == record['bytes'] and digest(p) == record['sha256'], name
    c = read(root / 'notes/contextual_rl_contract.json')
    originals = read(root / 'ORIGINAL_INPUTS_AND_RAW_SHA256.json')
    assert digest(root / 'notes/contextual_rl_contract.json') == originals['notes/contextual_rl_contract.json']
    for name, expected in c['input_sha256'].items():
        assert originals[name] == expected
        if (root / name).is_file(): assert digest(root / name) == expected
    for name in ['math_train.json', 'tokenizer-base.json']:
        assert digest(args.data / name) == c['input_sha256']['code/data/' + name]
    data = read(args.data / 'math_train.json')
    split = read(root / c['split'])
    train = table(root / 'evidence/training.json')
    native = table(root / 'evidence/native.json')
    group_records = table(root / 'evidence/groups.json')
    assert len(train) == 8192 and len(native) == 2560 and len(group_records) == 1024
    result = read(root / 'analysis/result.json')
    assert result['status'] == 'COMPLETE_AUDITED'
    bygroup = defaultdict(list)
    for r in train:
        assert r['arm'] in c['arms'] and 1 <= r['step'] <= 64 and 0 <= r['position'] < 4
        assert r['index'] == split['train_indices'][(r['step']-1)*4+r['position']]
        assert 1 <= r['length'] <= 512 and r['eos'] in (0,1) and r['reward'] in (0,1)
        bygroup[r['arm'],r['step'],r['position']].append(r)
    active_batches = set()
    seen = set()
    for g in group_records:
        key = g['arm'],g['step'],g['position']
        assert key not in seen
        seen.add(key)
        rr = sorted(bygroup[key], key=lambda r:r['sample'])
        assert [r['sample'] for r in rr] == list(range(8))
        rewards = np.array([r['reward'] for r in rr], dtype=np.float32)
        adv = (rewards-rewards.mean())/(rewards.std(ddof=1)+1e-4)*np.array([r['eos'] for r in rr])
        assert np.allclose(adv, g['advantages'], atol=3e-7, rtol=3e-7)
        donor = next((i for i,r in enumerate(rr) if r['reward'] and r['eos']), None)
        receivers = [] if donor is None else [i for i,r in enumerate(rr) if not r['reward'] and r['eos']]
        assert g['donor'] == donor and g['receivers'] == receivers
        dense = g['arm'] in ('C','E')
        assert g['dense_tokens'] == (sum(rr[i]['length'] for i in receivers) if dense else 0)
        active = [i for i in range(8) if adv[i] != 0 or (dense and i in receivers) or (g['arm']=='H' and i==donor)]
        assert len(active) == g['backward_calls']
        if active: active_batches.add((g['arm'],g['step']))
    assert set(bygroup) == seen
    for arm in c['arms']:
        rr = [r for r in train if r['arm']==arm]
        expected = result['training'][arm]
        assert len(rr)==2048 and sum(r['reward'] for r in rr)==expected['correct']
        assert sum(r['eos'] for r in rr)==expected['eos']
        assert sum(a==arm for a,s in active_batches)==expected['optimizer_updates']
    bynative = {}
    values = {}
    for arm in ['W']+c['arms']:
        rr = [r for r in native if r['arm']==arm]
        byid = {r['index']:r for r in rr}
        assert len(rr)==len(byid)==512 and set(byid)==set(split['development_indices'])
        assert all(1<=r['length']<=2048 and r['correct'] in (0,1) for r in rr)
        bynative.update({(arm,i):r for i,r in byid.items()})
        values[arm] = np.array([byid[i]['correct'] for i in split['development_indices']])
        assert int(values[arm].sum()) == result['counts'][arm]
    contrasts = {}
    for other in ['W','G','E','H']:
        delta = values['C']-values[other]
        rng = np.random.default_rng(c['bootstrap_seed'])
        bs = np.concatenate([delta[rng.integers(512,size=(1000,512))].mean(1) for _ in range(100)])
        got = {'difference':float(delta.mean()), 'net_correct':int(delta.sum()),
               'ci95':np.quantile(bs,[.025,.975]).tolist(),
               'gains':int((delta==1).sum()), 'losses':int((delta==-1).sum())}
        assert got == result['contrasts']['C-'+other]
        contrasts['C-'+other] = got
    gate = contrasts['C-W']['net_correct']>=16 and contrasts['C-W']['ci95'][0]>0 and contrasts['C-G']['ci95'][0]>0
    assert bool(gate) == result['development_gate']
    from run_baseline import Qwen3Tokenizer, extract_final_candidate, render_prompt
    from reasoning_from_scratch.ch03 import normalize_text
    from run_with_isolated_grader import IsolatedGrader
    tok = Qwen3Tokenizer(tokenizer_file_path=args.data / 'tokenizer-base.json')
    lookup = {(r['arm'],r['step'],r['position'],r['sample']):r for r in train}
    snippets = read(root / 'evidence/raw_samples.json')
    assert len(snippets)==26
    expected_samples = {('train',arm,chunk) for arm in c['arms'] for chunk in range(c['chunks'])}
    expected_samples |= {('eval',arm,chunk) for arm in ['W']+c['arms'] for chunk in range(2)}
    assert {(s['kind'],s['arm'],s['chunk']) for s in snippets} == expected_samples
    native_prompt_lengths = {i:len(tok.encode(render_prompt(data[i]['problem']))) for i in split['development_indices']}
    native_order = sorted(split['development_indices'], key=native_prompt_lengths.__getitem__)
    for sample in snippets:
        r = sample['record']
        if sample['kind']=='train':
            assert (r['step'],r['position'],r['sample']) == (sample['chunk']*c['updates_per_chunk']+1,0,0)
            assert r['prompt']==tok.encode(render_prompt(data[r['index']]['problem']))
            assert bool(r['eos'])==(r['tokens'][-1]==tok.eos_token_id)
        else:
            assert r['index']==native_order[sample['chunk']*256]
            assert r['prompt_tokens']==native_prompt_lengths[r['index']]
        assert tok.eos_token_id not in r['tokens'][:-1]
        assert tok.decode(r['tokens'])==r['text']
        pred = extract_final_candidate(r['text'], fallback=None)
        expected = lookup[sample['arm'],r['step'],r['position'],r['sample']] if sample['kind']=='train' else bynative[sample['arm'],r['index']]
        assert r['index']==expected['index']
        if sample['kind']=='train': assert int(r['eos'])==expected['eos']
        assert pred==expected['prediction'] and len(r['tokens'])==expected['length']
    out.mkdir(parents=True)
    grader = IsolatedGrader(normalize_text, out / 'grading.jsonl')
    try:
        for rr, field in [(train,'reward'), (native,'correct')]:
            for r in rr:
                got = bool(r['prediction']) and grader.grade(r['prediction'], data[r['index']]['answer'])
                assert bool(got)==bool(r[field]), (r['arm'],r['index'],field)
    finally:
        grader.close()
    output = {'status':'PASS', 'training_candidates':len(train), 'native_candidates':len(native),
              'raw_sample_reextractions':len(snippets), 'counts':result['counts'],
              'contrasts':contrasts, 'development_gate':bool(gate),
              'scope':'All extracted candidates regraded; fixed26 samples re-extracted. Full raw re-extraction was author-audited; omitted raw responses and neural optimizer execution are not independently checked by this compact audit.'}
    (out / 'result.json').write_text(json.dumps(output,indent=2))
    print(json.dumps(output,indent=2))


if __name__=='__main__':
    main()
