#!/usr/bin/env python3
'''Run 18 fixed-R2G perturbation prior-constrained node-graph solves.'''
from __future__ import annotations
import argparse,json,sys
from concurrent.futures import ProcessPoolExecutor,as_completed
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 imu_lidar.geometry import make_transform
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_free_lever_many,summarize_fixed_state_values)
from tools.audit_rtk_imu_factor_consistency import _jsonable
from tools.run_rtk_imu_node_graph_free_selected import _restore_segments

_CONTEXT={}
MANUAL=np.array([-.45072,-.25682,.73208])
STD=np.array([.02,.02,.03])
COV=np.diag(STD**2)

def _initialize_worker(segments,states,rpy,max_nfev,sigma,checkpoint_dir):
    _CONTEXT.update(segments=segments,states=states,rpy=np.asarray(rpy),
        max_nfev=max_nfev,sigma=sigma,checkpoint_dir=Path(checkpoint_dir))

def _solve(task):
    axis,sign,angle=task
    name=f'{axis}_{sign*angle:+.1f}deg'
    path=_CONTEXT['checkpoint_dir']/f'{name}.json'
    if path.exists(): return json.loads(path.read_text(encoding='utf-8'))
    vector=np.zeros(3); vector['xyz'.index(axis)]=np.deg2rad(sign*angle)
    nominal_R=Rotation.from_euler('xyz',_CONTEXT['rpy'],degrees=True).as_matrix()
    perturbed_R=Rotation.from_rotvec(vector).as_matrix()@nominal_R
    problems=[build_problem(segment,perturbed_R,MANUAL,_CONTEXT['sigma'])
              for segment in _CONTEXT['segments']]
    result,states=solve_free_lever_many(problems,MANUAL,_CONTEXT['max_nfev'],
        _CONTEXT['states'],lever_prior_mean_m=MANUAL,
        lever_prior_covariance_m2=COV,return_state_values=True)
    result=asdict(result); lever=np.asarray(result['final_l_I_m'])
    data=summarize_fixed_state_values(problems,states,lever)
    transform=make_transform(-perturbed_R@lever,perturbed_R)
    payload={'name':name,'axis':axis,'signed_angle_deg':sign*angle,
        'success':result['success'],'message':result['message'],'nfev':result['nfev'],
        'final_l_I_m':lever,'delta_l_from_nominal_m':None,
        'map_cost':result['final_cost'],'data_cost':data['cost'],
        'global_chi_square_per_dof':data['chi_square_per_dof'],
        'residual_by_factor':data['residual_by_factor'],
        'physical_residual':{
            'BEST_position_m':data['best_position_physical_m'],
            'Doppler_m_s':data['doppler_physical_m_s'],
            'HPR_rad':data['hpr_physical_rad']},
        'posterior_covariance_m2':result['lever_covariance_m2'],
        'posterior_std_m':np.sqrt(np.diag(result['lever_covariance_m2'])),
        'prior_pull_sigma':(lever-MANUAL)/STD,
        'T_RTK_IMU':transform}
    path.write_text(json.dumps(_jsonable(payload),ensure_ascii=False,indent=2,
                               allow_nan=False)+'\n',encoding='utf-8')
    return _jsonable(payload)

