"""Cached LeRobot v3 batches with explicit episode-safe windows."""
from concurrent.futures import ThreadPoolExecutor
import json
from pathlib import Path
import numpy as np
from prepare import DELTA, normalize

IMAGE_NAMES=('base_0_rgb','left_wrist_0_rgb','right_wrist_0_rgb')

class Dataset:
    def __init__(self,root,workers=32):
        self.root=Path(root);self.cache=self.root/'cache/prepared'
        self.norm=json.loads((self.root/'config/normalization.json').read_text())
        assert json.loads((self.root/'config/preparation_receipt.json').read_text())['status']=='passed'
        for name in ['states_raw','states','actions_raw','episode_ids','frame_ids','tokens','token_masks']:
            setattr(self,name,np.load(self.cache/f'{name}.npy',mmap_mode='r'))
        self.rgb={i:np.load(self.cache/f'rgb_{i:03d}.npy',mmap_mode='r') for i in range(259)}
        self.pool=ThreadPoolExecutor(max_workers=workers)
        self.prefetch=ThreadPoolExecutor(max_workers=1)

    def starts(self,split):return np.load(self.cache/f'{split}_starts.npy')

    def _image(self,i):return self.rgb[int(self.episode_ids[i])][int(self.frame_ids[i])]

    def batch(self,starts,batch_size):
        n=len(starts);idx=np.pad(starts,(0,batch_size-n),mode='edge')
        images=np.stack(list(self.pool.map(self._image,idx)))
        actions=np.array(self.actions_raw[idx[:,None]+np.arange(30)],copy=True)
        assert np.all(self.episode_ids[idx[:,None]+np.arange(30)]==self.episode_ids[idx,None])
        actions[:,:,DELTA]-=self.states_raw[idx][:,None,DELTA]
        actions=np.pad(normalize(actions,self.norm['actions']),((0,0),(0,0),(0,9)))
        batch={'images':{name:images[:,k] for k,name in enumerate(IMAGE_NAMES)},'state':np.asarray(self.states[idx]),
            'tokens':np.asarray(self.tokens[idx]),'token_mask':np.asarray(self.token_masks[idx]),'actions':actions,
            'ids':idx.astype(np.int32)}
        return batch,{'n':n,'episode_ids':np.asarray(self.episode_ids[starts]),'starts':starts}

    def eval_batches(self,split,batch_size):
        starts=self.starts(split)
        chunks=[starts[i:i+batch_size] for i in range(0,len(starts),batch_size)]
        future=self.prefetch.submit(self.batch,chunks[0],batch_size)
        for i in range(len(chunks)):
            result=future.result()
            if i+1<len(chunks):future=self.prefetch.submit(self.batch,chunks[i+1],batch_size)
            yield result

    def train_batches(self,batch_size,seed,start_step=0):
        starts=self.starts('train');per_epoch=len(starts)//batch_size;step=start_step
        while True:
            epoch=step//per_epoch;offset=step%per_epoch
            perm=np.random.default_rng(np.random.SeedSequence([seed,epoch])).permutation(starts)
            for b in range(offset,per_epoch):
                yield self.batch(perm[b*batch_size:(b+1)*batch_size],batch_size)[0]
                step+=1

    def inverse_actions(self,pred,starts):
        n=self.norm['actions'];pred=np.asarray(pred)[...,:23]
        out=(pred+1)/2*(np.array(n['q99'])-np.array(n['q01'])+1e-6)+np.array(n['q01'])
        out[:,:,DELTA]+=self.states_raw[starts][:,None,DELTA]
        return out

    def close(self):self.prefetch.shutdown();self.pool.shutdown()