#!/usr/bin/env python3
'''Select non-overlapping qualified windows by incremental lever information.'''
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 _all_hpr,_height_reference
from rtk_imu.rtk_imu_multisource import load_unified_sessions
from rtk_imu.rtk_imu_node_graph import build_problem,linearized_lever_information
from tools.audit_rtk_imu_factor_consistency import _build_segment,_jsonable,_qualified_runs
from tools.run_rtk_imu_node_graph_fixed_lever import _select_window


STD_GATE=np.array([.15,.15,.20])


def _overlap(left,right):
    return left['session_id']==right['session_id'] and not (
        left['end_s']<right['start_s'] or right['end_s']<left['start_s'])


def _summary(information):
    covariance=np.linalg.pinv(information,rcond=1e-9)
    singular=np.linalg.svd(information,compute_uv=False)
    sign,logdet=np.linalg.slogdet(information)
    return {'std_m':np.sqrt(np.maximum(np.diag(covariance),0.)),
        'singular_values':singular,'lambda_min':singular[-1],
        'logdet':float(logdet) if sign>0 else -np.inf}


def _window_candidates(sessions,reference,period_s,duration_s):
    candidates=[]
    for session,run,segment_id in _qualified_runs(sessions,reference,period_s):
        times=np.asarray([node.t_s for node in run])
        start=0
        while start<len(run):
            end=int(np.searchsorted(times,times[start]+duration_s))
            if end>=len(run): break
            window=tuple(run[start:end+1])
            if 10.<=window[-1].t_s-window[0].t_s<=20.:
                candidate_id=f'{segment_id}:info_window_{start:04d}'
                segment=_build_segment(session,window,candidate_id)
                if segment is not None:
                    candidates.append({'candidate_id':candidate_id,
                        'session_id':session.session_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),'segment':segment,'session':session})
            start=end+1
    return candidates


