#!/usr/bin/env python3
'''Run fixed-mechanical and soft-prior solutions on immutable 47-window baseline.'''
from __future__ import annotations
import argparse,hashlib,json,sys
from concurrent.futures import ThreadPoolExecutor
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 (
    NODE_DOF,build_problem,fit_states_at_fixed_lever,solve_free_lever_many,
    summarize_fixed_state_values)
from tools.audit_rtk_imu_factor_consistency import MECHANICAL_L_I_M,_jsonable
from tools.run_rtk_imu_node_graph_free_selected import _restore_segments

MANUAL_STD_M=np.array([.02,.02,.03])
MANUAL_COVARIANCE_M2=np.diag(MANUAL_STD_M**2)


def _factor_delta(candidate,baseline):
    output={}
    for key in ('best_position','doppler','hpr','imu_preintegration'):
        left,right=candidate['residual_by_factor'][key],baseline['residual_by_factor'][key]
        output[key]={'rms_delta':left['rms']-right['rms'],
            'p95_abs_delta':left['p95_abs']-right['p95_abs'],
            'nis_per_dof_delta':left['chi_square_per_dof']-right['chi_square_per_dof']}
    return output


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--manifest',type=Path,required=True)
    parser.add_argument('--selection',type=Path,required=True)
    parser.add_argument('--free-baseline',type=Path,required=True)
    parser.add_argument('--output',type=Path,required=True)
    parser.add_argument('--state-output',type=Path)
    parser.add_argument('--sample-period-s',type=float,default=1.)
    parser.add_argument('--fixed-max-nfev',type=int,default=120)
    parser.add_argument('--prior-max-nfev',type=int,default=120)
    parser.add_argument('--workers',type=int,default=4)
    parser.add_argument('--hpr-direct-sigma-rad',type=float,default=.006)
    parser.add_argument('--rotation-rpy-deg',nargs=3,type=float,
                        default=[.4543066225,-.0026392019,.0122384129])
    args=parser.parse_args()
    selection=json.loads(args.selection.read_text(encoding='utf-8'))
    baseline=json.loads(args.free_baseline.read_text(encoding='utf-8'))
    if len(selection['selected_windows'])!=47:
        raise RuntimeError('engineering branch requires immutable 47-window selection')
    ids={item['session_id'] for item in selection['selected_windows']}
    sessions=load_unified_sessions(args.manifest,selected_session_ids=ids)
    reference=_height_reference(sessions)
    segments=_restore_segments(sessions,reference,selection['selected_windows'],
                               args.sample_period_s)
    rotation=Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    mechanical=MECHANICAL_L_I_M.copy()
    problems=[build_problem(segment,rotation,mechanical,args.hpr_direct_sigma_rad)
              for segment in segments]
    with ThreadPoolExecutor(max_workers=args.workers) as executor:
        fitted=list(executor.map(lambda problem:fit_states_at_fixed_lever(
            problem,mechanical,args.fixed_max_nfev),problems))
    fixed_states=[item[0] for item in fitted]
    fixed_optimizer=[item[1] for item in fitted]
    fixed_summary=summarize_fixed_state_values(problems,fixed_states,mechanical)
    prior_result,prior_states=solve_free_lever_many(
        problems,mechanical,args.prior_max_nfev,fixed_states,
        lever_prior_mean_m=mechanical,
        lever_prior_covariance_m2=MANUAL_COVARIANCE_M2,
        return_state_values=True)
    prior_result=asdict(prior_result)
    prior_data=summarize_fixed_state_values(
        problems,prior_states,np.asarray(prior_result['final_l_I_m']))
    bias_values={}
    for segment,state in zip(segments,prior_states):
        states=np.asarray(state).reshape(-1,NODE_DOF)
        values=bias_values.setdefault(segment.session_id,{'bg':[],'ba':[]})
        values['bg'].extend(states[:,9:12])
        values['ba'].extend(states[:,12:15])
    calibration_bias={session_id:{
        'gyro_bias_rad_s':np.median(values['bg'],axis=0),
        'accel_bias_m_s2':np.median(values['ba'],axis=0),
        'node_count':len(values['bg'])}
        for session_id,values in bias_values.items()}
    free=baseline['solutions']['mechanical']
    posterior_cov=np.asarray(prior_result['lever_covariance_m2'])
    variance_ratio=np.diag(posterior_cov)/np.diag(MANUAL_COVARIANCE_M2)
    prior_pull=(np.asarray(prior_result['final_l_I_m'])-mechanical)/MANUAL_STD_M
    translation_refined=bool(np.all(variance_ratio<=.90))
    comparisons={'fixed_vs_free':{
        'cost_delta':fixed_summary['cost']-free['final_cost'],
        'relative_cost_delta':fixed_summary['cost']/free['final_cost']-1.,
        'factor_residual_delta':_factor_delta(fixed_summary,free)},
        'prior_data_vs_free':{
            'cost_delta':prior_data['cost']-free['final_cost'],
            'relative_cost_delta':prior_data['cost']/free['final_cost']-1.,
            'factor_residual_delta':_factor_delta(prior_data,free)}}
    baseline_hash=hashlib.sha256(args.free_baseline.read_bytes()).hexdigest()
    payload={'scope':'47-window mechanical-prior engineering branch',
        'data_only_translation_accepted':False,
        'immutable_free_baseline':{'path':str(args.free_baseline),
            'sha256':baseline_hash,'solution':free},
        'selected_window_count':len(segments),'selection_path':str(args.selection),
        'manual_l_I_m':mechanical,'manual_l_I_std_m':MANUAL_STD_M,
        'manual_l_I_covariance_m2':MANUAL_COVARIANCE_M2,
        'fixed_mechanical_solution':{'l_I_m':mechanical,
            'window_optimizer':fixed_optimizer,'data_summary':fixed_summary},
        'prior_constrained_solution':{'result':prior_result,
            'data_only_summary_excluding_prior_factor':prior_data,
            'posterior_covariance_m2':posterior_cov,
            'posterior_prior_variance_ratio':variance_ratio,
            'prior_pull_sigma':prior_pull,
            'calibration_only_frozen_bias_by_session':calibration_bias},
        'comparisons':comparisons,
        'translation_refinement_gate':{
            'required_max_axis_variance_ratio':.90,
            'passed':translation_refined},
        'translation_refined_by_data':translation_refined,
        'heldout_validation_called':False,
        'engineering_translation_accepted':False}
    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')
    if args.state_output is not None:
        sizes=np.asarray([len(value) for value in prior_states],dtype=int)
        args.state_output.parent.mkdir(parents=True,exist_ok=True)
        np.savez_compressed(args.state_output,
            states=np.concatenate(prior_states),offsets=np.cumsum(np.r_[0,sizes]),
            candidate_ids=np.asarray(
                [item['candidate_id'] for item in selection['selected_windows']]))
    print(json.dumps(_jsonable({'fixed':fixed_summary,
        'prior_l_I_m':prior_result['final_l_I_m'],
        'prior_data_cost':prior_data['cost'],'prior_map_cost':prior_result['final_cost'],
        'variance_ratio':variance_ratio,'prior_pull_sigma':prior_pull,
        'translation_refined_by_data':translation_refined,
        'fixed_optimizer_failed_count':sum(not item['success'] for item in fixed_optimizer)}),
        ensure_ascii=False,indent=2))
    return 0


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