#!/usr/bin/env python3
"""
fopen_turn_finder.py — Find sharpest turns across all INTERFRAME.corr-xml sequences.

Scans every scene directory under a linked data root, parses the
*-INTERFRAME.corr-xml inside the latest version subdirectory, and
computes angular motion metrics for each sequence.

A,T,R axis convention note (Elphel / imagej-elphel):
  Originally named for a forward-looking driving camera:
    A (azimuth) = left/right steering
    T (tilt)    = uphill/downhill
    R (roll)    = rollover (wheels tipping)

  For the nadir-looking Cessna with –X along the flight direction:
    A = pitch    (nose up/down)
    T = banking  (wings tip left/right)
    R = heading  (course change — the "sharpest turn" axis)

  So to find the sharpest change of course, rank by R.
  All three axes are reported so angular dynamics are fully visible.

Velocity metrics:
  _dt-based (PRIMARY — reliable):
    The XML stores per-frame *_dt entries: the algorithm's own smooth
    polynomial derivative (dx/dt, dy/dt, dz/dt, da/dt, dt/dt, dr/dt).
    These are unaffected by single-frame pose glitches and represent
    the algorithm's best estimate of instantaneous velocity.
    max_dt_r/t/a = max |angular rate| in deg/s from _dt entries.

  finite-diff (SECONDARY — glitch indicator):
    Consecutive finite differences of the pose values, normalised by Δt.
    For clean data these agree with _dt.  Large disagreement (ratio >> 1)
    indicates pose glitches — algorithm failed, typically at low altitude
    with dense canopy where feature matching breaks down.

  glitch_ratio = max_fd_vel / max_dt_vel  (per axis).
    ~1   → clean data, no glitches
    >>1  → pose jumps present; _dt is still usable but sequence is suspect

  range = max − min angle over the full sequence (deg).
    Large range confirms a genuine manoeuvre rather than vibration.

Note on sequence length:
  Number of frames per sequence (~40–500) varies with flight altitude
  (~25 m to ~250 m); lower altitude → smaller overlap window → shorter
  sequences.  All velocity metrics are Δt-normalised and are comparable
  across sequence lengths.

Usage:
  python3 scripts/fopen_turn_finder.py \\
      --root   /path/to/linked_1763232117-1763234145-v88 \\
      [--top   20]            # top sequences to print per axis (default 20)
      [--out   turns.tsv]     # summary TSV (optional)
      [--sort  R_dt]          # sort column (default R_dt)
      [--max_glitch 5.0]      # hide sequences with glitch_ratio above this
"""

import argparse
import os
import sys
import xml.etree.ElementTree as ET
import numpy as np
import csv


# ──────────────────────────── XML parser ─────────────────────────────────── #

def parse_xml(path):
    """
    Parse one *-INTERFRAME.corr-xml.

    Returns (pose_arr, dt_arr) where each is shape (N, 7):
        [timestamp, x, y, z, a, t, r]
    sorted by timestamp.  Returns (None, None) if fewer than 2 scenes found.

    pose_arr  — absolute scene positions/orientations
    dt_arr    — per-frame velocity estimates (algorithm's smooth derivatives)
                dx/dt dy/dt dz/dt  da/dt dt/dt dr/dt  (m/s and rad/s)
    """
    PREFIX    = "EYESIS_DCT_AUX."
    SCENE_KEY = PREFIX + "scenes_"
    try:
        tree    = ET.parse(path)
        entries = {e.get("key"): e.text for e in tree.getroot().findall("entry")}
    except Exception:
        return None, None

    raw_pose = {}
    raw_dt   = {}
    for key, value in entries.items():
        if not key.startswith(SCENE_KEY):
            continue
        ts_str = key[len(SCENE_KEY):]
        is_dt  = ts_str.endswith("_dt")
        if ts_str.endswith("_d2t"):
            continue
        if is_dt:
            ts_str = ts_str[:-3]
        try:
            vals = [float(v) for v in value.split(",")]
        except (TypeError, ValueError):
            continue
        if len(vals) != 6:
            continue
        # "1763233715_106222"  →  1763233715.106222 s
        # Replace only the first underscore; the microsecond part has no underscores.
        ts = float(ts_str.replace("_", ".", 1))
        if is_dt:
            raw_dt[ts]   = vals
        else:
            raw_pose[ts] = vals

    if len(raw_pose) < 2:
        return None, None

    def to_arr(d):
        rows = sorted(d.items())
        return np.array([[ts] + vals for ts, vals in rows])

    pose_arr = to_arr(raw_pose)
    dt_arr   = to_arr(raw_dt) if len(raw_dt) >= 2 else None
    return pose_arr, dt_arr


