#!/usr/bin/env python3
'''Run one 10-20 s per-node state graph with a fixed mechanical lever.'''
from __future__ import annotations
import argparse, json, sys
from dataclasses import asdict
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, solve_fixed_lever
from tools.audit_rtk_imu_factor_consistency import (
    MECHANICAL_L_I_M, _build_segment, _jsonable, _qualified_runs)

def _select_window(sessions,reference,period,target_duration):
    best = None
    for session,run,segment_id in _qualified_runs(sessions,reference,period):
        for start in range(len(run)):
            target_t = run[start].t_s+target_duration
            end = int(np.searchsorted([node.t_s for node in run],target_t))
            if end >= len(run): continue
            window = tuple(run[start:end+1])
            duration = window[-1].t_s-window[0].t_s
            if not 10. <= duration <= 20.: continue
            best_count = sum(node.source == 'BESTNAVA' for node in window)
            doppler_count = sum(node.velocity_enu_m_s is not None for node in window)
            hpr_count = sum(node.hpr_factor_valid for node in window)
            gyro_score = sum(np.linalg.norm(node.gyro_rad_s) for node in window)
            score = 10.*best_count+10.*doppler_count+2.*hpr_count+gyro_score
            candidate = (score,session,window,f'{segment_id}:window_{start:03d}',
                         best_count,doppler_count,hpr_count)
            if best is None or candidate[0] > best[0]: best = candidate
    if best is None: raise RuntimeError('no 10-20 s qualified window')
    score,session,window,segment_id,best_count,doppler_count,hpr_count = best
    segment = _build_segment(session,window,segment_id)
    if segment is None: raise RuntimeError('selected window preintegration failed')
    return segment,{'score':score,'session_id':session.session_id,
        'segment_id':segment_id,'start_s':window[0].t_s,'end_s':window[-1].t_s,
        'duration_s':window[-1].t_s-window[0].t_s,'node_count':len(window),
        'best_count':best_count,'doppler_count':doppler_count,'hpr_factor_count':hpr_count}

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('--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('--max-nfev',type=int,default=30)
    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()
    sessions = load_unified_sessions(args.manifest,selected_session_ids={args.session})
    reference = _height_reference(sessions)
    if reference is None: raise RuntimeError('no BEST reference')
    segment,selection = _select_window(
        sessions,reference,args.sample_period_s,args.target_duration_s)
    rotation = Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    problem = build_problem(segment,rotation,np.asarray(args.mechanical_l_I_m))
    result = solve_fixed_lever(problem,args.max_nfev)
    payload = {'scope':'single 10-20 s node-state graph; fixed mechanical lever',
        'translation_variable_enabled':False,'free_prior_loo_bootstrap_sensitivity_called':False,
        'selection':selection,'result':asdict(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(payload),ensure_ascii=False,indent=2))
    return 0

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