#!/usr/bin/env python3
'''Independent prediction/innovation noise audit for node-state RTK/IMU factors.'''
from __future__ import annotations
import argparse, json, math, 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.geometry import so3_log
from imu_lidar.imu_preintegration import apply_bias_correction_imu,preintegrate_gyro
from rtk_imu.rtk_imu_engineering import (
    G_ENU,HPR_DIRECT_ANGULAR_SIGMA_RAD,_all_hpr,_height_reference,
    _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 MECHANICAL_L_I_M,_jsonable
from tools.run_rtk_imu_node_graph_fixed_lever import _select_window

BEST_SIGMA = np.array([.06,.06,.12])
DOPPLER_SIGMA = np.array([.15,.15,.30])

def _whiten(value,covariance):
    covariance = .5*(covariance+covariance.T)+np.eye(len(value))*1e-12
    return np.linalg.solve(np.linalg.cholesky(covariance),value)

def _bin_speed(value):
    if value < .2: return 'speed_lt_0p2'
    if value < 1.: return 'speed_0p2_to_1'
    return 'speed_ge_1'

def _bin_gyro(value):
    if value < .02: return 'gyro_lt_0p02'
    if value < .10: return 'gyro_0p02_to_0p10'
    return 'gyro_ge_0p10'

def _bin_gap(value):
    if not np.isfinite(value): return 'gap_unavailable'
    if value <= .12: return 'gap_direct'
    if value <= .3: return 'gap_0p12_to_0p3'
    if value <= .5: return 'gap_0p3_to_0p5'
    return 'gap_gt_0p5'

def _distribution(values):
    a = np.asarray(values,dtype=float).reshape(-1)
    if not a.size:
        return {'count':0,'rms':np.nan,'p50_abs':np.nan,'p95_abs':np.nan,'p99_abs':np.nan}
    return {'count':len(a),'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.))}

def _autocorrelation(vectors):
    a = np.asarray(vectors,dtype=float)
    result = np.full(a.shape[1] if a.ndim == 2 else 0,np.nan)
    if a.ndim != 2 or len(a) < 3: return result
    for axis in range(a.shape[1]):
        left,right = a[:-1,axis],a[1:,axis]
        if np.std(left)>1e-12 and np.std(right)>1e-12:
            result[axis] = np.corrcoef(left,right)[0,1]
    return result

def _summarize(records):
    if not records: return {'sample_count':0}
    residual = np.asarray([item['residual'] for item in records])
    normalized = np.asarray([item['normalized'] for item in records])
    factor_only = np.asarray([item['normalized_factor_only'] for item in records])
    nis = np.sum(normalized*normalized,axis=1)
    dimension = residual.shape[1]
    effective_dimension = int(records[0].get('effective_dimension',dimension))
    total_dof = len(records)*effective_dimension
    norm = np.linalg.norm(residual,axis=1)
    centered_t = np.asarray([item['t_s'] for item in records],dtype=float)
    centered_t -= np.mean(centered_t)
    temporal_drift = np.zeros(dimension)
    if len(records)>=3 and np.ptp(centered_t)>1e-9:
        temporal_drift = np.asarray([
            np.polyfit(centered_t,residual[:,axis],1)[0]
            for axis in range(dimension)])
    return {'sample_count':len(records),'dimension':dimension,
        'effective_dof_per_sample':effective_dimension,
        'innovation_bias':np.mean(residual,axis=0),
        'innovation_distribution':_distribution(residual),
        'axis_rms':np.sqrt(np.mean(residual*residual,axis=0)),
        'axis_p50_abs':np.percentile(np.abs(residual),50.,axis=0),
        'axis_p95_abs':np.percentile(np.abs(residual),95.,axis=0),
        'axis_p99_abs':np.percentile(np.abs(residual),99.,axis=0),
        'vector_rms':float(np.sqrt(np.mean(norm*norm))),
        'vector_p50':float(np.percentile(norm,50.)),
        'vector_p95':float(np.percentile(norm,95.)),
        'vector_p99':float(np.percentile(norm,99.)),
        'normalized_distribution':_distribution(normalized),
        'empirical_covariance':np.cov(residual,rowvar=False),
        'nis_distribution':_distribution(nis),
        'nis_total':float(np.sum(nis)),'nis_per_dof':float(np.sum(nis)/total_dof),
        'alpha_factor':float(np.sum(nis)/total_dof),
        'alpha_composite_prediction':float(np.sum(nis)/total_dof),
        'alpha_factor_only_upper_bound':float(np.sum(factor_only*factor_only)/total_dof),
        'temporal_autocorrelation_lag1':_autocorrelation(residual),
        'temporal_linear_drift_per_s':temporal_drift}

def _group(records,key):
    values = {}
    for item in records: values.setdefault(str(item[key]),[]).append(item)
    return {name:_summarize(items) for name,items in values.items()}

def _base_record(session_id,motion,t,speed,gyro_norm,gap):
    return {'session':session_id,'motion':motion,'t_s':float(t),
            'speed_bin':_bin_speed(speed),'gyro_bin':_bin_gyro(gyro_norm),
            'hpr_gap_bin':_bin_gap(gap)}

def _interval_innovations(problem,motion,bg=None,ba=None):
    records = {'best_position':[],'doppler':[],'imu_preintegration':[]}
    lever = problem.fixed_l_I_m
    bg=np.zeros(3) if bg is None else np.asarray(bg,dtype=float)
    ba=np.zeros(3) if ba is None else np.asarray(ba,dtype=float)
    nodes = problem.segment.nodes
    for index,(left,right,pre) in enumerate(zip(nodes[:-1],nodes[1:],problem.segment.preintegrations)):
        if not (left.hpr_factor_valid and right.hpr_factor_valid): continue
        if left.velocity_enu_m_s is None or right.velocity_enu_m_s is None: continue
        R0 = _world_rtk(left.baseline_enu)@problem.R_RTK_IMU
        R1 = _world_rtk(right.baseline_enu)@problem.R_RTK_IMU
        p_i0 = left.p_enu_m-R0@lever
        p_i1 = right.p_enu_m-R1@lever
        v_i0 = left.velocity_enu_m_s-R0@np.cross(left.gyro_rad_s-bg,lever)
        v_i1 = right.velocity_enu_m_s-R1@np.cross(right.gyro_rad_s-bg,lever)
        dt = pre.duration_s
        delta_R,delta_v,delta_p=apply_bias_correction_imu(pre,bg,ba)
        pred_p_i1 = p_i0+v_i0*dt+.5*G_ENU*dt*dt+R0@delta_p
        pred_v_i1 = v_i0+G_ENU*dt+R0@delta_v
        pred_R1 = R0@delta_R
        p_innovation = pred_p_i1+pred_R1@lever-right.p_enu_m
        v_innovation = pred_v_i1+pred_R1@np.cross(
            right.gyro_rad_s-bg,lever)-right.velocity_enu_m_s
        speed = float(np.linalg.norm(left.velocity_enu_m_s))
        gyro_norm = float(pre.mean_gyro_norm)
        gap = max(left.hpr_support_gap_s,right.hpr_support_gap_s)
        base = _base_record(problem.segment.session_id,motion,right.t_s,speed,gyro_norm,gap)
        base.update({'interval_id':f'{problem.segment.segment_id}:{index}',
                     'dt_s':float(dt),'R0_WI':R0})
        p_cov = (np.diag(BEST_SIGMA**2)+dt*dt*np.diag(DOPPLER_SIGMA**2)
                 +R0@pre.cov[6:9,6:9]@R0.T+np.diag(BEST_SIGMA**2))
        v_cov = (np.diag(DOPPLER_SIGMA**2)+R0@pre.cov[3:6,3:6]@R0.T
                 +np.diag(DOPPLER_SIGMA**2))
        records['best_position'].append({**base,'residual':p_innovation,
            'normalized':_whiten(p_innovation,p_cov),
            'normalized_factor_only':p_innovation/BEST_SIGMA})
        records['doppler'].append({**base,'residual':v_innovation,
            'normalized':_whiten(v_innovation,v_cov),
            'normalized_factor_only':v_innovation/DOPPLER_SIGMA})
        imu_error = np.concatenate([
            so3_log(delta_R.T@R0.T@R1),
            R0.T@(v_i1-v_i0-G_ENU*dt)-delta_v,
            R0.T@(p_i1-p_i0-v_i0*dt-.5*G_ENU*dt*dt)-delta_p])
        anchored_cov = pre.cov.copy()
        hpr_var = left.hpr_angular_sigma_rad**2+right.hpr_angular_sigma_rad**2
        anchored_cov[:3,:3] += np.eye(3)*hpr_var
        anchored_cov[3:6,3:6] += R0.T@np.diag(2.*DOPPLER_SIGMA**2)@R0
        anchored_cov[6:9,6:9] += R0.T@np.diag(
            2.*BEST_SIGMA**2+dt*dt*DOPPLER_SIGMA**2)@R0
        records['imu_preintegration'].append({**base,'residual':imu_error,
            'normalized':_whiten(imu_error,anchored_cov),
            'normalized_factor_only':_whiten(imu_error,pre.cov)})
    return records

def _hpr_innovations(session,problem,motion):
    hpr = _all_hpr(session)
    indices = [int(i) for i in hpr.valid_indices
               if problem.segment.nodes[0].t_s <= hpr.t_s[i] <= problem.segment.nodes[-1].t_s]
    records = []
    baseline_I = problem.R_RTK_IMU.T[:,0]
    node_times = np.asarray([node.t_s for node in problem.segment.nodes])
    for position in range(1,len(indices)-1):
        left,index,right = indices[position-1],indices[position],indices[position+1]
        t0,t,t1 = hpr.t_s[left],hpr.t_s[index],hpr.t_s[right]
        total_gap = float(t1-t0)
        if not 0. < total_gap <= .5: continue
        fraction = float((t-t0)/total_gap)
        predicted = (1.-fraction)*hpr.baseline_enu[left]+fraction*hpr.baseline_enu[right]
        predicted /= np.linalg.norm(predicted)
        innovation = np.cross(predicted,hpr.baseline_enu[index])
        sigma = HPR_DIRECT_ANGULAR_SIGMA_RAD*np.sqrt(
            1.+(1.-fraction)**2+fraction**2)
        nearest = int(np.argmin(np.abs(node_times-t)))
        node = problem.segment.nodes[nearest]
        speed = float(np.linalg.norm(node.velocity_enu_m_s)) if node.velocity_enu_m_s is not None else 0.
        gyro_norm = float(np.linalg.norm(node.gyro_rad_s))
        base = _base_record(session.session_id,motion,t,speed,gyro_norm,total_gap)
        records.append({**base,'method':'leave_one_out_interpolation',
            'effective_dimension':2,'residual':innovation,
            'normalized':innovation/sigma,
            'normalized_factor_only':innovation/HPR_DIRECT_ANGULAR_SIGMA_RAD})
        gyro_pre = preintegrate_gyro(session.imu.t_s,session.imu.gyro_rad_s,float(t0),float(t))
        R0 = _world_rtk(hpr.baseline_enu[left])@problem.R_RTK_IMU
        propagated = R0@gyro_pre.delta_R@baseline_I
        gyro_innovation = np.cross(propagated,hpr.baseline_enu[index])
        gyro_sigma = np.sqrt(2.*HPR_DIRECT_ANGULAR_SIGMA_RAD**2
                             +float(np.trace(gyro_pre.cov))/3.)
        records.append({**base,'method':'gyro_propagation',
            'effective_dimension':2,'residual':gyro_innovation,
            'normalized':gyro_innovation/gyro_sigma,
            'normalized_factor_only':gyro_innovation/HPR_DIRECT_ANGULAR_SIGMA_RAD})
    return records

def _factor_report(records):
    return {'overall':_summarize(records),
        'by_session':_group(records,'session'),'by_motion':_group(records,'motion'),
        'by_speed':_group(records,'speed_bin'),'by_gyro_norm':_group(records,'gyro_bin'),
        'by_hpr_bridge_gap':_group(records,'hpr_gap_bin'),
        'by_method':_group(records,'method') if records and 'method' in records[0] else {}}

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--manifest',type=Path,required=True)
    parser.add_argument('--output',type=Path,required=True)
    parser.add_argument('--circle-session',required=True)
    parser.add_argument('--left-right-session',required=True)
    parser.add_argument('--slope-session',required=True)
    parser.add_argument('--sample-period-s',type=float,default=1.)
    parser.add_argument('--target-duration-s',type=float,default=15.)
    parser.add_argument('--rotation-rpy-deg',nargs=3,type=float,
                        default=[.4543066225,-.0026392019,.0122384129])
    parser.add_argument('--mechanical-l-I-m',nargs=3,type=float,
                        default=MECHANICAL_L_I_M.tolist())
    args = parser.parse_args()
    categories = {'circle':args.circle_session,'left_right':args.left_right_session,
                  'slope':args.slope_session}
    sessions = load_unified_sessions(args.manifest,selected_session_ids=set(categories.values()))
    reference = _height_reference(sessions)
    if reference is None: raise RuntimeError('no BEST reference')
    rotation = Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    lever = np.asarray(args.mechanical_l_I_m)
    selections, all_records = {}, {
        'best_position':[],'doppler':[],'hpr':[],'imu_preintegration':[]}
    for motion,session_id in categories.items():
        session = next(item for item in sessions if item.session_id == session_id)
        segment,selection = _select_window(
            [session],reference,args.sample_period_s,args.target_duration_s)
        selections[motion] = selection
        problem = build_problem(segment,rotation,lever)
        interval = _interval_innovations(problem,motion)
        for key,records in interval.items(): all_records[key].extend(records)
        all_records['hpr'].extend(_hpr_innovations(session,problem,motion))
    factors = {key:_factor_report(records) for key,records in all_records.items()}
    payload = {
        'scope':'independent prediction innovations; no node optimization/covariance writeback',
        'least_squares_called':False,'covariance_parameters_modified':False,
        'fixed_l_I_m':lever,'selections':selections,
        'current_physical_sigma':{
            'best_position_xyz_m':BEST_SIGMA,
            'doppler_xyz_m_s':DOPPLER_SIGMA,
            'hpr_direct_angular_rad':HPR_DIRECT_ANGULAR_SIGMA_RAD,
            'bias_random_walk_source':'unchanged device/static/Allan noise model'},
        'alpha_semantics':{
            'alpha_factor':'innovation NIS/dof using composite prediction covariance',
            'alpha_factor_only_upper_bound':(
                'innovation divided only by current factor covariance; includes predictor and '
                'endpoint-anchor noise and must not be written back directly')},
        'factors':factors,
        'recommended_covariance_writeback':False,
        'covariance_freeze_assessment':{
            'best_position':'predictor-confounded; unchanged',
            'doppler':'predictor-confounded; unchanged',
            'imu_preintegration':'endpoint-anchor dominated; unchanged',
            'bias_random_walk':'static/Allan/device-model based; unchanged',
            'hpr':{
                'identifiable':True,
                'evidence':'264 withheld samples; LOO and gyro predictions agree',
                'frozen_direct_sigma_rad':.006,
                'writeback_policy':'explicit audit decision, not automatic alpha'}}}
    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:{'count':value['overall'].get('sample_count',0),
        'bias':value['overall'].get('innovation_bias'),
        'rms':value['overall'].get('innovation_distribution',{}).get('rms'),
        'p95':value['overall'].get('innovation_distribution',{}).get('p95_abs'),
        'alpha_composite':value['overall'].get('alpha_composite_prediction'),
        'alpha_factor_only_upper_bound':value['overall'].get('alpha_factor_only_upper_bound'),
        'autocorrelation':value['overall'].get('temporal_autocorrelation_lag1')}
        for key,value in factors.items()}
    print(json.dumps(_jsonable(compact),ensure_ascii=False,indent=2))
    return 0

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