"""Read-only dataset audit and reproducible episode split; does not train."""
import hashlib
import json
from pathlib import Path
import numpy as np
import pyarrow.parquet as pq

ROOT = Path('/home/eai/Downloads/gpt6_b1k/task00_turning_on_radio/datasets/task00_success_v1')
OUT = Path(__file__).resolve().parent
SEED = 20260921
HORIZON = 30
rows = [json.loads(s) for s in (ROOT/'generation_manifest.jsonl').read_text().splitlines()]
assert len(rows) == 259 and len({r['instance_id'] for r in rows}) == 259
assert all(r['success'] for r in rows)
# Consecutive groups of ten in duration order; hold out one instance per group.
ordered = sorted(rows, key=lambda r: (r['steps'], r['instance_id']))
val = set()
for start in range(0, len(ordered), 10):
    group = ordered[start:start+10]
    chosen = min(group, key=lambda r: hashlib.sha256(f"{SEED}:{r['instance_id']}".encode()).hexdigest())
    val.add(chosen['episode_index'])
train = {r['episode_index'] for r in rows} - val
assert len(val)==26 and len(train)==233 and not train & val
metadata = []
for p in sorted(ROOT.glob('meta/episodes/**/*.parquet')):
    metadata.extend(pq.read_table(p).to_pylist())
assert len(metadata)==259
by_id = {r['episode_index']:r for r in rows}
video_paths=set()
for e in metadata:
    assert e['length']==by_id[e['episode_index']]['steps']
    assert e['dataset_to_index']-e['dataset_from_index']==e['length']
    for key in ['observation.rgb.zed_link_camera_0','observation.rgb.left_realsense_link_camera_0','observation.rgb.right_realsense_link_camera_0']:
        path = ROOT / f"videos/{key}/chunk-{e[f'videos/{key}/chunk_index']:03d}/file-{e[f'videos/{key}/file_index']:03d}.mp4"
        assert path.is_file(), path
        video_paths.add(path)
actions=[]; states=[]; frame_ids=[]
for p in sorted(ROOT.glob('data/**/*.parquet')):
    t=pq.read_table(p,columns=['action','observation.state','episode_index'])
    actions.extend(t['action'].to_pylist());states.extend(t['observation.state'].to_pylist());frame_ids.extend(t['episode_index'].to_pylist())
a=np.asarray(actions);s=np.asarray(states)
assert a.shape==(148475,23) and s.shape==(148475,61)
assert np.isfinite(a).all() and np.isfinite(s).all()
counts=dict(zip(*np.unique(frame_ids,return_counts=True)))
assert all(counts[r['episode_index']]==r['steps'] for r in rows)
result={
 'seed':SEED,'source_root':str(ROOT),'selection':'success-only; group by unique instance_id',
 'method':'sort by (steps, instance_id); blocks of 10; one validation instance per block chosen by lowest sha256(seed:instance_id)',
 'action_horizon':HORIZON,'valid_window_rule':'start in [0, episode_length - horizon], inclusive; no boundary padding',
 'manifest_sha256':hashlib.sha256((ROOT/'generation_manifest.jsonl').read_bytes()).hexdigest(),
 'splits':{},
}
for name,ids in [('train',train),('validation',val)]:
    selected=sorted((r for r in rows if r['episode_index'] in ids),key=lambda r:r['episode_index'])
    result['splits'][name]={'episodes':len(selected),'frames':sum(r['steps'] for r in selected),'valid_action_windows':sum(max(0,r['steps']-HORIZON+1) for r in selected),'episode_ids':[r['episode_index'] for r in selected],'instance_ids':[r['instance_id'] for r in selected],'episode_lengths':[r['steps'] for r in selected]}
(OUT/'split_manifest.json').write_text(json.dumps(result,indent=2)+'\n')
files=[]
for p in sorted(ROOT.rglob('*')):
    if p.is_file(): files.append({'path':str(p.relative_to(ROOT)),'bytes':p.stat().st_size,'sha256':hashlib.sha256(p.read_bytes()).hexdigest()})
(OUT/'dataset_sha256.json').write_text(json.dumps(files,indent=2)+'\n')
audit={'episodes':259,'unique_instances':259,'frames':len(a),'raw_state_dim':61,'raw_action_dim':23,'all_numeric_rows_finite':True,'rgb_video_files_exist':len(video_paths),'full_video_decode':'not performed in planning phase','source_files':len(files),'source_bytes':sum(r['bytes'] for r in files),'split_counts':{k:{x:v[x] for x in ['episodes','frames','valid_action_windows']} for k,v in result['splits'].items()}}
(OUT/'dataset_audit.json').write_text(json.dumps(audit,indent=2)+'\n')
print(json.dumps(audit,indent=2))