# ─────────────────────────── metrics ─────────────────────────────────────── #

def sequence_metrics(pose_arr, dt_arr):
    """
    Compute per-sequence metrics.

    Returns dict with keys:
      n_frames, duration_s
      For each axis in (a, t, r):
        {ax}_range(deg)
        {ax}_dt_max(deg/s)    — max |_dt angular rate|  (PRIMARY, reliable)
        {ax}_dt_ts            — timestamp of that frame
        {ax}_fd_max(deg/s)    — max consecutive finite-diff angular rate
        {ax}_glitch           — fd_max / dt_max  (1=clean, >>1=glitches)
    """
    ts   = pose_arr[:, 0]
    pose = np.degrees(pose_arr[:, 4:7])   # a, t, r in degrees
    n    = len(pose_arr)

    result = {
        'n_frames':   n,
        'duration_s': float(ts[-1] - ts[0]),
    }

    axis_names = ('a', 't', 'r')

    # ── finite-diff rates (glitch-sensitive) ──────────────────────────────
    dts = np.diff(ts)
    dts = np.where(dts > 1e-9, dts, np.nan)
    fd  = np.abs(np.diff(pose, axis=0)) / dts[:, np.newaxis]  # (N-1, 3) deg/s

    for ci, ax in enumerate(axis_names):
        result[f'{ax}_range(deg)'] = float(pose[:, ci].max() - pose[:, ci].min())
        result[f'{ax}_fd_max(deg/s)'] = float(np.nanmax(fd[:, ci]))

    # ── _dt rates (smooth, algorithm's own estimate) ──────────────────────
    if dt_arr is not None and len(dt_arr) >= 2:
        # align _dt timestamps to pose timestamps (inner join)
        pose_ts_set = set(np.round(ts, 6))
        dt_matched  = dt_arr[np.array([round(t, 6) in pose_ts_set
                                       for t in dt_arr[:, 0]])]
        if len(dt_matched) >= 2:
            dt_ang = np.degrees(np.abs(dt_matched[:, 4:7]))  # deg/s, all positive
            for ci, ax in enumerate(axis_names):
                best = int(np.argmax(dt_ang[:, ci]))
                result[f'{ax}_dt_max(deg/s)'] = float(dt_ang[best, ci])
                result[f'{ax}_dt_ts']         = float(dt_matched[best, 0])
        else:
            for ax in axis_names:
                result[f'{ax}_dt_max(deg/s)'] = float('nan')
                result[f'{ax}_dt_ts']         = float('nan')
    else:
        for ax in axis_names:
            result[f'{ax}_dt_max(deg/s)'] = float('nan')
            result[f'{ax}_dt_ts']         = float('nan')

    # ── glitch ratio: fd_max / dt_max ─────────────────────────────────────
    # Values >> 1 indicate pose jumps not reflected in the smooth _dt fit.
    # A ratio of ~5 matches a single inter-frame spike (vel5 ≈ vel1/5).
    for ax in axis_names:
        fd_max = result[f'{ax}_fd_max(deg/s)']
        dt_max = result[f'{ax}_dt_max(deg/s)']
        if dt_max and dt_max > 0.01:   # avoid divide-by-near-zero
            result[f'{ax}_glitch'] = float(fd_max / dt_max)
        else:
            result[f'{ax}_glitch'] = float('nan')

    return result


# ──────────────────────────── discovery ──────────────────────────────────── #

