#!/usr/bin/env python3
'''Root-cause audit for paired BEST/Doppler propagation innovations.'''
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 rtk_imu.rtk_imu_engineering import G0,G_ENU,_height_reference,_nodes,_world_rtk
from rtk_imu.rtk_imu_multisource import load_unified_sessions
from rtk_imu.rtk_imu_node_graph import build_problem
from tools.audit_rtk_imu_factor_consistency import _jsonable
from tools.audit_rtk_imu_heldout_innovation import _calibration_biases
from tools.audit_rtk_imu_innovation_noise import _factor_report,_interval_innovations
from tools.run_rtk_imu_node_graph_free_selected import _restore_segments

def _vector_summary(values):
    a=np.asarray(values,dtype=float).reshape(-1,3)
    if not len(a): return {'count':0}
    norm=np.linalg.norm(a,axis=1)
    return {'count':len(a),'bias':np.mean(a,axis=0),
        '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)),
        'empirical_covariance':np.cov(a,rowvar=False)}

def _correlation(left,right):
    a,b=np.asarray(left),np.asarray(right); result=np.full(3,np.nan)
    for axis in range(3):
        if len(a)>2 and np.std(a[:,axis])>1e-12 and np.std(b[:,axis])>1e-12:
            result[axis]=np.corrcoef(a[:,axis],b[:,axis])[0,1]
    return result

def _root_report(position,velocity):
    p={x['interval_id']:x for x in position}; v={x['interval_id']:x for x in velocity}
    ids=sorted(set(p)&set(v)); records=[]
    for key in ids:
        left,right=p[key],v[key]; dt=float(left['dt_s'])
        a_p=2.*np.asarray(left['residual'])/(dt*dt)
        a_v=np.asarray(right['residual'])/dt
        R=np.asarray(left['R0_WI'])
        records.append({**{k:left[k] for k in (
            'interval_id','session','motion','t_s','speed_bin','gyro_bin')},
            'dt_s':dt,'a_position':a_p,'a_velocity':a_v,
            'difference':a_v-a_p,'a_common':.5*(a_p+a_v),
            'a_common_body':R.T@(.5*(a_p+a_v))})
    def summarize(items):
        if not items:
            return {'interval_count':0,
                'a_err_from_position_m_s2':{'count':0},
                'a_err_from_velocity_m_s2':{'count':0},
                'per_axis_correlation':[np.nan]*3,
                'mean_direction_cosine':np.nan,
                'mean_magnitude_ratio_velocity_over_position':np.nan,
                'difference_velocity_minus_position_m_s2':{'count':0},
                'common_acceleration_world_m_s2':{'count':0},
                'common_acceleration_body_m_s2':{'count':0},
                'equivalent_horizontal_tilt_rad':np.nan,
                'equivalent_horizontal_tilt_deg':np.nan}
        ap=[x['a_position'] for x in items]; av=[x['a_velocity'] for x in items]
        diff=[x['difference'] for x in items]; common=[x['a_common'] for x in items]
        body=[x['a_common_body'] for x in items]
        mean_p=np.mean(ap,axis=0); mean_v=np.mean(av,axis=0)
        denom=np.linalg.norm(mean_p)*np.linalg.norm(mean_v)
        cosine=float(mean_p@mean_v/denom) if denom>1e-12 else np.nan
        ratio=float(np.linalg.norm(mean_v)/max(np.linalg.norm(mean_p),1e-12))
        body_bias=np.mean(body,axis=0)
        body_std=np.std(body,axis=0)
        horizontal=float(np.linalg.norm(body_bias[:2]))
        return {'interval_count':len(items),
            'a_err_from_position_m_s2':_vector_summary(ap),
            'a_err_from_velocity_m_s2':_vector_summary(av),
            'per_axis_correlation':_correlation(ap,av),
            'mean_direction_cosine':cosine,'mean_magnitude_ratio_velocity_over_position':ratio,
            'difference_velocity_minus_position_m_s2':_vector_summary(diff),
            'common_acceleration_world_m_s2':_vector_summary(common),
            'common_acceleration_body_m_s2':_vector_summary(body),
            'body_bias_stability_std_over_bias_norm':float(
                np.linalg.norm(body_std)/max(np.linalg.norm(body_bias),1e-12)),
            'constant_body_accelerometer_bias_direction_stable':bool(
                np.linalg.norm(body_std)<=np.linalg.norm(body_bias)),
            'equivalent_horizontal_tilt_rad':horizontal/G0,
            'equivalent_horizontal_tilt_deg':float(np.degrees(horizontal/G0)),
            'gravity_tilt_leakage_magnitude_le_0p5deg':bool(
                np.degrees(horizontal/G0)<=.5)}
    def grouped(key):
        groups={}
        for record in records: groups.setdefault(record[key],[]).append(record)
        return {name:summarize(items) for name,items in groups.items()}
    overall=summarize(records)
    detection_checks={}
    if records:
        mp=np.asarray(overall['a_err_from_position_m_s2']['bias'])
        mv=np.asarray(overall['a_err_from_velocity_m_s2']['bias'])
        detection_checks={'mean_direction_cosine_ge_0p95':
                overall['mean_direction_cosine']>=.95,
            'mean_magnitude_ratio_in_0p75_1p25':
                .75<=overall['mean_magnitude_ratio_velocity_over_position']<=1.25,
            'mean_acceleration_difference_norm_le_0p05_m_s2':
                np.linalg.norm(mv-mp)<=.05}
        detected=bool(all(detection_checks.values()))
    else:
        detected=False
    return {'overall':overall,'per_session':grouped('session'),
        'by_motion_class':grouped('motion'),'by_speed':grouped('speed_bin'),
        'by_gyro_norm':grouped('gyro_bin'),
        'common_acceleration_detection_gate':{
            'thresholds':{'direction_cosine_min':.95,
                'magnitude_ratio_range':[.75,1.25],
                'mean_difference_norm_max_m_s2':.05},
            'checks':detection_checks,'passed':detected},
        'common_constant_acceleration_error_detected':detected},records

