#!/usr/bin/env python3
'''Audit RTK/IMU factor conventions without running an optimizer.'''
from __future__ import annotations
import argparse, json, math, sys
from dataclasses import asdict, is_dataclass
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.imu_preintegration import preintegrate_imu
from rtk_imu.rtk_imu_engineering import (
    G_ENU, MIN_SEGMENT_DURATION_S, MIN_SEGMENT_NODE_COUNT,
    NUISANCE_DOF_PER_SEGMENT, _Segment, _all_hpr, _audit,
    _dense_colored_jacobian, _enu, _height_reference,
    _hpr_factor_observation, _initial_parameters, _marginal_lever_information,
    _motion_flags, _nodes, _position_valid, _residual,
    _segment_residual_size, _world_rtk)
from rtk_imu.rtk_imu_multisource import _f, _truth, load_unified_sessions
MECHANICAL_L_I_M = np.array([-0.45072, -0.25682, 0.73208])

def _jsonable(value):
    if isinstance(value, np.ndarray): return _jsonable(value.tolist())
    if isinstance(value, np.generic): return _jsonable(value.item())
    if isinstance(value, float): return value if math.isfinite(value) else None
    if is_dataclass(value): return _jsonable(asdict(value))
    if isinstance(value, dict): return {str(k): _jsonable(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)): return [_jsonable(v) for v in value]
    return value

def _summary(error):
    a = np.asarray(error, dtype=float)
    return _jsonable(_audit(list(a.reshape(-1, 3)), 3)) if a.size else _jsonable(_audit([], 3))

def _corr(a, b):
    out = np.full(3, np.nan)
    for axis in range(3):
        if len(a) >= 3 and np.std(a[:, axis]) > 1e-10 and np.std(b[:, axis]) > 1e-10:
            out[axis] = np.corrcoef(a[:, axis], b[:, axis])[0, 1]
    return out

def _best_arrays(session, reference):
    rows = []
    for row in session.rtk_by_type.get('BESTNAVA', []):
        v = np.array([_f(row, 'velocity_east_m_s'), _f(row, 'velocity_north_m_s'),
                      _f(row, 'vertical_speed_m_s')])
        if _position_valid(row, 'BESTNAVA') and _truth(row, 'doppler_velocity_valid') and np.all(np.isfinite(v)):
            rows.append(row)
    rows.sort(key=lambda row: _f(row, 't_device_s'))
    rows = [row for i, row in enumerate(rows) if i == 0 or _f(row, 't_device_s') > _f(rows[i-1], 't_device_s')]
    t = np.asarray([_f(row, 't_device_s') for row in rows])
    p = np.asarray([_enu(row, 'BESTNAVA', reference)[0] for row in rows]).reshape(-1, 3)
    v = np.asarray([[_f(row, 'velocity_east_m_s'), _f(row, 'velocity_north_m_s'),
                     _f(row, 'vertical_speed_m_s')] for row in rows]).reshape(-1, 3)
    return rows, t, p, v

def _gnss_audit(sessions, reference, max_dt):
    all_dpdt, all_v, per_session = [], [], {}
    for session in sessions:
        _, t, p, v = _best_arrays(session, reference)
        dt = np.diff(t); keep = (dt > 0) & (dt <= max_dt)
        dpdt = np.diff(p, axis=0)[keep] / dt[keep, None]
        v_avg = .5 * (v[:-1] + v[1:])[keep]
        all_dpdt.append(dpdt); all_v.append(v_avg)
        per_session[session.session_id] = {
            'interval_count': len(dpdt), 'nominal_correlation_xyz': _corr(dpdt, v_avg),
            'nominal_error_m_s': _summary(dpdt-v_avg)}
    dpdt, velocity = np.vstack(all_dpdt), np.vstack(all_v)
    hypotheses = []
    for swap in (False, True):
        base = velocity[:, [1,0,2]] if swap else velocity
        for sx in (-1.,1.):
            for sy in (-1.,1.):
                for sz in (-1.,1.):
                    transformed = base * [sx,sy,sz]
                    hypotheses.append({
                        'mapping': ('[N,E,Z]' if swap else '[E,N,Z]')+f'*[{sx:+.0f},{sy:+.0f},{sz:+.0f}]',
                        'correlation_xyz': _corr(dpdt, transformed),
                        'error_m_s': _summary(dpdt-transformed)})
    hypotheses.sort(key=lambda item: item['error_m_s']['vector_rms'])
    nominal = next(x for x in hypotheses if x['mapping'] == '[E,N,Z]*[+1,+1,+1]')
    return {'interval_count': len(dpdt), 'nominal': nominal, 'best_mapping': hypotheses[0],
            'all_hypotheses': hypotheses, 'per_session': per_session}