def find_xml_files(root):
    """
    Yield (scene_ts_str, xml_path, hist_path_or_None, ims_path_or_None)
    for every scene dir under root.
    Picks the latest version subdirectory by name (e.g. v88 > v80).
    Skips directories whose name ends with '-index'.

    hist_path  — *-elevations_histogram.csv inside the version subdir (AGL)
    ims_path   — *-ims.corr-xml in the scene root dir (GPS lat/lon/alt)
    """
    try:
        scene_dirs = sorted(os.listdir(root))
    except OSError as e:
        print(f"[ERROR] cannot list {root}: {e}", file=sys.stderr)
        return

    for scene in scene_dirs:
        if scene.endswith('-index'):
            continue
        scene_path = os.path.join(root, scene)
        if not os.path.isdir(scene_path):
            continue

        # find version subdirs (e.g. v88, v80) — pick latest by name
        try:
            subdirs = [d for d in os.listdir(scene_path)
                       if os.path.isdir(os.path.join(scene_path, d))
                       and d != 'jp4']
        except OSError:
            continue
        if not subdirs:
            continue
        ver = sorted(subdirs)[-1]

        ver_path  = os.path.join(scene_path, ver)
        xml_path  = os.path.join(ver_path,   f"{scene}-INTERFRAME.corr-xml")
        hist_path = os.path.join(ver_path,   f"{scene}-elevations_histogram.csv")
        ims_path  = os.path.join(scene_path, f"{scene}-ims.corr-xml")
        if os.path.isfile(xml_path):
            yield (scene, xml_path,
                   hist_path if os.path.isfile(hist_path) else None,
                   ims_path  if os.path.isfile(ims_path)  else None)


def parse_agl(hist_path):
    """
    Read AGL (flight altitude above ground level, metres) from
    *-elevations_histogram.csv.

    Format:
      Row 1 (header): AGL  ELEVATION  FRACTION_LOWER
      Row 2:          <agl_value>  <elev>  <frac>   ← AGL is here
      Row 3+:         (empty first col, cumulative histogram continues)

    Returns float AGL, or nan if the file cannot be parsed.
    """
    try:
        with open(hist_path) as f:
            next(f)                       # skip header
            first_data = next(f).split('\t')
            return float(first_data[0])
    except Exception:
        return float('nan')


def parse_lla(ims_path):
    """
    Read latitude, longitude, altitude from *-ims.corr-xml.

    The file lives in the scene root directory (not the version subdir)
    and always contains a single representative GPS fix for the sequence.

    Key priority (most to least accurate):
      did_ins_2-lla   — INS fused GPS+IMU position  (preferred)
      did_gps1_pos-lla — raw GPS1 position           (fallback)

    Format of each entry value: "lat, lon, alt_m"

    Returns (lat, lon) as floats, or (nan, nan) on failure.
    """
    nan = float('nan')
    try:
        tree    = ET.parse(ims_path)
        entries = {e.get("key"): e.text
                   for e in tree.getroot().findall("entry")}
    except Exception:
        return nan, nan

    for key in ("did_ins_2-lla", "did_gps1_pos-lla"):
        val = entries.get(key, "")
        if val:
            try:
                parts = [float(v) for v in val.split(",")]
                if len(parts) >= 2:
                    return parts[0], parts[1]
            except ValueError:
                continue
    return nan, nan


# ─────────────────────────────── main ────────────────────────────────────── #