def _outside_targets(t,intervals):
    return all(not (start-1.<=t<=end+1.) for start,end in intervals)

def _physical_static_biases(sessions,reference,R,heldout):
    ranges={}
    for item in heldout:
        ranges.setdefault(item['session_id'],[]).append((item['start_s'],item['end_s']))
    result={}
    for session in sessions:
        estimates=[]; times=[]
        for node in _nodes(session,reference,1.):
            if (node.zupt_static and node.hpr_factor_valid and
                    _outside_targets(node.t_s,ranges.get(session.session_id,[]))):
                R_WI=_world_rtk(node.baseline_enu)@R
                estimates.append(node.accel_m_s2-R_WI.T@(-G_ENU)); times.append(node.t_s)
        if len(estimates)>=3:
            a=np.asarray(estimates)
            result[session.session_id]={'available':True,'sample_count':len(a),
                'time_min_s':min(times),'time_max_s':max(times),
                'physical_accel_bias_m_s2':np.median(a,axis=0),
                'sample_axis_std_m_s2':np.std(a,axis=0),
                'method':('target-excluded zupt_static + fixed R2G/HPR gravity; '
                    'horizontal components remain gravity-tilt confounded')}
        else:
            result[session.session_id]={'available':False,'sample_count':len(estimates),
                'reason':'fewer than 3 target-excluded independent static nodes'}
    return result