def _nearest_imu(session, t, tolerance=.03):
    right = int(np.searchsorted(session.imu.t_s, t))
    candidates = [i for i in (right-1,right) if 0 <= i < len(session.imu.t_s)]
    if not candidates: return None
    index = min(candidates, key=lambda i: abs(session.imu.t_s[i]-t))
    return index if abs(session.imu.t_s[index]-t) <= tolerance else None

def _static_audit(sessions, rotation):
    current, opposite, per_session = [], [], {}
    for session in sessions:
        hpr, local = _all_hpr(session), []
        last_t = -np.inf
        for hpr_index in hpr.valid_indices:
            t = float(hpr.t_s[hpr_index])
            if t-last_t < 1.: continue
            last_t = t
            gravity, _ = _motion_flags(session,t,np.zeros(0),np.zeros((0,3)))
            baseline, _, valid, _, _ = _hpr_factor_observation(hpr,t)
            imu_index = _nearest_imu(session,t)
            if not (gravity and valid and imu_index is not None):
                continue
            R_WI = _world_rtk(baseline) @ rotation
            accel = session.imu.acc_m_s2[imu_index]
            value = R_WI @ accel + G_ENU
            local.append(value); current.append(value); opposite.append(R_WI @ accel-G_ENU)
        per_session[session.session_id] = {'count':len(local),'current_formula_m_s2':_summary(local)}
    return {'formula':'a_W_linear = R_WI @ specific_force_I + G_ENU',
            'current_formula_m_s2':_summary(current),
            'opposite_gravity_sign_m_s2':_summary(opposite),'per_session':per_session}

def _closure_audit(sessions, reference, rotation, lever, max_dt):
    all_p, all_v, per_session = [], [], {}
    for session in sessions:
        _, t, p_ant, v_ant = _best_arrays(session, reference)
        hpr, local_p, local_v = _all_hpr(session), [], []
        for index, dt in enumerate(np.diff(t)):
            if not .5 <= dt <= max_dt: continue
            baseline, _, valid, _, _ = _hpr_factor_observation(hpr, t[index])
            i0, i1 = _nearest_imu(session, t[index]), _nearest_imu(session, t[index+1])
            if not valid or i0 is None or i1 is None: continue
            pre = preintegrate_imu(session.imu.t_s, session.imu.gyro_rad_s,
                                   session.imu.acc_m_s2, t[index], t[index+1])
            if pre.duration_s <= 0 or abs(pre.duration_s-dt) > 1e-6: continue
            R0 = _world_rtk(baseline) @ rotation
            p_i0 = p_ant[index] - R0 @ lever
            v_i0 = v_ant[index] - R0 @ np.cross(session.imu.gyro_rad_s[i0], lever)
            p_i1 = p_i0 + v_i0*dt + .5*G_ENU*dt**2 + R0@pre.delta_p
            v_i1 = v_i0 + G_ENU*dt + R0@pre.delta_v
            R1 = R0 @ pre.delta_R
            local_p.append(p_i1 + R1@lever - p_ant[index+1])
            local_v.append(v_i1 + R1@np.cross(session.imu.gyro_rad_s[i1],lever) - v_ant[index+1])
        all_p.extend(local_p); all_v.extend(local_v)
        per_session[session.session_id] = {
            'interval_count':len(local_p),'position_m':_summary(local_p),
            'velocity_m_s':_summary(local_v)}
    return {'method':'single-step forward closure; fixed R2G/mechanical lever/BEST p-v; no least_squares',
            'position_m':_summary(all_p),'velocity_m_s':_summary(all_v),
            'per_session':per_session}

