#!/usr/bin/env python3
'''P0.5 fixed-lever node-graph covariance and motion-window audit.'''
from __future__ import annotations
import argparse, json, sys
from dataclasses import asdict
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 rtk_imu.rtk_imu_engineering import _height_reference
from rtk_imu.rtk_imu_multisource import load_unified_sessions
from rtk_imu.rtk_imu_node_graph import (
    build_problem,solve_fixed_lever,solve_fixed_lever_many)
from tools.audit_rtk_imu_factor_consistency import MECHANICAL_L_I_M,_jsonable
from tools.run_rtk_imu_node_graph_fixed_lever import _select_window

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('--circle-session',required=True)
    parser.add_argument('--left-right-session',required=True)
    parser.add_argument('--slope-session',required=True)
    parser.add_argument('--sample-period-s',type=float,default=1.)
    parser.add_argument('--target-duration-s',type=float,default=15.)
    parser.add_argument('--max-nfev',type=int,default=50)
    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('--hpr-direct-sigma-rad',type=float,default=.006)
    args = parser.parse_args()
    categories = {'circle':args.circle_session,'left_right':args.left_right_session,
                  'slope':args.slope_session}
    sessions = load_unified_sessions(args.manifest,selected_session_ids=set(categories.values()))
    reference = _height_reference(sessions)
    rotation = Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    lever = np.asarray(args.mechanical_l_I_m)
    problems, selections, separate = [], {}, {}
    for category,session_id in categories.items():
        local = [session for session in sessions if session.session_id == session_id]
        segment,selection = _select_window(
            local,reference,args.sample_period_s,args.target_duration_s)
        problem = build_problem(segment,rotation,lever,args.hpr_direct_sigma_rad)
        problems.append(problem); selections[category] = selection
        separate[category] = asdict(solve_fixed_lever(problem,args.max_nfev))
    joint = asdict(solve_fixed_lever_many(problems,args.max_nfev))
    fixed_pass = bool(
        all(value['success'] for value in separate.values()) and joint['success']
        and joint['chi_square_per_dof'] >= .25
        and joint['chi_square_per_dof'] <= 4.
        and joint['final_position_residual_m']['vector_p95'] <= .20
        and joint['final_velocity_residual_m_s']['vector_p95'] <= .50)
    payload = {'scope':'P0.5 three-motion fixed-lever only',
        'translation_variable_enabled':False,
        'manual_prior_loo_bootstrap_sensitivity_called':False,
        'covariance_model':{
            'source':'independent_innovation_audit_three_motion',
            'best_position_xyz_m':[.06,.06,.12],
            'doppler_xyz_m_s':[.15,.15,.30],
            'hpr_direct_angular_rad':args.hpr_direct_sigma_rad,
            'hpr_bridge_rule':'sqrt(direct_sigma^2 + dropout_extra_variance)',
            'imu_preintegration':'unchanged physical covariance',
            'bias_random_walk':'unchanged static/Allan/device model',
            'postfit_global_scale_applied':False},
        'fixed_l_I_m':lever,'selections':selections,
        'separate_fixed_lever':separate,'joint_fixed_lever':joint,
        'fixed_lever_covariance_gate':{
            'chi_square_per_dof_range':[.25,4.],
            'position_vector_p95_max_m':.20,
            'velocity_vector_p95_max_m_s':.50,
            'passed':fixed_pass},
        'whitening_diagnosis':{
            'normalized_scale_consistent':joint['chi_square_per_dof'] >= .25,
            'joint_preintegration_nis_per_dof':
                joint['final_residual_by_factor']['imu_preintegration']['chi_square_per_dof'],
            'joint_preintegration_normalized_p95':
                joint['final_residual_by_factor']['imu_preintegration']['p95_abs'],
            'preintegration_sigma_distribution':joint['preintegration_covariance_sigma'],
            'interpretation':(
                'absolute preintegration sigma is small, but per-node states satisfy process '
                'factors almost exactly while BEST/Doppler/HPR normalized residuals are also '
                'well below one; current factor covariance set is collectively overconservative'
            )},
        'free_lever_unlocked':fixed_pass}
    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')
    compact = {'selections':selections,
        'separate':{key:{'success':value['success'],
            'chi_square_per_dof':value['chi_square_per_dof'],
            'position_p95_m':value['final_position_residual_m']['vector_p95'],
            'velocity_p95_m_s':value['final_velocity_residual_m_s']['vector_p95'],
            'factor_stats':value['final_residual_by_factor']}
            for key,value in separate.items()},
        'joint':{'success':joint['success'],'chi_square_per_dof':joint['chi_square_per_dof'],
            'position_p95_m':joint['final_position_residual_m']['vector_p95'],
            'velocity_p95_m_s':joint['final_velocity_residual_m_s']['vector_p95'],
            'factor_stats':joint['final_residual_by_factor']},
        'free_lever_unlocked':fixed_pass}
    print(json.dumps(_jsonable(compact),ensure_ascii=False,indent=2))
    return 0

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