def _bias_distribution(values,key):
    a=np.asarray([x[key] for x in values.values()],dtype=float)
    return {'session_count':len(a),'mean':np.mean(a,axis=0),'std':np.std(a,axis=0),
            'min':np.min(a,axis=0),'max':np.max(a,axis=0),
            'peak_to_peak':np.ptp(a,axis=0)}

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('--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'))
    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)}')
    lever=np.asarray(engineering['prior_constrained_solution']['result']['final_l_I_m'])
    frozen=np.array([-.4518015159,-.2644749820,.7314656115])
    if not np.allclose(lever,frozen,atol=1e-10):
        raise RuntimeError('candidate lever differs from frozen root-cause value')
    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,1.)
    R=Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    problems=[build_problem(segment,R,lever,.006) for segment in segments]
    calibration_bias=_calibration_biases(engineering,problems)
    physical_bias=_physical_static_biases(sessions,reference,R,heldout)
    motion_map={'0808_20260808_092827':'circle',
        '0808_20260808_082148':'left_right',
        '0815_20260812_123424':'slope'}

    def evaluate(mode):
        output={'best_position':[],'doppler':[],'imu_preintegration':[]}
        for problem in problems:
            session_id=problem.segment.session_id
            if mode=='calibration_frozen':
                bg=calibration_bias[session_id]['gyro_bias_rad_s']
                ba=calibration_bias[session_id]['accel_bias_m_s2']
            elif mode=='zero':
                bg=np.zeros(3); ba=np.zeros(3)
            else:
                source=physical_bias[session_id]
                if not source['available']: continue
                bg=np.zeros(3); ba=source['physical_accel_bias_m_s2']
            motion=motion_map.get(session_id,'other_recovered_dynamic')
            records=_interval_innovations(problem,motion,bg,ba)
            for key,value in records.items(): output[key].extend(value)
        root,pairs=_root_report(output['best_position'],output['doppler'])
        return {'bias_source':mode,
            'BEST_position_innovation':_factor_report(output['best_position']),
            'Doppler_innovation':_factor_report(output['doppler']),
            'IMU_preintegration_innovation':_factor_report(output['imu_preintegration']),
            'propagation_acceleration_consistency':root},pairs,output

    modes={}; pairs={}; raw={}
    for mode in ('calibration_frozen','zero','session_static_physical'):
        modes[mode],pairs[mode],raw[mode]=evaluate(mode)
    static_ids={x['interval_id'] for x in pairs['session_static_physical']}
    common_subset={}
    for mode in ('calibration_frozen','zero','session_static_physical'):
        pos=[x for x in raw[mode]['best_position'] if x['interval_id'] in static_ids]
        vel=[x for x in raw[mode]['doppler'] if x['interval_id'] in static_ids]
        root,_=_root_report(pos,vel)
        common_subset[mode]={'BEST_position_innovation':_factor_report(pos),
            'Doppler_innovation':_factor_report(vel),
            'propagation_acceleration_consistency':root}
    def bias_norm(report,factor):
        return float(np.linalg.norm(report[factor]['overall']['innovation_bias']))
    A,C=common_subset['calibration_frozen'],common_subset['session_static_physical']
    static_available=bool(static_ids)
    best_reduction=(bias_norm(C,'BEST_position_innovation')/
                    max(bias_norm(A,'BEST_position_innovation'),1e-12)
                    if static_available else np.inf)
    doppler_reduction=(bias_norm(C,'Doppler_innovation')/
                       max(bias_norm(A,'Doppler_innovation'),1e-12)
                       if static_available else np.inf)
    nuisance_transfer=bool(static_available and best_reduction<=.5 and doppler_reduction<=.5)
    high_gyro={}
    for factor in ('BEST_position_innovation','Doppler_innovation'):
        summary=modes['calibration_frozen'][factor]['by_gyro_norm'].get('gyro_ge_0p10',{})
        limit=.5
        high_gyro[factor]={'sample_count':summary.get('sample_count',0),
            'bias_norm':float(np.linalg.norm(summary.get('innovation_bias',[np.inf]*3))),
            'vector_p95':summary.get('vector_p95',np.inf),
            'passed':bool(summary.get('sample_count',0)>=20 and
                np.linalg.norm(summary['innovation_bias'])<=.10 and
                summary['vector_p95']<=limit)}
    common_detected=modes['calibration_frozen'][
        'propagation_acceleration_consistency'][
            'common_constant_acceleration_error_detected']
    extrinsic_sensitive=bool(common_detected and
                             all(x['passed'] for x in high_gyro.values()))
    propagation_passed=bool(nuisance_transfer)
    graph_bias=engineering['prior_constrained_solution'][
        'calibration_only_frozen_bias_by_session']
    graph_distribution=_bias_distribution(graph_bias,'accel_bias_m_s2')
    physical_available={k:v for k,v in physical_bias.items() if v['available']}
    physical_values={k:{'accel_bias_m_s2':v['physical_accel_bias_m_s2']}
                     for k,v in physical_available.items()}
    payload={'scope':'propagation bias root-cause only; frozen extrinsic and covariance',
        'lever_reoptimized':False,'R2G_refit':False,'covariance_retuned':False,
        'R0_parser_modified':False,'new_window_selection':False,
        'data_only_free_bootstrap_LOO_called':False,
        'fixed_l_I_m':lever,'fixed_rotation_rpy_deg':args.rotation_rpy_deg,
        'calibration_window_count':47,'heldout_window_count':267,
        'target_GNSS_observation_used_for_bias_estimation':False,
        'bias_semantics':{
            'graph_nuisance_accel_bias':(
                'node-graph nuisance absorbing IMU/model/attitude effects; '
                'not assumed transferable physical sensor zero bias'),
            'physical_IMU_accel_bias':(
                'target-excluded session static estimate; horizontal components '
                'remain gravity-tilt confounded')},
        'graph_nuisance_accel_bias_distribution_m_s2':graph_distribution,
        'per_session_graph_nuisance_bias':graph_bias,
        'per_session_static_physical_bias':physical_bias,
        'static_physical_bias_distribution_m_s2':(
            _bias_distribution(physical_values,'accel_bias_m_s2')
            if physical_values else {'session_count':0}),
        'bias_source_ablation':modes,
        'common_static_interval_subset_ablation':common_subset,
        'static_common_subset_interval_count':len(static_ids),
        'physical_over_calibration_bias_norm_ratio':{
            'BEST_position':best_reduction,'Doppler':doppler_reduction},
        'common_constant_acceleration_error_detected':common_detected,
        'nuisance_bias_transfer_failure_detected':nuisance_transfer,
        'physical_ba_ablation_available':static_available,
        'nuisance_bias_transfer_assessment':(
            'confirmed' if nuisance_transfer else
            'not_testable_no_independent_static_segments' if not static_available
            else 'not_confirmed_by_static_ablation'),
        'lever_sensitive_high_gyro_validation':high_gyro,
        'independent_propagation_validation_passed':propagation_passed,
        'independent_extrinsic_sensitive_validation_passed':extrinsic_sensitive,
        'engineering_translation_accepted':False,
        'acceptance_modified_by_this_audit':False}
    args.output.write_text(json.dumps(_jsonable(payload),ensure_ascii=False,indent=2,
                                      allow_nan=False)+'\n',encoding='utf-8')
    compact={'common_constant_acceleration_error_detected':common_detected,
        'static_session_count':len(physical_available),
        'static_common_subset_interval_count':len(static_ids),
        'bias_norm_ratio':payload['physical_over_calibration_bias_norm_ratio'],
        'nuisance_bias_transfer_failure_detected':nuisance_transfer,
        'high_gyro':high_gyro,
        'independent_propagation_validation_passed':propagation_passed,
        'independent_extrinsic_sensitive_validation_passed':extrinsic_sensitive,
        'engineering_translation_accepted':False,
        'calibration_frozen_acceleration':
            modes['calibration_frozen']['propagation_acceleration_consistency']['overall']}
    print(json.dumps(_jsonable(compact),ensure_ascii=False,indent=2))
    return 0
if __name__=='__main__': raise SystemExit(main())
