#!/usr/bin/env python3
'''Multi-horizon open-loop RTK/IMU propagation audit without optimization.'''
from __future__ import annotations
import argparse, json, sys
from pathlib import Path
import numpy as np
from scipy.spatial.transform import Rotation
ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path: sys.path.insert(0,str(ROOT))
from imu_lidar.imu_preintegration import preintegrate_imu
from rtk_imu.rtk_imu_engineering import (
    G_ENU, _all_hpr, _height_reference, _hpr_factor_observation, _world_rtk)
from rtk_imu.rtk_imu_multisource import load_unified_sessions
from tools.audit_rtk_imu_factor_consistency import (
    MECHANICAL_L_I_M, _best_arrays, _jsonable, _nearest_imu, _summary)

HORIZONS_S = (1.,2.,5.,10.,20.)

def _angle_deg(left,right):
    return float(np.degrees(np.arccos(np.clip(np.dot(left,right),-1.,1.))))

def _scalar_summary(values):
    a = np.asarray(values,dtype=float)
    if not a.size:
        return {'count':0,'rms':np.nan,'p95_abs':np.nan}
    return {'count':len(a),'rms':float(np.sqrt(np.mean(a*a))),
            'p95_abs':float(np.percentile(np.abs(a),95.))}

def _evaluate_session(session,reference,rotation,lever,bg,ba):
    _,times,p_ant,v_ant = _best_arrays(session,reference)
    hpr = _all_hpr(session)
    baseline_I = rotation.T[:,0]
    values = {h:{'position':[],'velocity':[],'attitude':[],'baseline':[]}
              for h in HORIZONS_S}
    pair_pre = []
    for left,right in zip(times[:-1],times[1:]):
        if not .5 <= right-left <= 1.75:
            pair_pre.append(None)
        else:
            pair_pre.append(preintegrate_imu(
                session.imu.t_s,session.imu.gyro_rad_s,session.imu.acc_m_s2,
                float(left),float(right),bg,ba))
    for start,t0 in enumerate(times):
        baseline0,_,valid0,_,_ = _hpr_factor_observation(hpr,float(t0))
        imu0 = _nearest_imu(session,float(t0))
        if not valid0 or imu0 is None: continue
        R0 = _world_rtk(baseline0) @ rotation
        p_i0 = p_ant[start]-R0@lever
        v_i0 = v_ant[start]-R0@np.cross(session.imu.gyro_rad_s[imu0]-bg,lever)
        delta_R, delta_v, delta_p = np.eye(3), np.zeros(3), np.zeros(3)
        elapsed, matched = 0., set()
        for end in range(start+1,len(times)):
            pre = pair_pre[end-1]
            if pre is None: break
            delta_p = delta_p + delta_v*pre.duration_s + delta_R@pre.delta_p
            delta_v = delta_v + delta_R@pre.delta_v
            delta_R = delta_R@pre.delta_R
            elapsed = float(times[end]-t0)
            if elapsed > max(HORIZONS_S)+.25: break
            candidates = [h for h in HORIZONS_S if h not in matched and abs(elapsed-h) <= .25]
            if not candidates: continue
            horizon = min(candidates,key=lambda h:abs(elapsed-h))
            matched.add(horizon)
            baseline1,_,valid1,_,_ = _hpr_factor_observation(hpr,float(times[end]))
            imu1 = _nearest_imu(session,float(times[end]))
            if not valid1 or imu1 is None: continue
            p_i1 = p_i0+v_i0*elapsed+.5*G_ENU*elapsed**2+R0@delta_p
            v_i1 = v_i0+G_ENU*elapsed+R0@delta_v
            R1 = R0@delta_R
            predicted_p = p_i1+R1@lever
            predicted_v = v_i1+R1@np.cross(session.imu.gyro_rad_s[imu1]-bg,lever)
            observed_R1 = _world_rtk(baseline1)@rotation
            values[horizon]['position'].append(predicted_p-p_ant[end])
            values[horizon]['velocity'].append(predicted_v-v_ant[end])
            values[horizon]['attitude'].append(
                np.degrees(Rotation.from_matrix(observed_R1.T@R1).magnitude()))
            values[horizon]['baseline'].append(_angle_deg(R1@baseline_I,baseline1))
    summary = {str(int(h)):{
        'position_error_m':_summary(values[h]['position']),
        'velocity_error_m_s':_summary(values[h]['velocity']),
        'attitude_level_completed_error_deg':_scalar_summary(values[h]['attitude']),
        'observable_baseline_angular_error_deg':_scalar_summary(values[h]['baseline'])}
        for h in HORIZONS_S}
    return summary, values

