"""World hand/camera motion, in metres and seconds. No interpolation or smoothing."""
import hashlib
import json
import numpy as np
from scipy.spatial.transform import Rotation
from scipy.stats import spearmanr

VERSION = 'world_motion_v2.0'
PARAMETERS = dict(epsilon_m_s=.001, low_speed_gate_m_s=.02, max_gap_factor=1.5,
    min_duration_s=2., min_coverage=.9, stationary_translation_m_s=.02,
    stationary_rotation_deg_s=3., hand_high_m_s=.3, camera_high_m_s=.15,
    camera_high_deg_s=25., relative_high_m_s=.3,
    threshold_status='provisional screening; sensitivity required, not calibrated sensor accuracy',
    smoothing='none', aggregation='any = maximum side P90; both requires both eligible sides')
PARAMETER_SHA = hashlib.sha256(json.dumps(PARAMETERS,sort_keys=True).encode()).hexdigest()


def valid_pose(pose):
    pose=np.asarray(pose,float)
    finite=np.isfinite(pose).all(axis=(-2,-1));safe=np.where(finite[...,None,None],pose,np.eye(4));r=safe[...,:3,:3]
    return finite & (np.max(abs(safe[...,3,:]-[0,0,0,1]),axis=-1)<1e-6) & (np.max(abs(np.swapaxes(r,-1,-2)@r-np.eye(3)),axis=(-2,-1))<.005) & (abs(np.linalg.det(r)-1)<.005)


def stats(values,dt,total):
    good=np.isfinite(values)
    return dict(median=float(np.median(values[good])) if good.any() else None,
        p90=float(np.quantile(values[good],.9)) if good.any() else None,
        valid_pairs=int(good.sum()),valid_duration_s=float(dt[good].sum()),
        coverage=float(dt[good].sum()/total) if total>0 else 0.)


