#!/usr/bin/env python3
'''Refine non-overlapping windows after comparing predicted and actual 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 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,linearized_lever_information
from tools.audit_rtk_imu_factor_consistency import _jsonable
from tools.run_rtk_imu_node_graph_free_selected import _restore_segments
from tools.select_rtk_imu_windows_by_lever_information import (
    STD_GATE,_overlap,_sample_keys,_summary,_window_candidates)


def _sqrt_psd(matrix,inverse=False):
    values,vectors=np.linalg.eigh(.5*(matrix+matrix.T))
    values=np.maximum(values,1e-12)
    scale=1./np.sqrt(values) if inverse else np.sqrt(values)
    return (vectors*scale)@vectors.T


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--manifest',type=Path,required=True)
    parser.add_argument('--base-selection',type=Path,required=True)
    parser.add_argument('--actual-result',type=Path,required=True)
    parser.add_argument('--output',type=Path,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=400)
    parser.add_argument('--hpr-direct-sigma-rad',type=float,default=.006)
    parser.add_argument('--rotation-rpy-deg',nargs=3,type=float,
                        default=[.4543066225,-.0026392019,.0122384129])
    args=parser.parse_args()
    base=json.loads(args.base_selection.read_text(encoding='utf-8'))
    actual=json.loads(args.actual_result.read_text(encoding='utf-8'))
    actual_solution=next(iter(actual['solutions'].values()))
    lever=np.asarray(actual_solution['final_l_I_m'],dtype=float)
    actual_cov=np.asarray(actual_solution['lever_covariance_m2'],dtype=float)
    actual_information=np.linalg.pinv(actual_cov,rcond=1e-9)
    predicted_information=sum((np.asarray(item['single_window_information'])
        for item in base['selected_windows']),np.zeros((3,3)))
    correction=_sqrt_psd(actual_information)@_sqrt_psd(predicted_information,True)
    sessions=load_unified_sessions(args.manifest)
    reference=_height_reference(sessions)
    rotation=Rotation.from_euler('xyz',args.rotation_rpy_deg,degrees=True).as_matrix()
    restored=_restore_segments(sessions,reference,base['selected_windows'],
                               args.sample_period_s)
    session_by_id={session.session_id:session for session in sessions}
    selected=[]
    for item,segment in zip(base['selected_windows'],restored):
        selected.append({**item,'segment':segment,
            'session':session_by_id[item['session_id']]})
    candidates=[entry for entry in _window_candidates(
        sessions,reference,args.sample_period_s,args.window_duration_s)
        if not any(_overlap(entry,item) for item in selected)]
    for entry in candidates:
        problem=build_problem(entry['segment'],rotation,lever,args.hpr_direct_sigma_rad)
        raw,_,singular,weak=linearized_lever_information(problem,lever)
        entry['raw_information']=raw
        entry['information']=correction@raw@correction.T
        entry['single_window_singular_values']=singular
        entry['single_window_weakest_direction_I']=weak
    total=actual_information.copy()
    curve=[{'window_count':len(selected),'added_candidate_id':'actual_base',
            'information_source':'actual_joint_free_schur',**_summary(total)}]
    remaining=list(candidates); saturation_count=0
    stop_reason='candidate_exhausted'
    while remaining and len(selected)<len(base['selected_windows'])+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,'information_source':'actual-calibrated prediction',
            **summary})
        saturation_count=saturation_count+1 if gain<.01 else 0
        if np.all(np.asarray(summary['std_m'])<=STD_GATE):
            stop_reason='actual_calibrated_observability_gate_reached'; break
        if saturation_count>=3:
            stop_reason='incremental_lambda_min_gain_saturated'; break
    else:
        if len(selected)>=len(base['selected_windows'])+args.max_additional_windows:
            stop_reason='max_additional_windows_reached'
    seen={'imu':set(),'gnss':set(),'hpr':set()}; duplicate={key:0 for key in seen}
    output=[]
    for order,entry in enumerate(selected):
        keys=_sample_keys(entry); shared={key:len(value&seen[key]) for key,value in keys.items()}
        for key,value in keys.items(): duplicate[key]+=shared[key]; seen[key].update(value)
        information=np.asarray(entry['information'] if 'information' in entry
                               else entry['single_window_information'])
        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':order<len(base['selected_windows']),
            'single_window_information':information,
            'raw_single_window_information':entry.get('raw_information'),
            'sample_count':{key:len(value) for key,value in keys.items()},
            'shared_sample_count_with_previous':shared})
    overlap={'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':{key:len(value) for key,value in seen.items()}}
    payload={'scope':'actual-H calibrated incremental lever-information selection',
        'manual_prior_used':False,'base_window_count':len(base['selected_windows']),
        'candidate_count':len(candidates),'selected_window_count':len(selected),
        'actual_base_l_I_m':lever,'actual_base_information':actual_information,
        'predicted_base_information':predicted_information,
        'information_congruence_correction':correction,
        'selection_objective':'maximize lambda_min(H_l); break ties by logdet(H_l)',
        'std_gate_m':STD_GATE,'stop_reason':stop_reason,
        'selected_windows':output,'lever_std_vs_information_curve':curve,
        'sample_overlap_audit':overlap,
        'final_calibrated_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({'selected_window_count':len(selected),
        'additional_window_count':len(selected)-len(base['selected_windows']),
        'stop_reason':stop_reason,'final_curve':curve[-1],
        'sample_overlap_audit':overlap}),ensure_ascii=False,indent=2))
    return 0


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