"""YOR-only waypoint filtering and lossless interleaved temporal subsampling."""
import numpy as np
from datasets.core import TrajectoryDataset

DEFAULT_THRESHOLDS = dict(left_xyz_m=.005,right_xyz_m=.005,left_rotation_deg=3.,
    right_rotation_deg=3.,left_hand_rad=.08,right_hand_rad=.08,lift_m=.005,
    base_xy_m=.01,base_yaw_rad=.03)
GROUPS = dict(left_xyz_m=slice(4,7),right_xyz_m=slice(11,14),
    left_hand_rad=slice(14,30),right_hand_rad=slice(34,50),lift_m=slice(54,55),
    base_xy_m=slice(55,57),base_yaw_rad=slice(57,58))

def validate_thresholds(thresholds):
    unknown=set(thresholds)-set(DEFAULT_THRESHOLDS)
    if unknown: raise ValueError(f'Unknown motion thresholds: {unknown}')
    for key,t in thresholds.items():
        if t is None: continue
        t=np.asarray(t,float)
        width=GROUPS[key].stop-GROUPS[key].start if key in GROUPS else 1
        if t.ndim>1 or t.size not in (1,width) or not np.isfinite(t).all() or (t<=0).any():
            raise ValueError(f'{key}: use null to disable, or positive scalar/{width}-vector')

def motion_scores(actions,reference,thresholds):
    """Largest per-axis normalized change. Quaternion change is sign-invariant."""
    a=np.asarray(actions,float);b=np.asarray(reference,float);scores=[]
    for key,t in thresholds.items():
        if t is None: continue
        if key in GROUPS:
            d=np.abs(a[...,GROUPS[key]]-b[GROUPS[key]])
            if key=='base_yaw_rad':d=np.abs((a[...,57:58]-b[57:58]+np.pi)%(2*np.pi)-np.pi)
            scores.append(np.max(d/np.asarray(t),axis=-1))
        else:
            sl=slice(0,4) if key=='left_rotation_deg' else slice(7,11)
            q=a[...,sl];r=b[sl];den=np.linalg.norm(q,axis=-1)*np.linalg.norm(r)
            if (den<1e-9).any():raise ValueError('Invalid quaternion in motion selection')
            angle=np.degrees(2*np.arccos(np.clip(np.abs(np.sum(q*r,axis=-1))/den,0,1)))
            scores.append(angle/float(np.asarray(t).reshape(-1)[0]))
    return np.max(np.stack(scores),axis=0) if scores else np.zeros(a.shape[:-1])

def select_waypoints(actions,enabled=False,thresholds=None,max_gap_frames=30):
    """Keep first/last, threshold crossings vs last kept, and optional hold anchors.

    max_gap_frames=0 removes the hold-anchor rule. It does not enforce minimum
    motion on anchors or endpoints. Source data and source indices are unchanged.
    """
    a=np.asarray(actions);n=len(a)
    if a.ndim!=2 or a.shape[1]!=58 or not np.isfinite(a).all():raise ValueError('Expected finite [T,58] actions')
    if int(max_gap_frames)!=max_gap_frames or max_gap_frames<0:raise ValueError('max_gap_frames must be a nonnegative integer')
    t=dict(DEFAULT_THRESHOLDS if thresholds is None else thresholds);validate_thresholds(t)
    if not enabled or n<2:return np.arange(n,dtype=np.int64)
    kept=[0]
    for i in range(1,n):
        if motion_scores(a[i],a[kept[-1]],t)>=1 or (max_gap_frames and i-kept[-1]>=max_gap_frames):kept.append(i)
    if kept[-1]!=n-1:kept.append(n-1)
    return np.asarray(kept,dtype=np.int64)

def phase_indices(indices,frame_skip):
    if int(frame_skip)!=frame_skip or frame_skip<1:raise ValueError('frame_skip must be a positive integer')
    return [np.asarray(indices)[p::int(frame_skip)] for p in range(min(int(frame_skip),len(indices)))]

class YorTemporalDataset(TrajectoryDataset):
    """Wrap AFTER episode splitting and embedding; phases cannot leak across splits."""
    def __init__(self,dataset,frame_skip=1,motion_filter=None):
        self.dataset=dataset;self.frame_skip=frame_skip;self.motion_filter=dict(motion_filter or {})
        self.sequences=[]
        for i in range(len(dataset)):
            a=dataset.get_frames(i,range(dataset.get_seq_length(i)))[1]
            if hasattr(a,'detach'):a=a.detach().cpu().numpy()
            ids=select_waypoints(a,**self.motion_filter)
            self.sequences.extend((i,phase,frames) for phase,frames in enumerate(phase_indices(ids,frame_skip)))
    def __len__(self):return len(self.sequences)
    def get_seq_length(self,idx):return len(self.sequences[idx][2])
    def get_frames(self,idx,frames):
        i,phase,ids=self.sequences[idx]
        return self.dataset.get_frames(i,ids[np.asarray(list(frames),dtype=int)].tolist())
    def __getitem__(self,idx):return self.get_frames(idx,range(self.get_seq_length(idx)))
    def get_all_actions(self):
        import torch
        return torch.cat([self[i][1] for i in range(len(self))])


from datasets.core import TrajectorySlicerDataset

class YorTemporalSlicer(TrajectorySlicerDataset):
    """Pad short phases and optionally predict actions after each observation.

    action_offset is measured in sampled sequence rows. With offset 1 the
    current observation at t predicts t+1 onward. Exclude anchors that have no
    next target, but repeat the last target when a chunk reaches the end.
    """
    def __init__(self,dataset,action_offset=0,**kwargs):
        if int(action_offset)!=action_offset or action_offset<0:
            raise ValueError('action_offset must be a nonnegative integer')
        self.action_offset=int(action_offset)
        super().__init__(dataset,**kwargs)
        self.slices=[(i,t) for i in range(len(dataset))
                     for t in range(max(0,dataset.get_seq_length(i)-self.action_offset))
                     if self.pad_seq_length or t+self.action_offset+self.action_window<=dataset.get_seq_length(i)]
    def source_indices(self,idx):
        i,t=self.slices[idx];T=self.dataset.get_seq_length(i)
        obs=np.clip(np.arange(t-self.window+1,t+1),0,T-1)
        start=t if self.vqbet_get_future_action_chunk else t-self.window+1
        act=np.clip(np.arange(start,t+self.action_window)+self.action_offset,self.action_offset,T-1)
        return i,obs,act
    def __getitem__(self,idx):
        i,oi,ai=self.source_indices(idx)
        obs,act,*others=self.dataset[i]
        values=(obs[oi],act[ai],*[other[oi] for other in others])
        return self.transform(values) if self.transform else values