def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('--root',       required=True,
                    help='Linked data root')
    ap.add_argument('--top',        type=int,   default=20,
                    help='Top sequences to print per axis (default 20)')
    ap.add_argument('--out',        default=None,
                    help='Output TSV path (optional)')
    ap.add_argument('--sort',       default='R_dt',
                    choices=['R_dt','T_dt','A_dt','R_range','T_range','A_range'],
                    help='Column to rank by (default R_dt)')
    ap.add_argument('--max_glitch', type=float, default=None,
                    help='Hide sequences whose R glitch ratio exceeds this '
                         '(e.g. 5.0 keeps only clean data; default: show all)')
    ap.add_argument('--min_agl',   type=float, default=None,
                    help='Hide sequences with AGL below this value in metres '
                         '(e.g. 50 skips low-altitude failure-prone sequences)')
    args = ap.parse_args()

    print(f"Scanning {args.root} …", flush=True)
    rows = []
    for scene_ts, xml_path, hist_path, ims_path in find_xml_files(args.root):
        pose_arr, dt_arr = parse_xml(xml_path)
        if pose_arr is None:
            continue
        m = sequence_metrics(pose_arr, dt_arr)
        m['scene_ts'] = scene_ts
        m['xml_path'] = xml_path
        m['agl_m']    = parse_agl(hist_path) if hist_path else float('nan')
        lat, lon      = parse_lla(ims_path)  if ims_path  else (float('nan'), float('nan'))
        m['lat']      = lat
        m['lon']      = lon
        rows.append(m)

    print(f"Parsed {len(rows)} sequences.", flush=True)
    if not rows:
        return

    sort_map = {
        'R_dt':    'r_dt_max(deg/s)',
        'T_dt':    't_dt_max(deg/s)',
        'A_dt':    'a_dt_max(deg/s)',
        'R_range': 'r_range(deg)',
        'T_range': 't_range(deg)',
        'A_range': 'a_range(deg)',
    }
    sort_key = sort_map[args.sort]

    # apply filters
    display = rows
    if args.max_glitch is not None:
        display = [r for r in display
                   if (r.get('r_glitch') or 0) <= args.max_glitch]
        print(f"After glitch filter (r_glitch ≤ {args.max_glitch}): "
              f"{len(display)} sequences remain.")
    if args.min_agl is not None:
        display = [r for r in display
                   if np.isfinite(r.get('agl_m', float('nan')))
                   and r['agl_m'] >= args.min_agl]
        print(f"After AGL filter (≥ {args.min_agl} m): "
              f"{len(display)} sequences remain.")

    display.sort(key=lambda r: r.get(sort_key) or 0, reverse=True)

    # ── console report ────────────────────────────────────────────────────
    hdr = (f"  {'Scene':>22}  {'fr':>4}  {'AGL(m)':>7}  {'lat':>11}  {'lon':>12}  "
           f"{'range(°)':>8}  {'dt_max(°/s)':>11}  {'at timestamp':>20}  "
           f"{'glitch':>7}  {'fd_max(°/s)':>11}")

    for axis, label in (('r', 'R (heading/course)'),
                        ('t', 'T (banking)'),
                        ('a', 'A (pitch)')):
        ranked = sorted(display,
                        key=lambda r: r.get(f'{axis}_dt_max(deg/s)') or 0,
                        reverse=True)
        print(f"\n{'═'*90}")
        print(f"  {label}  —  top {args.top} by _dt angular rate  "
              f"(glitch = fd_max / dt_max; ~1 = clean)")
        print(f"{'═'*90}")
        print(hdr)
        for r in ranked[:args.top]:
            agl  = r.get('agl_m', float('nan'))
            lat  = r.get('lat',   float('nan'))
            lon  = r.get('lon',   float('nan'))
            agl_s = f"{agl:7.1f}" if np.isfinite(agl) else "    n/a"
            lat_s = f"{lat:11.6f}" if np.isfinite(lat) else "        n/a"
            lon_s = f"{lon:12.6f}" if np.isfinite(lon) else "         n/a"
            print(f"  {r['scene_ts']:>22}  {r['n_frames']:>4}  {agl_s}  "
                  f"{lat_s}  {lon_s}  "
                  f"{r[f'{axis}_range(deg)']:>8.2f}  "
                  f"{r[f'{axis}_dt_max(deg/s)']:>11.3f}  "
                  f"{r[f'{axis}_dt_ts']:>20.3f}  "
                  f"{r[f'{axis}_glitch']:>7.1f}  "
                  f"{r[f'{axis}_fd_max(deg/s)']:>11.1f}")

    # ── TSV output ────────────────────────────────────────────────────────
    if args.out:
        fieldnames = [
            'scene_ts', 'n_frames', 'duration_s', 'agl_m', 'lat', 'lon',
            'r_dt_max(deg/s)', 'r_dt_ts', 'r_range(deg)', 'r_glitch', 'r_fd_max(deg/s)',
            't_dt_max(deg/s)', 't_dt_ts', 't_range(deg)', 't_glitch', 't_fd_max(deg/s)',
            'a_dt_max(deg/s)', 'a_dt_ts', 'a_range(deg)', 'a_glitch', 'a_fd_max(deg/s)',
            'xml_path',
        ]
        sep = '\t' if args.out.endswith('.tsv') else ','
        with open(args.out, 'w', newline='') as f:
            w = csv.DictWriter(f, fieldnames=fieldnames, delimiter=sep,
                               extrasaction='ignore')
            w.writeheader()
            # write all rows (unfiltered) sorted by chosen key
            for r in sorted(rows, key=lambda r: r.get(sort_key) or 0, reverse=True):
                w.writerow(r)
        print(f"\nSaved: {args.out}  ({len(rows)} rows, unfiltered)")


if __name__ == '__main__':
    main()
