"""Independent full chunks from one model, using immutable shared observations."""
import argparse
import json
import logging
import time
import jax
import numpy as np
from serve_checkpoint import load_policy
from common import ROOT, EVAL, local_only, conditions, sha, write_json, save_npz, progress

def main(variant):
    local_only()
    manifest = json.loads((ROOT/'observations/manifest.json').read_text())
    for ident in conditions():
        assert sha(ROOT/'observations'/f'{ident}.npz') == manifest[ident]['sha256']
    start=time.monotonic()
    wrapper=load_policy(EVAL.parent/'b1k_sft'/variant)
    base=wrapper.base_policy
    tap={}; transform=base._output_transform
    def save_raw(outputs):
        tap['raw']=np.asarray(outputs['actions'],dtype=np.float32).copy()
        return transform(outputs)
    base._output_transform=save_raw
    outdir=ROOT/'samples'/variant; outdir.mkdir(parents=True,exist_ok=True)
    norm=EVAL.parent/'b1k_sft'/variant/'checkpoint-14999/assets/IliaLarchenko/behavior_224_rgb/norm_stats.json'
    norm_sha=sha(norm)
    elapsed=[]
    for ci, ident in enumerate(conditions()):
        path=ROOT/'observations'/f'{ident}.npz'
        meta=manifest[ident]
        with np.load(path) as data:
            obs={k:data[k] for k in data.files}
        wrapper.current_stage=meta['stage']
        batch=wrapper.prepare_batch_for_pi_behavior(wrapper.process_obs(obs))
        assert int(batch['subtask_state']) == meta['stage']
        def draw(seed):
            base._rng=jax.random.key(int(seed))
            t=time.monotonic()
            result=base.infer(batch, initial_actions=None)
            a=np.asarray(result['actions'],dtype=np.float32)
            raw=tap['raw'].copy()
            assert a.shape==(30,23) and raw.shape==(30,32)
            assert np.isfinite(a).all() and np.isfinite(raw).all()
            return a,raw,time.monotonic()-t,np.asarray(result['subtask_logits'])
        if ci==0:
            a,r,compile_s,logits=draw(90000)
            b,s,_,_=draw(90000)
            assert np.array_equal(a,b) and np.array_equal(r,s),'Same-seed repeatability failed'
            check=[draw(90001+i) for i in range(32)]
            preflight=dict(same_seed_exact=True,different_seed_physical_maxdiff=float(np.max(np.abs(check[0][0]-check[1][0]))),
                different_seed_raw_maxdiff=float(np.max(np.abs(check[0][1]-check[1][1]))),
                median_seconds=float(np.median([x[2] for x in check])),first_inference_seconds=compile_s,
                load_and_preflight_seconds=time.monotonic()-start, normalizer_sha256=norm_sha,
                initial_actions=None,num_flow_steps=20,wrapper_act_used=False,stage_fixed=True)
            write_json(outdir/'preflight.json',preflight)
            print('PREFLIGHT',json.dumps(preflight),flush=True)
        final=outdir/f'{ident}.npz'; partial=outdir/f'{ident}.partial.npz'
        actions=[]; raw_actions=[]; durations=[]
        existing=final if final.exists() else partial
        if existing.exists():
            with np.load(existing) as old:
                assert str(old['observation_sha256'])==meta['sha256'] and str(old['normalizer_sha256'])==norm_sha
                n=len(old['seeds']); assert np.array_equal(old['seeds'],np.arange(100000,100000+n))
                actions=list(old['actions']);raw_actions=list(old['raw_actions']);durations=list(old['seconds'])
            if final.exists():
                assert n==512
                write_json(final.with_suffix('.json'),dict(variant=variant,condition=ident,samples=512,
                    sha256=sha(final),observation_sha256=meta['sha256'],normalizer_sha256=norm_sha,
                    num_flow_steps=20,initial_actions=None,physical_shape=[512,30,23],raw_shape=[512,30,32]))
                continue
        for j in range(len(actions),512):
            a,r,d,logits=draw(100000+j)
            actions.append(a);raw_actions.append(r);durations.append(d);elapsed.append(d)
            if (j+1)%32==0:
                save_npz(partial,actions=np.stack(actions),raw_actions=np.stack(raw_actions),
                    seeds=np.arange(100000,100001+j),seconds=np.array(durations),
                    observation_sha256=np.array(meta['sha256']),normalizer_sha256=np.array(norm_sha),
                    stage=np.array(meta['stage']),subtask_logits=logits)
                progress('sampling_progress.json',variant=variant,condition=ident,condition_index=ci,
                         in_condition=j+1,model_samples=ci*512+j+1,model_target=10240,
                         median_seconds=float(np.median(elapsed[-128:])))
        partial.replace(final)
        write_json(final.with_suffix('.json'),dict(variant=variant,condition=ident,samples=512,
            sha256=sha(final),observation_sha256=meta['sha256'],normalizer_sha256=norm_sha,
            num_flow_steps=20,initial_actions=None,physical_shape=[512,30,23],raw_shape=[512,30,32]))
        print('SAMPLED',variant,ident,512,flush=True)
    progress(f'logs/{variant}_sampling_complete.json',state='complete',samples=10240)

if __name__=='__main__':
    logging.basicConfig(level=logging.INFO)
    p=argparse.ArgumentParser();p.add_argument('variant');main(p.parse_args().variant)