#!/usr/bin/env python3
'''Validate an engineering lever on disjoint held-out windows.'''
from __future__ import annotations
import argparse,json,sys
from concurrent.futures import ProcessPoolExecutor,as_completed
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 imu_lidar.rtk_imu_engineering import _height_reference
from imu_lidar.rtk_imu_multisource import load_unified_sessions
from imu_lidar.rtk_imu_node_graph import build_problem,fit_states_at_fixed_lever,residual
from tools.audit_rtk_imu_factor_consistency import _jsonable
from tools.run_rtk_imu_node_graph_free_selected import _restore_segments

FACTORS=('best_position','doppler','hpr','imu_preintegration')
PHYSICAL=('best_position_physical','doppler_physical','hpr_physical')

def _fit_checkpoint(task):
    index,problem,lever,max_nfev,path_text=task
    path=Path(path_text)
    if path.exists():
        data=np.load(path,allow_pickle=False)
        return index,data['state'],json.loads(str(data['optimizer']))
    state,meta=fit_states_at_fixed_lever(problem,lever,max_nfev)
    np.savez_compressed(path,state=state,optimizer=json.dumps(meta))
    return index,state,meta

def _stats(values,dof=None):
    a=np.asarray(values,dtype=float).reshape(-1)
    if not a.size:
        return {'count':0,'dof':0,'rms':np.nan,'p50_abs':np.nan,
                'p95_abs':np.nan,'p99_abs':np.nan,'nis':np.nan,
                'chi_square_per_dof':np.nan}
    dof=len(a) if dof is None else max(int(dof),1); nis=float(a@a)
    return {'count':len(a),'dof':dof,'rms':float(np.sqrt(np.mean(a*a))),
            'p50_abs':float(np.percentile(np.abs(a),50)),
            'p95_abs':float(np.percentile(np.abs(a),95)),
            'p99_abs':float(np.percentile(np.abs(a),99)),
            'nis':nis,'chi_square_per_dof':nis/dof}

def _vectors(values):
    a=np.asarray(values,dtype=float).reshape(-1,3)
    if not a.size:
        return {'count':0,'axis_rms':[np.nan]*3,'axis_p95_abs':[np.nan]*3,
                'vector_rms':np.nan,'vector_p95':np.nan}
    norm=np.linalg.norm(a,axis=1)
    return {'count':len(a),'axis_rms':np.sqrt(np.mean(a*a,axis=0)),
            'axis_p95_abs':np.percentile(np.abs(a),95,axis=0),
            'vector_rms':float(np.sqrt(np.mean(norm*norm))),
            'vector_p95':float(np.percentile(norm,95))}