def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--manifest',type=Path,required=True)
    p.add_argument('--selection',type=Path,required=True)
    p.add_argument('--engineering-result',type=Path,required=True)
    p.add_argument('--nominal-states',type=Path,required=True)
    p.add_argument('--output',type=Path,required=True)
    p.add_argument('--checkpoint-dir',type=Path,required=True)
    p.add_argument('--workers',type=int,default=3)
    p.add_argument('--max-nfev',type=int,default=120)
    p.add_argument('--hpr-direct-sigma-rad',type=float,default=.006)
    p.add_argument('--rotation-rpy-deg',nargs=3,type=float,
                   default=[.4543066225,-.0026392019,.0122384129])
    args=p.parse_args()
    engineering=json.loads(args.engineering_result.read_text(encoding='utf-8'))
    selection=json.loads(args.selection.read_text(encoding='utf-8'))
    if len(selection['selected_windows'])!=47:
        raise RuntimeError('sensitivity requires immutable 47-window selection')
    sessions=load_unified_sessions(args.manifest,
        selected_session_ids={x['session_id'] for x in selection['selected_windows']})
    segments=_restore_segments(sessions,_height_reference(sessions),
                               selection['selected_windows'],1.)
    saved=np.load(args.nominal_states,allow_pickle=False)
    offsets=saved['offsets']; flat=saved['states']
    states=[flat[offsets[i]:offsets[i+1]] for i in range(len(offsets)-1)]
    if list(saved['candidate_ids'])!=[x['candidate_id'] for x in selection['selected_windows']]:
        raise RuntimeError('nominal state checkpoint does not match selection order')
    args.checkpoint_dir.mkdir(parents=True,exist_ok=True)
    tasks=[(axis,sign,angle) for axis in 'xyz'
           for angle in (.1,.3,.5) for sign in (-1.,1.)]
    results=[]
    with ProcessPoolExecutor(max_workers=args.workers,initializer=_initialize_worker,
            initargs=(segments,states,args.rotation_rpy_deg,args.max_nfev,
                      args.hpr_direct_sigma_rad,str(args.checkpoint_dir))) as executor:
        futures=[executor.submit(_solve,task) for task in tasks]
        for future in as_completed(futures):
            result=future.result(); results.append(result)
            print('completed',result['name'],flush=True)
    results.sort(key=lambda x:('xyz'.index(x['axis']),x['signed_angle_deg']))
    nominal=engineering['prior_constrained_solution']
    nominal_l=np.asarray(nominal['result']['final_l_I_m'])
    nominal_data=nominal['data_only_summary_excluding_prior_factor']
    nominal_R=Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    nominal_T=make_transform(-nominal_R@nominal_l,nominal_R)
    for result in results:
        result['delta_l_from_nominal_m']=np.asarray(result['final_l_I_m'])-nominal_l
        result['delta_l_norm_m']=float(np.linalg.norm(result['delta_l_from_nominal_m']))
        result['transform_translation_delta_m']=(
            np.asarray(result['T_RTK_IMU'])[:3,3]-nominal_T[:3,3])
        result['transform_translation_delta_norm_m']=float(np.linalg.norm(
            result['transform_translation_delta_m']))
        result['transform_rotation_delta_deg']=abs(result['signed_angle_deg'])
        result['relative_data_cost_delta']=result['data_cost']/nominal_data['cost']-1.
        result['factor_p95_delta_sigma']={key:
            result['residual_by_factor'][key]['p95_abs']-
            nominal_data['residual_by_factor'][key]['p95_abs']
            for key in ('best_position','doppler','hpr','imu_preintegration')}
    delta=np.asarray([x['delta_l_from_nominal_m'] for x in results])
    at_point3=[x for x in results if abs(x['signed_angle_deg'])==.3]
    checks={'all_18_complete_and_converged':
                len(results)==18 and all(x['success'] for x in results),
        'max_delta_norm_at_0p3deg_le_0p10m':
                max(x['delta_l_norm_m'] for x in at_point3)<=.10,
        'max_delta_norm_all_le_0p15m':
                max(x['delta_l_norm_m'] for x in results)<=.15,
        'relative_data_cost_increase_all_le_0p05':
                max(x['relative_data_cost_delta'] for x in results)<=.05,
        'factor_normalized_p95_increase_all_le_0p5sigma':
                max(v for x in results for v in x['factor_p95_delta_sigma'].values())<=.5}
    passed=bool(all(checks.values()))
    payload={'scope':'18 prior-constrained fixed-R2G perturbation solves',
        'data_only_free_called':False,'bootstrap_called':False,'loo_called':False,
        'parser_R0_covariance_modified':False,'calibration_window_count':47,
        'nominal_rotation_rpy_deg':args.rotation_rpy_deg,
        'nominal_l_I_m':nominal_l,'nominal_T_RTK_IMU':nominal_T,
        'perturbations':results,'summary':{
            'max_abs_delta_l_xyz_m':np.max(np.abs(delta),axis=0),
            'max_delta_l_norm_m':float(np.max(np.linalg.norm(delta,axis=1))),
            'max_transform_translation_delta_norm_m':max(
                x['transform_translation_delta_norm_m'] for x in results)},
        'diagnostic_gate_thresholds_frozen_before_run':{
            'max_delta_norm_at_0p3deg_m':.10,'max_delta_norm_all_m':.15,
            'relative_data_cost_increase':.05,
            'factor_normalized_p95_increase_sigma':.5},
        'gate_checks':checks,'rotation_sensitivity_passed':passed}
    args.output.write_text(json.dumps(_jsonable(payload),ensure_ascii=False,indent=2,
                                      allow_nan=False)+'\n',encoding='utf-8')
    print(json.dumps(_jsonable({'summary':payload['summary'],
        'checks':checks,'rotation_sensitivity_passed':passed}),indent=2))
    return 0
if __name__=='__main__': raise SystemExit(main())
