"""Load exact Shared-FM comparison inputs. Requires Python3 and numpy only.

python load_inputs.py --case case15 --chunk 0 --output input.npz
python load_inputs.py --case case02 --history current_repeat --variant full_stop --seed 314
Use --base /path/to/downloaded/site for offline use. Chunk indices are zero-based.
"""
import argparse, hashlib, io, json, urllib.request
from pathlib import Path
import numpy as np

DEFAULT_BASE='https://vpp2-sharedfm-eval.pages.dev/'

def read(base,relative):
    if base.startswith(('https://','http://')):
        request=urllib.request.Request(base.rstrip('/')+'/'+relative,headers={'User-Agent':'Mozilla/5.0 SharedFM-Reproduction/1.0'})
        with urllib.request.urlopen(request,timeout=120) as r:return r.read()
    return (Path(base)/relative).read_bytes()

def load_asset(base,record,cache):
    p=Path(cache)/record['sha256'];p.parent.mkdir(parents=True,exist_ok=True)
    if not p.exists():p.write_bytes(read(base,record['path']))
    raw=p.read_bytes()
    assert len(raw)==record['bytes'] and hashlib.sha256(raw).hexdigest()==record['sha256'],'Input asset SHA mismatch'
    with np.load(io.BytesIO(raw),allow_pickle=False) as z:return {k:z[k] for k in z.files}

def load_case(case_id,chunk=0,base=DEFAULT_BASE,history='true_history',variant=None,seed=None,cache='.sharedfm-input-cache'):
    """Returns GT-conditioned inputs/targets; never use later GT for autoregression.

    For intervention tests set variant (e.g. full_original/full_stop) explicitly.
    A normal rollout seed is case-specific; diagnostic seeds are42/314.
    The returned target_future_rgb and target_action are scoring-only.
    """
    c=json.loads(read(base,f'reproduction/{case_id}/case.json'));assert 0<=chunk<8
    x=c['chunks'][chunk];state=load_asset(base,c['assets']['state_arrays'],cache)
    rawparts=[load_asset(base,r,cache) for r in c['assets']['frame_parts']]
    frames=np.concatenate([r['frames'] for r in rawparts]);ticks=np.concatenate([r['indices'] for r in rawparts]);lookup={int(i):f for i,f in zip(ticks,frames)}
    assert hashlib.sha256(frames.tobytes()).hexdigest()==c['frame_pixels_sha256']
    ids=list(x['history_ticks']);prompt=x['prompt'];actual_seed=x['seed'] if seed is None else seed
    if variant is not None:
        assert chunk==0,'Intervention records are first-chunk only'
        actual_seed=42 if seed is None else seed
        matching=[d for d in c['diagnostics'] if d['history']==history and d['variant']==variant and d['seed']==actual_seed]
        if len(matching)!=1:raise ValueError('This exact case/history/variant/seed was not evaluated')
        prompt=matching[0]['prompt']
    elif history!='true_history':raise ValueError('Set variant explicitly for a recorded history intervention')
    if history=='current_repeat':ids=[ids[0]]+[ids[-1]]*8
    elif history=='reverse_past':ids=[ids[0]]+list(reversed(ids[1:-1]))+[ids[-1]]
    elif history!='true_history':raise ValueError(history)
    rgb=np.stack([lookup[i] for i in ids]);j=chunk*32;raw=state['states'][j];valid=state['valid'][j]
    proprio=np.where(valid,(raw-state['mean'])/state['std'],0).astype('float32')
    return dict(condition_rgb=rgb,condition_video=rgb.transpose(3,0,1,2).astype('float32')/127.5-1,
        history_ticks=np.asarray(ids),proprio_raw=raw,proprio_normalized=proprio,proprio_valid=valid,
        prompt=np.asarray(prompt),seed=np.asarray(actual_seed),mean=state['mean'],std=state['std'],
        target_action=state['states'][j+1:j+33],target_valid=state['valid'][j+1:j+33],
        target_future_rgb=np.stack([lookup[i] for i in x['future_video_ticks']]))

def generated_feedback_history(requested_ticks,initial_tick,real_frame_lookup,generated_frame_lookup):
    """Reproduce previous-sample-hold RGB conditioning using OWN predictions.

    generated_frame_lookup maps25Hz episode ticks to predicted uint8 RGB images.
    Initial/earlier real frames are allowed. Future GT is never read here.
    """
    result=[]
    for tick in requested_ticks:
        if tick<=initial_tick:result.append(real_frame_lookup[tick])
        else:
            keys=[k for k in generated_frame_lookup if k<=tick]
            result.append(generated_frame_lookup[max(keys)] if keys else real_frame_lookup[initial_tick])
    return np.stack(result)

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('--base',default=DEFAULT_BASE);p.add_argument('--case',required=True);p.add_argument('--chunk',type=int,default=0)
    p.add_argument('--history',default='true_history',choices=['true_history','current_repeat','reverse_past']);p.add_argument('--variant');p.add_argument('--seed',type=int)
    p.add_argument('--cache',default='.sharedfm-input-cache');p.add_argument('--output',default='input.npz');a=p.parse_args()
    x=load_case(a.case,a.chunk,a.base,a.history,a.variant,a.seed,a.cache);np.savez_compressed(a.output,**x)
    print(json.dumps({'case':a.case,'chunk':a.chunk,'history_ticks':x['history_ticks'].tolist(),'seed':int(x['seed']),'condition_shape':list(x['condition_video'].shape),'output':a.output}))