"""Offline whole-sequence temporal regularization with explicit distortion budgets."""
import argparse
import json
import shutil
from pathlib import Path
import xml.etree.ElementTree as ET

import mujoco
import numpy as np
from scipy import sparse
from scipy.sparse.linalg import splu
from scipy.spatial.transform import Rotation

ROOT = Path(__file__).resolve().parents[1]


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--source', type=Path, default=ROOT/'output/l20_15886123')
    p.add_argument('--output', type=Path, default=ROOT/'output/l20_15886123_stable')
    p.add_argument('--side', choices=['right','left'], default='right')
    a = p.parse_args()
    a.output.mkdir(parents=True, exist_ok=True)
    with np.load(a.source/'motion.npz') as f:
        m = {k:f[k] for k in f.files}
    n = len(m['qpos'])
    D2 = sparse.diags([np.ones(n-2), -2*np.ones(n-2), np.ones(n-2)], [0,1,2], shape=(n-2,n),format='csc')
    D3 = sparse.diags([-np.ones(n-3),3*np.ones(n-3),-3*np.ones(n-3),np.ones(n-3)], [0,1,2,3],shape=(n-3,n),format='csc')
    eye = sparse.eye(n,format='csc')
    names, active = list(m['joint_names']), list(m['active_joint_names'])
    ai = [names.index(k) for k in active]
    original = ET.parse(ROOT/f'third_party/l20_assets/L20/{a.side.upper()}/linkerhand_g20_{a.side}.urdf').getroot()
    limits = {j.get('name'):np.array([float(j.find('limit').get(k)) for k in ['lower','upper']]) for j in original.findall('joint') if j.get('type')!='fixed'}
    low, high = np.array([limits[k] for k in active]).T.copy()
    mimic = {}
    for j in original.findall('joint'):
        mi = j.find('mimic')
        if mi is None:
            continue
        par, sc, off = mi.get('joint'),float(mi.get('multiplier','1')),float(mi.get('offset','0'))
        mimic[j.get('name')] = par,sc,off
        lo,hi = np.sort((limits[j.get('name')]-off)/sc)
        k = active.index(par)
        low[k],high[k] = max(low[k],lo),min(high[k],hi)

    def smooth(x, weight, bounds=None):
        # Sample-domain objective ||x-y||² + ||D²x||² + weight*||D³x||².
        A = eye + D2.T@D2 + weight*(D3.T@D3)
        result = splu(A).solve(x)
        if bounds is not None:
            lo,hi = bounds
            # Explicit post-solve projection; evaluate smoothness after projection.
            # This is not claimed to solve the box-constrained temporal QP exactly.
            result = np.clip(result,lo,hi)
        return result

    def d2(x):
        return float(np.sqrt(np.mean(np.diff(x[30:],n=2,axis=0)**2)))

    rot = Rotation.from_quat(m['wrist_quat_wxyz'][:,[1,2,3,0]])
    rv = (rot[0].inv()*rot).as_rotvec()
    assert np.linalg.norm(rv,axis=1).max() < np.pi/2
    choices, sweep = {}, {}
    # Budgets are engineering choices relative to prior output, not true-motion tolerances.
    components = dict(fingers=(m['qpos'][:,ai],(low,high),np.deg2rad(8)),
                      wrist_position=(m['wrist_pos'],None,.020),
                      wrist_rotation=(rv,None,np.deg2rad(6)))
    for key,(x,bounds,budget) in components.items():
        rows, candidates = [], []
        for weight in [100.,300.,1000.,3000.,10000.,30000.]:
            y = smooth(x,weight,bounds)
            if key=='fingers':
                dev = float(np.abs(y-x).max())
            elif key=='wrist_rotation':
                dev = float(np.linalg.norm((rot.inv()*(rot[0]*Rotation.from_rotvec(y))).as_rotvec(),axis=1).max())
            else:
                dev = float(np.linalg.norm(y-x,axis=1).max())
            ratio = d2(y)/d2(x)
            row = dict(weight_d2=1.,weight_d3=weight,max_deviation=dev,budget=float(budget),
                       second_difference_ratio=ratio,within_budget=bool(dev<=budget))
            rows.append(row)
            if dev<=budget:
                candidates.append((ratio,y,row))
        assert candidates, key
        _,chosen,row = min(candidates,key=lambda v:v[0])
        choices[key] = chosen
        sweep[key] = dict(selected=row,candidates=rows)
        print(key,row,flush=True)
    before = {k:m[k].copy() for k in ['qpos','wrist_pos','wrist_quat_wxyz','actual']}
    q = m['qpos'].copy()
    q[:,ai] = choices['fingers']
    for name,(parent,sc,off) in mimic.items():
        q[:,names.index(name)] = sc*q[:,names.index(parent)]+off
    m['qpos'] = q
    m['active_qpos'] = q[:,ai]
    m['wrist_pos'] = choices['wrist_position']
    newrot = rot[0]*Rotation.from_rotvec(choices['wrist_rotation'])
    m['wrist_quat_wxyz'] = newrot.as_quat()[:,[3,0,1,2]]
    for k,v in before.items():
        m['before_stabilization_'+k] = v
    m['stabilization_method'] = 'whole-sequence data fidelity + second/third difference regularization; component distortion budgets'
    m['stabilization_source'] = str(a.source.resolve())
    for filename in ['l20_moving.xml',f'l20_{a.side}.xml','l20_dex.urdf','dex_config.json']:
        shutil.copy2(a.source/filename,a.output/filename)
    model = mujoco.MjModel.from_xml_path(str(a.output/f'l20_{a.side}.xml'))
    data = mujoco.MjData(model)
    addr = [model.jnt_qposadr[model.joint(k).id] for k in names]
    sites = [model.site(f'landmark_{i:02d}').id for i in range(21)]
    actual = []
    for row in q:
        data.qpos[addr] = row
        mujoco.mj_forward(model,data)
        actual.append(data.site_xpos[sites].copy())
    m['actual'] = np.array(actual)
    np.savez_compressed(a.output/'motion.npz',**m)
    np.savetxt(a.output/'trajectory.csv',np.c_[m['time'],m['wrist_pos'],m['wrist_quat_wxyz'],q],delimiter=',',
               header=','.join(['time_s','wrist_x','wrist_y','wrist_z','qw','qx','qy','qz']+names),comments='')
    lim = np.array([limits[k] for k in names])
    violation = float(max(np.maximum(lim[:,0]-q,0).max(),np.maximum(q-lim[:,1],0).max()))
    assert violation < 1e-9
    mimic_error = max(float(np.abs(q[:,names.index(k)]-sc*q[:,names.index(par)]-off).max()) for k,(par,sc,off) in mimic.items())
    assert mimic_error < 1e-10

    def metrics(q,pos,quat,actual):
        r = Rotation.from_quat(quat[:,[1,2,3,0]])[30:]
        w = (r[:-1].inv()*r[1:]).as_rotvec()
        tips = np.linalg.norm(actual[:,[4,8,12,16,20]]-m['targets'][:,[4,8,12,16,20]],axis=-1)*1000
        return dict(joint_second_difference_rms_deg=np.rad2deg(d2(q)),
                    joint_max_step_deg=float(np.rad2deg(np.abs(np.diff(q[30:],axis=0)).max())),
                    wrist_second_difference_rms_mm=d2(pos)*1000,
                    wrist_max_step_mm=float(np.linalg.norm(np.diff(pos[30:],axis=0),axis=1).max()*1000),
                    wrist_rotation_velocity_difference_rms_deg=float(np.rad2deg(np.sqrt(np.mean(np.diff(w,axis=0)**2)))),
                    wrist_rotation_max_step_deg=float(np.rad2deg(np.linalg.norm(w,axis=1).max())),
                    wrist_range_m=np.ptp(pos,axis=0).tolist(),
                    tip_error_mean_mm=float(tips.mean()),tip_error_p95_mm=float(np.percentile(tips,95)))
    report = dict(frames=n,fps=float(m['fps']),evaluation_start_frame=30,
                  before=metrics(**dict(q=before['qpos'],pos=before['wrist_pos'],quat=before['wrist_quat_wxyz'],actual=before['actual'])),
                  after=metrics(q,m['wrist_pos'],m['wrist_quat_wxyz'],m['actual']),
                  selection=sweep,urdf_limit_violation_rad=violation,linear_mimic_error_rad=mimic_error,
                  max_joint_change_deg=float(np.rad2deg(np.abs(q-before['qpos']).max())),
                  max_wrist_change_mm=float(np.linalg.norm(m['wrist_pos']-before['wrist_pos'],axis=1).max()*1000),
                  max_wrist_rotation_change_deg=float(np.rad2deg(np.linalg.norm((rot.inv()*newrot).as_rotvec(),axis=1).max())),
                  uses_future_frames=True,boundary='Lower temporal derivatives do not prove true-motion accuracy; no contact/dynamics validation')
    assert all(np.isfinite(m[k]).all() for k in ['qpos','wrist_pos','wrist_quat_wxyz','actual'])
    (a.output/'stabilization_validation.json').write_text(json.dumps(report,indent=2))
    print(json.dumps({k:v for k,v in report.items() if k!='selection'},indent=2))


if __name__=='__main__':
    main()