import os,json,random
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor
from huggingface_hub import hf_hub_download
os.environ['HF_HUB_DISABLE_PROGRESS_BARS']='1'
files=json.loads(Path('/workspace/nextscene_ops/dataset_files.json').read_text()); fset=set(files)
rng=random.Random(20260930); chosen={}; jobs=[]
for sub in sorted({p.split('/')[0] for p in files if 'images_B/' in p}):
 candidates=[p for p in files if p.startswith(sub+'/images_B/') and Path(p).suffix.lower() in ['.jpg','.png','.webp','.jpeg'] and p.replace('images_B/','images_A/') in fset and str(Path(p).with_suffix('.txt')) in fset]
 selected=rng.sample(candidates,12)
 chosen[sub]=selected
 for p in selected:jobs.extend([p,p.replace('images_B/','images_A/'),str(Path(p).with_suffix('.txt'))])
 if sub.startswith('ds4'):
  smoke=[p for p in candidates if p not in selected][:16]
  chosen['smoke']=smoke
  for p in smoke: jobs.extend([p,p.replace('images_B/','images_A/'),str(Path(p).with_suffix('.txt'))])
Path('/workspace/nextscene_ops/priority_manifest.json').write_text(json.dumps(chosen,indent=2))
def dl(p):
 return hf_hub_download('AdwolfCzar/proxima_cena_grounded_original_dataset',p,repo_type='dataset',local_dir='/workspace/ds')
with ThreadPoolExecutor(16) as ex:list(ex.map(dl,jobs))
print('Priority pairs ready',len(jobs)//3,flush=True)