#!/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.

Metrics computed per sequence:
  range      — max minus min angle over the full sequence (deg)
  vel1       — max |Δangle| between consecutive frames / Δt  (deg/s)
               Fast but sensitive to single-frame glitches in the data.
  velW       — max |Δangle| over a W-frame window / Δt  (deg/s, default W=5)
               More robust: averages over several frames, suppresses outliers.

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.  Because velocities are normalised by elapsed time (Δt in
  seconds), the metrics are comparable across all sequence lengths.

Usage:
  python3 scripts/fopen_turn_finder.py \\
      --root  /home/elphel/lwir16-proc/NC/linked/linked_1763232117-1763234145-v88 \\
      [--window 5]          # frame offset for windowed velocity (default 5)
      [--top   20]          # how many top sequences to print per axis (default 20)
      [--out   turns.tsv]   # summary TSV (optional)
      [--sort  R_vel1]      # column to sort by (default R_vel1)
"""

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 array of shape (N, 7): [timestamp, x, y, z, a, t, r]
    sorted by timestamp.  Returns None if fewer than 2 scenes found.
    """
    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

    raw = {}
    for key, value in entries.items():
        if not key.startswith(SCENE_KEY):
            continue
        ts_str = key[len(SCENE_KEY):]
        if ts_str.endswith("_dt") or ts_str.endswith("_d2t"):
            continue
        try:
            vals = [float(v) for v in value.split(",")]
        except (TypeError, ValueError):
            continue
        if len(vals) != 6:
            continue
        # "1763233715_106222"  →  1763233715.106222 s
        ts = float(ts_str.replace("_", ".", 1))
        raw[ts] = vals

    if len(raw) < 2:
        return None

    rows = sorted(raw.items())
    arr  = np.array([[ts] + vals for ts, vals in rows])  # (N,7)
    return arr


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

def sequence_metrics(arr, window):
    """
    Compute metrics for one sequence array (N,7): [ts,x,y,z,a,t,r].

    Returns dict with keys:
      n_frames, duration_s,
      for each axis in (a, t, r):
        {axis}_range(deg)
        {axis}_vel1_max(deg/s)   — max |Δ| between consecutive frames / Δt
        {axis}_vel1_ts           — timestamp of that frame
        {axis}_velW_max(deg/s)   — max |Δ| over window frames / Δt
        {axis}_velW_ts           — timestamp of start frame
    """
    ts   = arr[:, 0]
    vals = np.degrees(arr[:, 4:7])   # a, t, r  in degrees
    n    = len(arr)

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

    axis_names = ('a', 't', 'r')
    for ci, axis in enumerate(axis_names):
        col = vals[:, ci]

        # range
        result[f'{axis}_range(deg)'] = float(col.max() - col.min())

        # consecutive velocity  (frame i → i+1)
        dt1   = np.diff(ts)
        dt1   = np.where(dt1 > 1e-9, dt1, np.nan)
        dang1 = np.abs(np.diff(col)) / dt1
        best1 = int(np.nanargmax(dang1))
        result[f'{axis}_vel1_max(deg/s)'] = float(dang1[best1])
        result[f'{axis}_vel1_ts']         = float(ts[best1])

        # windowed velocity  (frame i → i+window)
        if n > window:
            dtW   = ts[window:] - ts[:n - window]
            dtW   = np.where(dtW > 1e-9, dtW, np.nan)
            dangW = np.abs(col[window:] - col[:n - window]) / dtW
            bestW = int(np.nanargmax(dangW))
            result[f'{axis}_vel{window}_max(deg/s)'] = float(dangW[bestW])
            result[f'{axis}_vel{window}_ts']         = float(ts[bestW])
        else:
            result[f'{axis}_vel{window}_max(deg/s)'] = float('nan')
            result[f'{axis}_vel{window}_ts']         = float('nan')

    return result


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

def find_xml_files(root):
    """
    Yield (scene_ts_str, xml_path) for every scene directory under root.
    Picks the latest version subdirectory by name (e.g. v88 > v80).
    Skips directories whose name ends with '-index'.
    """
    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 the 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]   # lexicographic; "v88" > "v80" > "v20" etc.

        ver_path = os.path.join(scene_path, ver)
        xml_name = f"{scene}-INTERFRAME.corr-xml"
        xml_path = os.path.join(ver_path, xml_name)
        if os.path.isfile(xml_path):
            yield scene, xml_path


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