def _sample_keys(entry):
    session=entry['session']; start,end=entry['start_s'],entry['end_s']
    imu=set((session.session_id,int(round(t*1e6))) for t in session.imu.t_s
            if start<=t<=end)
    gnss=set((session.session_id,int(round(node.t_s*1e6)),node.source)
             for node in entry['segment'].nodes)
    hpr=_all_hpr(session)
    hpr_keys=set((session.session_id,int(round(hpr.t_s[i]*1e6)))
                 for i in hpr.valid_indices if start<=hpr.t_s[i]<=end)
    return {'imu':imu,'gnss':gnss,'hpr':hpr_keys}


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('--window-duration-s',type=float,default=15.)
    parser.add_argument('--max-additional-windows',type=int,default=100)
    parser.add_argument('--hpr-direct-sigma-rad',type=float,default=.006)
    parser.add_argument('--linearization-l-I-m',nargs=3,type=float,
                        default=[-.53243,-.42662,.74660])
    parser.add_argument('--rotation-rpy-deg',nargs=3,type=float,
                        default=[.4543066225,-.0026392019,.0122384129])
    args=parser.parse_args()
    sessions=load_unified_sessions(args.manifest)
    reference=_height_reference(sessions)
    rotation=Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    lever=np.asarray(args.linearization_l_I_m,dtype=float)
    categories={'circle':args.circle_session,'left_right':args.left_right_session,
                'slope':args.slope_session}
    selected=[]
    for motion,session_id in categories.items():
        session=[s for s in sessions if s.session_id==session_id]
        segment,meta=_select_window(
            session,reference,args.sample_period_s,args.window_duration_s)
        selected.append({**meta,'candidate_id':meta['segment_id'],'motion':motion,
            'segment':segment,'session':session[0],'seed_window':True})
    candidates=[entry for entry in _window_candidates(
        sessions,reference,args.sample_period_s,args.window_duration_s)
        if not any(_overlap(entry,seed) for seed in selected)]
    for entry in [*selected,*candidates]:
        problem=build_problem(entry['segment'],rotation,lever,args.hpr_direct_sigma_rad)
        information,_,singular,weak=linearized_lever_information(problem,lever)
        entry['problem']=problem; entry['information']=information
        entry['single_window_singular_values']=singular
        entry['single_window_weakest_direction_I']=weak
    total=sum((entry['information'] for entry in selected),np.zeros((3,3)))
    curve=[{'window_count':len(selected),'added_candidate_id':'seed_three_motion',
            **_summary(total)}]
    remaining=list(candidates)
    saturation_count=0
    stop_reason='candidate_exhausted'
    while remaining and len(selected)<3+args.max_additional_windows:
        feasible=[entry for entry in remaining
                  if not any(_overlap(entry,item) for item in selected)]
        if not feasible: break
        ranked=[]
        for entry in feasible:
            summary=_summary(total+entry['information'])
            ranked.append((summary['lambda_min'],summary['logdet'],entry,summary))
        _,_,choice,summary=max(ranked,key=lambda item:(item[0],item[1]))
        gain=summary['lambda_min']-curve[-1]['lambda_min']
        selected.append(choice); remaining.remove(choice); total+=choice['information']
        curve.append({'window_count':len(selected),'added_candidate_id':choice['candidate_id'],
            'delta_lambda_min':gain,**summary})
        saturation_count=saturation_count+1 if gain<.01 else 0
        if np.all(np.asarray(summary['std_m'])<=STD_GATE):
            stop_reason='linearized_observability_gate_reached'; break
        if saturation_count>=3:
            stop_reason='incremental_lambda_min_gain_saturated'; break
    else:
        if len(selected)>=3+args.max_additional_windows:
            stop_reason='max_additional_windows_reached'
    seen={'imu':set(),'gnss':set(),'hpr':set()}
    duplicate={'imu':0,'gnss':0,'hpr':0}
    selected_output=[]
    for order,entry in enumerate(selected):
        keys=_sample_keys(entry)
        shared={name:len(value&seen[name]) for name,value in keys.items()}
        for name,value in keys.items():
            duplicate[name]+=shared[name]; seen[name].update(value)
        selected_output.append({'selection_order':order,
            'candidate_id':entry['candidate_id'],'session_id':entry['session_id'],
            'start_s':entry['start_s'],'end_s':entry['end_s'],
            'duration_s':entry['duration_s'],'node_count':entry['node_count'],
            'seed_window':entry.get('seed_window',False),
            'single_window_information':entry['information'],
            'single_window_singular_values':entry['single_window_singular_values'],
            'single_window_weakest_direction_I':entry['single_window_weakest_direction_I'],
            'sample_count':{name:len(value) for name,value in keys.items()},
            'shared_sample_count_with_previous':shared})
    payload={'scope':'all recovered qualified dynamics; information-based window selection',
        'manual_prior_used':False,'linearization_l_I_m':lever,
        'linearization_point_is_not_prior_factor':True,
        'candidate_count':len(candidates),'selected_window_count':len(selected),
        'selection_objective':'maximize lambda_min(H_l); break ties by logdet(H_l)',
        'std_gate_m':STD_GATE,'stop_reason':stop_reason,
        'selected_windows':selected_output,'lever_std_vs_information_curve':curve,
        'sample_overlap_audit':{'duplicate_sample_count':duplicate,
            'all_selected_windows_time_nonoverlapping_within_session':not any(
                _overlap(a,b) for i,a in enumerate(selected) for b in selected[i+1:]),
            'unique_sample_count':{name:len(value) for name,value in seen.items()}},
        'final_linearized_observable':bool(np.all(
            np.asarray(curve[-1]['std_m'])<=STD_GATE))}
    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({'candidate_count':len(candidates),
        'selected_window_count':len(selected),'stop_reason':stop_reason,
        'final_curve':curve[-1],'sample_overlap_audit':payload['sample_overlap_audit']}),
        ensure_ascii=False,indent=2))
    return 0


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