def _qualified_runs(sessions, reference, period):
    result = []
    for session in sessions:
        nodes, start, number = _nodes(session, reference, period), 0, 0
        for end in range(1,len(nodes)+1):
            if end != len(nodes) and nodes[end].continuity_id == nodes[end-1].continuity_id:
                continue
            run, start = tuple(nodes[start:end]), end
            if (len(run) >= MIN_SEGMENT_NODE_COUNT
                    and run[-1].t_s-run[0].t_s >= MIN_SEGMENT_DURATION_S
                    and any(node.hpr_factor_valid for node in run)):
                result.append((session,run,f'{session.session_id}:{number:02d}'))
                number += 1
    return result

def _first_node_audit(runs, rotation, lever):
    records, ranges = [], {}
    for session, run, segment_id in runs:
        node = run[0]
        hpr_node = next(item for item in run if item.hpr_factor_valid)
        R_WI = _world_rtk(hpr_node.baseline_enu) @ rotation
        velocity = node.velocity_enu_m_s
        v = np.full(3,np.nan) if velocity is None else velocity
        records.append({
            'segment_id':segment_id,'source':node.source,'t_s':node.t_s,
            'first_position_global_enu_m':node.p_enu_m,'first_velocity_enu_m_s':v,
            'free_x0_l0_position_residual_m':np.zeros(3),
            'free_x0_l0_velocity_residual_m_s':-v,
            'mechanical_l_unshifted_x0_position_residual_m':R_WI@lever,
            'mechanical_l_unshifted_x0_velocity_residual_m_s':R_WI@np.cross(node.gyro_rad_s,lever)-v,
            'observation_seeded_mechanical_position_residual_m':np.zeros(3),
            'observation_seeded_mechanical_velocity_residual_m_s':
                np.zeros(3) if velocity is not None else np.full(3,np.nan)})
        ranges.setdefault(session.session_id,[]).append(node.p_enu_m)
    ranges = {key:{'count':len(value),'min_global_enu_m':np.min(value,axis=0),
                   'max_global_enu_m':np.max(value,axis=0)} for key,value in ranges.items()}
    return {'common_global_enu_reference':True,'segment_state_is_global_imu_position':True,
            'per_session_first_node_ranges':ranges,'segments':records}

def _gyro_score(session, run, category):
    keep = (session.imu.t_s >= run[0].t_s) & (session.imu.t_s <= run[-1].t_s)
    t, gyro = session.imu.t_s[keep], session.imu.gyro_rad_s[keep]
    if len(t) < 2: return 0.
    value = np.trapezoid(np.abs(gyro),t,axis=0)
    return float(value[2] if category != 'slope' else np.hypot(value[0],value[1]))

def _build_segment(session, run, segment_id):
    pre = tuple(preintegrate_imu(session.imu.t_s,session.imu.gyro_rad_s,
                                 session.imu.acc_m_s2,a.t_s,b.t_s)
                for a,b in zip(run[:-1],run[1:]))
    hpr = next(node for node in run if node.hpr_factor_valid)
    return _Segment(segment_id,session.session_id,run,pre,_world_rtk(hpr.baseline_enu))

def _subset_marginal(jacobian,residual,segments,indices):
    rows, row0 = [], 0
    for index,segment in enumerate(segments):
        count = _segment_residual_size(segment)
        if index in indices: rows.extend(range(row0,row0+count))
        row0 += count
    columns = [0,1,2]
    for index in indices:
        start = 3 + NUISANCE_DOF_PER_SEGMENT*index
        columns.extend(range(start,start+NUISANCE_DOF_PER_SEGMENT))
    rows, columns = np.asarray(rows), np.asarray(columns)
    return _marginal_lever_information(
        jacobian[np.ix_(rows,columns)],residual[rows])[0]