def _summarize_values(values):
    return {str(int(h)):{
        'position_error_m':_summary(values[h]['position']),
        'velocity_error_m_s':_summary(values[h]['velocity']),
        'attitude_level_completed_error_deg':_scalar_summary(values[h]['attitude']),
        'observable_baseline_angular_error_deg':_scalar_summary(values[h]['baseline'])}
        for h in HORIZONS_S}

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--manifest',type=Path,required=True)
    parser.add_argument('--output',type=Path,required=True)
    parser.add_argument('--session',action='append',required=True)
    parser.add_argument('--rotation-rpy-deg',nargs=3,type=float,
                        default=[.4543066225,-.0026392019,.0122384129])
    parser.add_argument('--mechanical-l-I-m',nargs=3,type=float,
                        default=MECHANICAL_L_I_M.tolist())
    parser.add_argument('--gyro-bias-rad-s',nargs=3,type=float,default=[0.,0.,0.])
    parser.add_argument('--accel-bias-m-s2',nargs=3,type=float,default=[0.,0.,0.])
    args = parser.parse_args()
    sessions = load_unified_sessions(args.manifest,selected_session_ids=set(args.session))
    reference = _height_reference(sessions)
    if reference is None: raise RuntimeError('no BEST height reference')
    rotation = Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    lever = np.asarray(args.mechanical_l_I_m)
    bg, ba = np.asarray(args.gyro_bias_rad_s), np.asarray(args.accel_bias_m_s2)
    combined = {h:{'position':[],'velocity':[],'attitude':[],'baseline':[]}
                for h in HORIZONS_S}
    per_session = {}
    for session in sessions:
        per_session[session.session_id], raw = _evaluate_session(
            session,reference,rotation,lever,bg,ba)
        for h in HORIZONS_S:
            for key in combined[h]: combined[h][key].extend(raw[h][key])
    aggregate = _summarize_values(combined)
    one, twenty = aggregate['1'], aggregate['20']
    growth = {
        'position_p95_ratio_20s_over_1s':twenty['position_error_m']['vector_p95']/one['position_error_m']['vector_p95'],
        'velocity_p95_ratio_20s_over_1s':twenty['velocity_error_m_s']['vector_p95']/one['velocity_error_m_s']['vector_p95'],
        'attitude_p95_ratio_20s_over_1s':twenty['observable_baseline_angular_error_deg']['p95_abs']/one['observable_baseline_angular_error_deg']['p95_abs']}
    significant = bool(
        growth['position_p95_ratio_20s_over_1s'] >= 3.
        and twenty['position_error_m']['vector_p95']-one['position_error_m']['vector_p95'] >= .5)
    payload = {
        'scope':'open-loop only; no least_squares/free/prior/LOO/bootstrap/sensitivity',
        'least_squares_called':False,'rotation_source':'R2G_gravity_level_prior',
        'mechanical_l_I_m':lever,'gyro_bias_rad_s':bg,'accel_bias_m_s2':ba,
        'bias_source':'current nominal engineering bias; configurable CLI; no fit artifact bias was persisted',
        'aggregate_by_horizon_s':aggregate,'per_session_by_horizon_s':per_session,
        'one_second_reanchored_control':{
            'definition':'each interval restarts from observed BEST p/v and HPR-completed R',
            'result':one},
        'growth_20s_over_1s':growth,
        'significant_accumulated_drift':significant,
        'legacy_long_segment_deterministic_model_deprecated':significant}
    args.output.parent.mkdir(parents=True,exist_ok=True)
    args.output.write_text(json.dumps(_jsonable(payload),ensure_ascii=False,indent=2,
                                      allow_nan=False)+'\n',encoding='utf-8')
    print(json.dumps(_jsonable({'aggregate':aggregate,'growth':growth,
        'legacy_deprecated':significant}),ensure_ascii=False,indent=2))
    return 0

if __name__ == '__main__':
    raise SystemExit(main())