def _summarize(records):
    factor={key:[] for key in FACTORS}; physical={key:[] for key in PHYSICAL}
    joined=[]; state_dimension=0; other_effective=0
    for record in records:
        joined.extend(record['all']); state_dimension+=record['state_dimension']
        other_effective+=record['other_effective']
        for key in FACTORS: factor[key].extend(record['factor'].get(key,[]))
        for key in PHYSICAL: physical[key].extend(record['physical'].get(key,[]))
    effective=sum(2*len(v)//3 if k=='hpr' else len(v)
                  for k,v in factor.items())+other_effective
    dof=max(effective-state_dimension,1); a=np.asarray(joined,dtype=float)
    converged=sum(bool(x['optimizer']['success']) for x in records)
    return {'window_count':len(records),'optimizer_converged_count':converged,
        'optimizer_converged_fraction':converged/max(len(records),1),
        'total_residual_dimension':len(a),'state_dimension':state_dimension,
        'statistical_dof':dof,'total_nis':float(a@a),
        'global_chi_square_per_dof':float((a@a)/dof),
        'residual_by_factor':{key:_stats(value,2*len(value)//3 if key=='hpr' else None)
                              for key,value in factor.items()},
        'best_position_physical_m':_vectors(physical['best_position_physical']),
        'doppler_physical_m_s':_vectors(physical['doppler_physical']),
        'hpr_physical_rad':_vectors(physical['hpr_physical'])}

def _gate(summary):
    factor=summary['residual_by_factor']
    checks={'optimizer_converged_fraction_ge_0p90':
                summary['optimizer_converged_fraction']>=.90,
        'global_chi_square_per_dof_in_0p25_4':
                .25<=summary['global_chi_square_per_dof']<=4.,
        'best_vector_p95_le_0p20_m':
                summary['best_position_physical_m']['vector_p95']<=.20,
        'doppler_vector_p95_le_0p50_m_s':
                summary['doppler_physical_m_s']['vector_p95']<=.50,
        'hpr_normalized_p95_le_4':factor['hpr']['p95_abs']<=4.,
        'preintegration_normalized_p95_le_3':
                factor['imu_preintegration']['p95_abs']<=3.}
    return {'uses_existing_P0p5_physical_and_statistical_health_gates':True,
            'checks':checks,'passed':bool(all(checks.values()))}

def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--manifest',type=Path,required=True)
    p.add_argument('--calibration-selection',type=Path,required=True)
    p.add_argument('--all-selection',type=Path,required=True)
    p.add_argument('--engineering-result',type=Path,required=True)
    p.add_argument('--output',type=Path,required=True)
    p.add_argument('--checkpoint-dir',type=Path,required=True)
    p.add_argument('--sample-period-s',type=float,default=1.)
    p.add_argument('--max-nfev',type=int,default=120)
    p.add_argument('--workers',type=int,default=4)
    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])
    p.add_argument('--circle-session',default='0808_20260808_092827')
    p.add_argument('--left-right-session',default='0808_20260808_082148')
    p.add_argument('--slope-session',default='0815_20260812_123424')
    args=p.parse_args()
    engineering=json.loads(args.engineering_result.read_text(encoding='utf-8'))
    calibration=json.loads(args.calibration_selection.read_text(encoding='utf-8'))
    all_windows=json.loads(args.all_selection.read_text(encoding='utf-8'))
    calibration_ids={x['candidate_id'] for x in calibration['selected_windows']}
    heldout=[x for x in all_windows['selected_windows']
             if x['candidate_id'] not in calibration_ids]
    if len(calibration_ids)!=47 or len(heldout)!=267:
        raise RuntimeError(f'expected 47+267 windows, got {len(calibration_ids)}+{len(heldout)}')
    shared=sum(x.get('shared_sample_count_with_previous',{}).get(k,0)
               for x in all_windows['selected_windows'] for k in ('imu','gnss','hpr'))
    if shared: raise RuntimeError(f'selection contains {shared} shared samples')
    lever=np.asarray(engineering['prior_constrained_solution']['result']['final_l_I_m'])
    sessions=load_unified_sessions(
        args.manifest,selected_session_ids={x['session_id'] for x in heldout})
    reference=_height_reference(sessions)
    segments=_restore_segments(sessions,reference,heldout,args.sample_period_s)
    rotation=Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    problems=[build_problem(x,rotation,lever,args.hpr_direct_sigma_rad) for x in segments]
    args.checkpoint_dir.mkdir(parents=True,exist_ok=True)
    fitted=[None]*len(problems)
    tasks=[(i,problem,lever,args.max_nfev,
            str(args.checkpoint_dir/f'{i:04d}.npz'))
           for i,problem in enumerate(problems)]
    with ProcessPoolExecutor(max_workers=args.workers) as executor:
        futures=[executor.submit(_fit_checkpoint,task) for task in tasks]
        completed=0
        for future in as_completed(futures):
            index,state,meta=future.result(); fitted[index]=(state,meta); completed+=1
            if completed%10==0: print(f'held-out checkpoint {completed}/{len(problems)}',flush=True)
    mapping={args.circle_session:'circle',args.left_right_session:'left_right',
             args.slope_session:'slope'}
    records=[]
    for item,problem,fit in zip(heldout,problems,fitted):
        state,optimizer=fit; detail={}
        all_residual=residual(problem,state,details=detail,lever_override=lever)
        known=sum(len(detail.get(k,[])) for k in FACTORS)
        records.append({'candidate_id':item['candidate_id'],
            'session_id':item['session_id'],
            'motion_class':mapping.get(item['session_id'],'other_recovered_dynamic'),
            'start_s':item['start_s'],'end_s':item['end_s'],
            'state_dimension':len(state),'optimizer':optimizer,
            'all':all_residual.tolist(),'other_effective':len(all_residual)-known,
            'factor':{k:detail.get(k,[]) for k in FACTORS},
            'physical':{k:detail.get(k,[]) for k in PHYSICAL}})
    overall=_summarize(records)
    by_session={key:_summarize([x for x in records if x['session_id']==key])
                for key in sorted({x['session_id'] for x in records})}
    by_motion={key:_summarize([x for x in records if x['motion_class']==key])
               for key in sorted({x['motion_class'] for x in records})}
    overall_gate=_gate(overall)
    session_gates={k:_gate(v) for k,v in by_session.items()}
    motion_gates={k:_gate(v) for k,v in by_motion.items()}
    heldout_passed=bool(overall_gate['passed'] and
        all(x['passed'] for x in session_gates.values()) and
        all(x['passed'] for x in motion_gates.values()))
    comparisons=engineering['comparisons']
    def nonconflicting(name):
        c=comparisons[name]
        return (c['relative_cost_delta']<=.05 and
                all(x['p95_abs_delta']<=.25 for x in c['factor_residual_delta'].values()))
    fit_nonconflict=nonconflicting('fixed_vs_free') and nonconflicting('prior_data_vs_free')
    T_rtk_imu=make_transform(-rotation@lever,rotation)
    T_imu_rtk=np.linalg.inv(T_rtk_imu)
    accepted=bool(fit_nonconflict and heldout_passed)
    payload={**engineering,
        'scope':'47-window mechanical-prior engineering branch + disjoint held-out validation',
        'heldout_validation_called':True,
        'engineering_lever_source':'prior_constrained_solution',
        'engineering_l_I_m':lever,'calibration_heldout_overlap_count':0,
        'heldout_window_count':len(heldout),
        'heldout_validation':{'lever_reoptimized':False,
            'nuisance_states_optimized_per_window':True,
            'motion_class_mapping':mapping,
            'other_sessions_class':'other_recovered_dynamic',
            'overall':overall,'by_session':by_session,'by_motion_class':by_motion,
            'overall_gate':overall_gate,'session_gates':session_gates,
            'motion_class_gates':motion_gates,'passed':heldout_passed},
        'calibration_fit_nonconflict_gate':{'relative_cost_increase_max':.05,
            'per_factor_normalized_p95_increase_max':.25,
            'fixed_passed':nonconflicting('fixed_vs_free'),
            'prior_passed':nonconflicting('prior_data_vs_free'),
            'passed':fit_nonconflict},
        'engineering_translation_accepted':accepted,
        'data_only_translation_accepted':False,
        'result_nature':'mechanical lever + dynamic-data consistency validation; not data-only translation calibration',
        'rotation_source':'R2G_gravity_level_prior',
        'translation_conditional_on_rotation':True,
        'candidate_T_RTK_IMU':T_rtk_imu,'candidate_T_IMU_RTK':T_imu_rtk,
        'T_RTK_IMU':T_rtk_imu if accepted else None,
        'T_IMU_RTK':T_imu_rtk if accepted else None,
        'transform_convention':('T_RTK_IMU maps IMU coordinates into ANT1 RTK frame; '
            'l_I=p_ANT1^I; T_IMU_RTK translation equals l_I')}
    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({'engineering_l_I_m':lever,
        'fit_nonconflict':fit_nonconflict,'heldout_overall':overall,
        'heldout_passed':heldout_passed,
        'engineering_translation_accepted':accepted}),ensure_ascii=False,indent=2))
    return 0
if __name__=='__main__': raise SystemExit(main())