def _jacobian_schur_audit(runs,categories,rotation,lever):
    segments, indices = [], {}
    for category in ('circle','left_right','slope'):
        choices = [item for item in runs if categories[item[0].session_id] == category]
        if not choices:
            indices[category] = []; continue
        best = max(choices,key=lambda item:_gyro_score(item[0],item[1],category))
        indices[category] = [len(segments)]
        segments.append(_build_segment(*best))
    x = _initial_parameters(segments); x[:3] = lever
    residual = _residual(x,segments,rotation)
    jacobian = _dense_colored_jacobian(x,segments,rotation)
    comparisons = []
    for axis in range(3):
        plus, minus = x.copy(), x.copy()
        plus[axis] += .001; minus[axis] -= .001
        numeric = (_residual(plus,segments,rotation)-_residual(minus,segments,rotation))/.002
        current = jacobian[:,axis]; delta = numeric-current
        comparisons.append({
            'axis':'XYZ'[axis],'perturbation_m':.001,
            'relative_difference':float(np.linalg.norm(delta)/max(np.linalg.norm(numeric),1e-12)),
            'difference_norm':float(np.linalg.norm(delta)),
            'max_absolute_difference':float(np.max(np.abs(delta))),
            'correlation':float(np.corrcoef(numeric,current)[0,1])})
    full = _marginal_lever_information(jacobian,residual)[0]
    parts = {key:_subset_marginal(jacobian,residual,segments,value)
             for key,value in indices.items() if value}
    summed = sum(parts.values(),np.zeros((3,3)))
    error = np.linalg.norm(summed-full,ord='fro')
    return {'linearization':'same observation-seeded nuisance state and mechanical lever; no optimizer',
            'selected_segments':[segment.segment_id for segment in segments],
            'finite_difference_vs_solver_jacobian':comparisons,
            'full_marginal_information':full,'category_marginal_information':parts,
            'category_sum':summed,'additivity_error_fro':float(error),
            'additivity_relative_error':float(error/max(np.linalg.norm(full,ord='fro'),1e-12))}

def _legacy_diagnostics(path):
    if path is None or not path.exists():
        return {'available':False,'reason':'no prior artifact supplied'}
    payload = json.loads(path.read_text(encoding='utf-8'))
    final_l = np.asarray(payload.get('free_solution',{}).get('l_I_m',[np.nan]*3))
    return {'available':False,'source':str(path),'initial_cost':None,'final_cost':None,
            'cost_reduction':None,'nfev':None,'optimality':None,'gradient_norm':None,
            'initial_l_I_m':[0.,0.,0.],'final_l_I_m':final_l,
            'l_step_norm_m':float(np.linalg.norm(final_l)) if np.all(np.isfinite(final_l)) else None,
            'reason':'legacy artifact did not persist optimizer state; prohibited solve was not rerun'}

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('--max-best-interval-s',type=float,default=1.75)
    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())
    parser.add_argument('--previous-free-result',type=Path)
    parser.add_argument('--level-static-session',action='append',
                        default=['0819_20260819_072130','0819_20260819_073045'])
    args = parser.parse_args()
    categories = {args.circle_session:'circle',args.left_right_session:'left_right',
                  args.slope_session:'slope'}
    selected_ids = set(categories) | set(args.level_static_session)
    all_sessions = load_unified_sessions(args.manifest,selected_session_ids=selected_ids)
    sessions = [session for session in all_sessions if session.session_id in categories]
    static_sessions = [session for session in all_sessions
                       if session.session_id in set(args.level_static_session)]
    reference = _height_reference(sessions)
    if reference is None: raise RuntimeError('no valid BESTNAVA height reference')
    rotation = Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    lever = np.asarray(args.mechanical_l_I_m)
    runs = _qualified_runs(sessions,reference,args.sample_period_s)
    payload = {
        'scope':'factor consistency only; no free solve/LOO/prior/bootstrap/sensitivity',
        'least_squares_called':False,'rotation_source':'R2G_gravity_level_prior',
        'rotation_rpy_deg':args.rotation_rpy_deg,'mechanical_l_I_m':lever,
        'common_enu_reference':reference,
        'gnss_position_difference_vs_doppler':_gnss_audit(
            sessions,reference,args.max_best_interval_s),
        'static_specific_force_gravity_sign':_static_audit(static_sessions,rotation),
        'one_step_imu_preintegration_closure':_closure_audit(
            sessions,reference,rotation,lever,args.max_best_interval_s),
        'first_node_origin_and_initial_residual':_first_node_audit(runs,rotation,lever),
        'lever_jacobian_and_category_schur':_jacobian_schur_audit(
            runs,categories,rotation,lever),
        'previous_free_fit_solver_diagnostics':_legacy_diagnostics(args.previous_free_result)}
    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({
        'gnss':payload['gnss_position_difference_vs_doppler'],
        'static':payload['static_specific_force_gravity_sign'],
        'closure':payload['one_step_imu_preintegration_closure'],
        'jacobian_schur':payload['lever_jacobian_and_category_schur']}),
        ensure_ascii=False,indent=2))
    return 0

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