def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('--root',   required=True,
                    help='Linked data root, e.g. .../linked_1763232117-1763234145-v88')
    ap.add_argument('--window', type=int, default=5,
                    help='Frame offset for windowed velocity (default 5)')
    ap.add_argument('--top',    type=int, default=20,
                    help='How many 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_vel1',
                    choices=['R_range', 'R_vel1', 'R_velW',
                             'T_range', 'T_vel1', 'T_velW',
                             'A_range', 'A_vel1', 'A_velW'],
                    help='Column to rank sequences by (default R_vel1)')
    args = ap.parse_args()

    W = args.window
    print(f"Scanning {args.root} …", flush=True)

    rows = []
    for scene_ts, xml_path in find_xml_files(args.root):
        arr = parse_xml(xml_path)
        if arr is None:
            continue
        m = sequence_metrics(arr, W)
        m['scene_ts'] = scene_ts
        m['xml_path'] = xml_path
        rows.append(m)

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

    # ── sort map ──────────────────────────────────────────────────────────
    sort_col_map = {
        'R_range': f'r_range(deg)',
        'R_vel1':  f'r_vel1_max(deg/s)',
        'R_velW':  f'r_vel{W}_max(deg/s)',
        'T_range': f't_range(deg)',
        'T_vel1':  f't_vel1_max(deg/s)',
        'T_velW':  f't_vel{W}_max(deg/s)',
        'A_range': f'a_range(deg)',
        'A_vel1':  f'a_vel1_max(deg/s)',
        'A_velW':  f'a_vel{W}_max(deg/s)',
    }
    sort_key = sort_col_map[args.sort]
    rows.sort(key=lambda r: r.get(sort_key, 0) or 0, reverse=True)

    # ── console report ────────────────────────────────────────────────────
    # For each axis print top-N by vel1 and by range
    for axis, label in (('r', 'R (heading/course)'),
                        ('t', 'T (banking)'),
                        ('a', 'A (pitch)')):
        print(f"\n{'═'*70}")
        print(f"  {label}  —  top {args.top} by consecutive velocity")
        print(f"{'═'*70}")
        print(f"  {'Scene':>22}  {'frames':>6}  {'range(°)':>9}  "
              f"{'vel1(°/s)':>10}  {'at timestamp':>20}  "
              f"{'vel{W}(°/s)':>10}  {'at timestamp':>20}".replace('{W}', str(W)))
        by_vel = sorted(rows, key=lambda r: r.get(f'{axis}_vel1_max(deg/s)', 0) or 0,
                        reverse=True)
        for r in by_vel[:args.top]:
            print(f"  {r['scene_ts']:>22}  {r['n_frames']:>6}  "
                  f"{r[f'{axis}_range(deg)']:>9.2f}  "
                  f"{r[f'{axis}_vel1_max(deg/s)']:>10.3f}  "
                  f"{r[f'{axis}_vel1_ts']:>20.3f}  "
                  f"{r[f'{axis}_vel{W}_max(deg/s)']:>10.3f}  "
                  f"{r[f'{axis}_vel{W}_ts']:>20.3f}")

    # ── TSV output ────────────────────────────────────────────────────────
    if args.out:
        fieldnames = [
            'scene_ts', 'n_frames', 'duration_s',
            'r_range(deg)', 'r_vel1_max(deg/s)', 'r_vel1_ts',
            f'r_vel{W}_max(deg/s)', f'r_vel{W}_ts',
            't_range(deg)', 't_vel1_max(deg/s)', 't_vel1_ts',
            f't_vel{W}_max(deg/s)', f't_vel{W}_ts',
            'a_range(deg)', 'a_vel1_max(deg/s)', 'a_vel1_ts',
            f'a_vel{W}_max(deg/s)', f'a_vel{W}_ts',
            '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()
            w.writerows(rows)
        print(f"\nSaved: {args.out}")


if __name__ == '__main__':
    main()