def compute_motion(a,parameters=None):
    """Return JSON-safe summary and per-interval arrays (NaN = unknown).

    Continuous quantities use forward translation differences and SO(3) logarithm
    at interval start. These are *estimates*, distinct from exact finite-step RHCM.
    Nominal dt / explicit pair_valid controls gaps; non-increasing time is rejected.
    """
    par=dict(PARAMETERS);par.update(parameters or {})
    t=np.asarray(a['timestamps_s'],float);dt=np.diff(t)
    if len(t)<2 or not np.isfinite(t).all() or np.any(dt<=0):raise ValueError('timestamps must be finite and strictly increasing')
    nominal=float(a.get('nominal_dt_s',np.median(dt)))
    interval=(dt<=par['max_gap_factor']*nominal)&np.asarray(a.get('pair_valid',np.ones(len(dt),bool)))
    total=float(t[-1]-t[0]);T=np.asarray(a['T_world_camera'],float)
    cv=np.asarray(a['camera_valid'],bool)&valid_pose(T)
    cp=cv[:-1]&cv[1:]&interval
    R=np.where(cv[:,None,None],T[:,:3,:3],np.eye(3));c=T[:,:3,3]
    omega=Rotation.from_matrix(R[1:]@R[:-1].transpose(0,2,1)).as_rotvec()/dt[:,None]
    vc=np.diff(c,axis=0)/dt[:,None]
    camera_translation=np.linalg.norm(vc,axis=-1);camera_translation[~cp]=np.nan
    camera_rotation=np.rad2deg(np.linalg.norm(omega,axis=-1));camera_rotation[~cp]=np.nan
    p=np.asarray(a['wrist_world_m'],float);hv=np.asarray(a['wrist_world_valid'],bool)&np.isfinite(p).all(-1)
    hp=hv[:-1]&hv[1:]&interval[:,None]&np.asarray(a.get('hand_pair_valid',np.ones((len(dt),2),bool)));rp=hp&cp[:,None]
    vh=np.diff(p,axis=0)/dt[:,None,None]
    u=vc[:,None,:]+np.cross(omega[:,None,:],p[:-1]-c[:-1,None,:])
    q=np.einsum('tji,tsj->tsi',R,p-c[:,None,:])
    arrays=dict(whm_m_s=np.linalg.norm(vh,axis=-1),cmh_m_s=np.linalg.norm(u,axis=-1),
        rhcm_continuous_estimate_m_s=np.linalg.norm(vh-u,axis=-1),
        rhcm_discrete_m_s=np.linalg.norm(np.diff(q,axis=0),axis=-1)/dt[:,None])
    arrays['whm_m_s'][~hp]=np.nan
    for k in ('cmh_m_s','rhcm_continuous_estimate_m_s','rhcm_discrete_m_s'):arrays[k][~rp]=np.nan
    denom=arrays['whm_m_s']+arrays['cmh_m_s']
    arrays['wms']=np.where(denom>=par['low_speed_gate_m_s'],arrays['rhcm_continuous_estimate_m_s']/(denom+par['epsilon_m_s']),np.nan)
    arrays['rham_deg_s']=np.full(hp.shape,np.nan)
    if 'R_world_wrist' in a:
        hr=np.asarray(a['R_world_wrist'],float)
        ht=np.tile(np.eye(4),(len(t),2,1,1));ht[:,:,:3,:3]=hr
        rv=np.asarray(a.get('wrist_rotation_valid',np.zeros(hv.shape,bool)))&valid_pose(ht)
        Q=R[:,None].transpose(0,1,3,2)@np.where(rv[...,None,None],hr,np.eye(3))
        relative=Q[1:]@Q[:-1].transpose(0,1,3,2)
        angles=np.rad2deg(Rotation.from_matrix(relative.reshape(-1,3,3)).magnitude()).reshape(-1,2)/dt[:,None]
        good=rv[:-1]&rv[1:]&cp[:,None]&np.asarray(a.get('wrist_rotation_pair_valid',np.ones((len(dt),2),bool)));angles[~good]=np.nan
        arrays['rham_deg_s']=angles
    summary=dict(metric_version=VERSION,parameters=par,parameter_sha256=hashlib.sha256(json.dumps(par,sort_keys=True).encode()).hexdigest(),
        training_allowed=False,camera_translation=stats(camera_translation,dt,total),camera_rotation=stats(camera_rotation,dt,total),sides={})
    def eligible(s):return s['coverage']>=par['min_coverage'] and s['valid_duration_s']>=par['min_duration_s']
    def classify(s,high,still):
        if not eligible(s):return 'unknown'
        return 'stationary' if s['p90']<=still else 'high' if s['p90']>=high else 'low'
    for j,side in enumerate(('left','right')):
        ss={k:stats(v[:,j],dt,total) for k,v in arrays.items()}
        good=rp[:,j];x=arrays['whm_m_s'][good,j];y=arrays['cmh_m_s'][good,j]
        ss['speed_spearman']=float(spearmanr(x,y).statistic) if len(x)>=3 and np.ptp(x)>1e-10 and np.ptp(y)>1e-10 else None
        ss['hand_class']=classify(ss['whm_m_s'],par['hand_high_m_s'],par['stationary_translation_m_s'])
        ss['relative_class']=classify(ss['rhcm_discrete_m_s'],par['relative_high_m_s'],par['stationary_translation_m_s'])
        angular=classify(ss['rham_deg_s'],par['camera_high_deg_s'],par['stationary_rotation_deg_s'])
        ss['relative_pose_class']='unknown' if 'unknown' in (ss['relative_class'],angular) else 'stationary' if ss['relative_class']==angular=='stationary' else 'moving'
        summary['sides'][side]=ss
    def aggregate(key):
        ss=[summary['sides'][s][key] for s in ('left','right')]
        valid=[s for s in ss if s['p90'] is not None]
        return max((s['p90'] for s in valid),default=None)
    def any_class(key):
        classes=[summary['sides'][s][key] for s in ('left','right')]
        # Missing one side cannot certify both stationary; any-side stationary is explicit.
        return next((v for v in ('high','low','stationary') if v in classes),'unknown')
    ct=summary['camera_translation'];cr=summary['camera_rotation']
    for key,name in [('hand_class','hand_both_class'),('relative_class','relative_both_class')]:
        classes=[summary['sides'][s][key] for s in ('left','right')]
        summary[name]='unknown' if 'unknown' in classes else classes[0] if classes[0]==classes[1] else 'mixed'
    pose_classes=[summary['sides'][s]['relative_pose_class'] for s in ('left','right')]
    summary['relative_pose_class']=next((v for v in ('moving','stationary') if v in pose_classes),'unknown')
    summary.update(camera_class='unknown' if not (eligible(ct) and eligible(cr)) else 'stationary' if ct['p90']<=par['stationary_translation_m_s'] and cr['p90']<=par['stationary_rotation_deg_s'] else 'high' if ct['p90']>=par['camera_high_m_s'] or cr['p90']>=par['camera_high_deg_s'] else 'low',
        hand_class=any_class('hand_class'),relative_class=any_class('relative_class'),relative_definition='wrist position relative to camera; orientation separately reported',
        camera_translation_speed_p90_m_s=ct['p90'],camera_angular_speed_p90_deg_s=cr['p90'],camera_valid_fraction=float(cv.mean()),camera_pairs=int(cp.sum()),
        hand_world_speed_p90_m_s=aggregate('whm_m_s'),hand_camera_speed_p90_m_s=aggregate('rhcm_discrete_m_s'),
        hand_world_valid_fraction=float(hv.mean()),hand_camera_valid_fraction=float((hv&cv[:,None]).mean()),
        cmh_p90_m_s=aggregate('cmh_m_s'),rhcm_p90_m_s=aggregate('rhcm_discrete_m_s'),rham_p90_deg_s=aggregate('rham_deg_s'),wms_p90=aggregate('wms'),
        gt_coverage=max(summary['sides'][s]['rhcm_discrete_m_s']['coverage'] for s in ('left','right')),
        hand_motion_definition='wrist translation, not finger articulation')
    arrays.update(camera_translation_m_s=camera_translation,camera_rotation_deg_s=camera_rotation,dt_s=dt,timestamps_s=t[:-1],pair_valid=rp)
    return summary,arrays