#!/usr/bin/env python3
'''Independent prediction innovations on the frozen 267 held-out windows.'''
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.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
from tools.audit_rtk_imu_factor_consistency import _jsonable
from tools.audit_rtk_imu_innovation_noise import (
    _factor_report,_hpr_innovations,_interval_innovations)
from tools.run_rtk_imu_node_graph_free_selected import _restore_segments

LIMITS={
    'best_position':{'bias_norm':.10,'vector_p95':.50},
    'doppler':{'bias_norm':.10,'vector_p95':.50},
    'hpr':{'bias_norm':.01,'vector_p95':.05},
}

def _finite_max_abs(values):
    a=np.asarray(values,dtype=float)
    return float(np.max(np.abs(a[np.isfinite(a)]))) if np.any(np.isfinite(a)) else np.inf

def _summary_gate(summary,limits,min_samples=20):
    if summary.get('sample_count',0)<min_samples:
        return {'evaluated':False,'passed':True,'reason':'insufficient_samples'}
    checks={'bias_norm':float(np.linalg.norm(summary['innovation_bias']))<=limits['bias_norm'],
            'vector_p95':summary['vector_p95']<=limits['vector_p95'],
            'lag1_autocorrelation_abs':_finite_max_abs(
                summary['temporal_autocorrelation_lag1'])<=.95}
    return {'evaluated':True,'thresholds':{**limits,'lag1_abs_max':.95},
            'checks':checks,'passed':bool(all(checks.values()))}

def _factor_gate(report,key):
    limits=LIMITS[key]; overall=_summary_gate(report['overall'],limits,1)
    grouped={}
    for group_name in ('by_session','by_motion','by_speed','by_gyro_norm'):
        grouped[group_name]={name:_summary_gate(value,limits)
                             for name,value in report[group_name].items()}
    evaluated=[gate for group in grouped.values() for gate in group.values()
               if gate['evaluated']]
    passed=overall['passed'] and all(gate['passed'] for gate in evaluated)
    return {'overall':overall,'grouped':grouped,'passed':bool(passed)}

def _calibration_biases(engineering,problems):
    source=engineering['prior_constrained_solution'].get(
        'calibration_only_frozen_bias_by_session')
    if not source:
        raise RuntimeError('engineering result lacks calibration-only frozen bias')
    by_date={}
    for session_id,value in source.items():
        by_date.setdefault(session_id.split('_')[1],[]).append(value)
    result={}
    for problem in problems:
        session_id=problem.segment.session_id
        if session_id in result: continue
        if session_id in source:
            value=source[session_id]; origin='same_calibration_session'
        else:
            values=by_date.get(session_id.split('_')[1],list(source.values()))
            weights=np.asarray([x['node_count'] for x in values],dtype=float)
            value={'gyro_bias_rad_s':np.average(
                       [x['gyro_bias_rad_s'] for x in values],axis=0,weights=weights),
                   'accel_bias_m_s2':np.average(
                       [x['accel_bias_m_s2'] for x in values],axis=0,weights=weights)}
            origin=('calibration_date_weighted_mean' if
                    session_id.split('_')[1] in by_date else 'calibration_global_weighted_mean')
        result[session_id]={'gyro_bias_rad_s':value['gyro_bias_rad_s'],
            'accel_bias_m_s2':value['accel_bias_m_s2'],'source':origin}
    return result

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('--sample-period-s',type=float,default=1.)
    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'))
    selected=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 selected['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 selected['selected_windows'] for k in ('imu','gnss','hpr'))
    if shared: raise RuntimeError(f'selection contains {shared} shared samples')
    lever=np.asarray(engineering.get('engineering_l_I_m',
        engineering['prior_constrained_solution']['result']['final_l_I_m']),dtype=float)
    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(segment,rotation,lever,args.hpr_direct_sigma_rad)
              for segment in segments]
    session_by_id={session.session_id:session for session in sessions}
    motion_by_session={args.circle_session:'circle',
        args.left_right_session:'left_right',args.slope_session:'slope'}
    records={key:[] for key in ('best_position','doppler','hpr','imu_preintegration')}
    biases=_calibration_biases(engineering,problems)
    for problem in problems:
        motion=motion_by_session.get(problem.segment.session_id,'other_recovered_dynamic')
        bias=biases[problem.segment.session_id]
        interval=_interval_innovations(problem,motion,
            bias['gyro_bias_rad_s'],bias['accel_bias_m_s2'])
        for key,value in interval.items(): records[key].extend(value)
        records['hpr'].extend(_hpr_innovations(
            session_by_id[problem.segment.session_id],problem,motion))
    factors={key:_factor_report(value) for key,value in records.items()}
    gates={key:_factor_gate(factors[key],key)
           for key in ('best_position','doppler','hpr')}
    passed=bool(all(value['passed'] for value in gates.values()))
    payload={'scope':'267-window independent held-out prediction innovations',
        'least_squares_called':False,'target_observation_used_by_predictor':False,
        'lever_reoptimized':False,'rotation_reoptimized':False,
        'covariance_parameters_modified':False,
        'fixed_l_I_m':lever,'fixed_rotation_rpy_deg':args.rotation_rpy_deg,
        'calibration_window_count':47,'heldout_window_count':267,
        'calibration_heldout_overlap_count':0,
        'motion_class_mapping':motion_by_session,
        'prediction_definition':{
            'BEST_position':'previous GNSS p/v + IMU preintegration + fixed lever',
            'Doppler':'previous GNSS velocity + IMU preintegration + omega-cross-lever',
            'HPR':'withheld direct HPR predicted by LOO interpolation and gyro propagation',
            'IMU_preintegration':'two independently GNSS/HPR-anchored endpoints'},
        'bias_source':('frozen prior-node bias from the disjoint 47-window calibration set; '
            'same-session when available, otherwise calibration date-weighted mean; '
            'no held-out target observation and no covariance writeback'),
        'per_session_frozen_bias':biases,
        'factors':factors,'validation_limits_are_physical_not_covariance_retuning':LIMITS,
        'factor_gates':gates,'independent_heldout_innovation_passed':passed}
    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={key:{'samples':value['overall'].get('sample_count',0),
        'bias':value['overall'].get('innovation_bias'),
        'vector_p95':value['overall'].get('vector_p95'),
        'nis_per_dof':value['overall'].get('nis_per_dof'),
        'gate':gates.get(key)} for key,value in factors.items()}
    print(json.dumps(_jsonable({'factors':compact,
        'independent_heldout_innovation_passed':passed}),ensure_ascii=False,indent=2))
    return 0
if __name__=='__main__': raise SystemExit(main())