#!/usr/bin/env python3
"""
fopen_export_colmap_poses.py — Export imagej-elphel scene poses as a COLMAP
sparse model (cameras.txt / images.txt / empty points3D.txt) so COLMAP can
triangulate / densify with OUR poses held fixed — no SfM pose estimation.

This is the "step 3" tool for FOPEN: on vegetation scenes where COLMAP's own
mapper fails, run feature extraction + sequential matching as usual, then

    colmap point_triangulator \
        --database_path colmap.db --image_path images \
        --input_path  <out_dir>  --output_path <out_dir>_tri \
        --Mapper.ba_refine_focal_length 0 \
        --Mapper.ba_refine_principal_point 0 \
        --Mapper.ba_refine_extra_params 0

point_triangulator keeps the given poses fixed and only triangulates points.
The resulting model seeds patch_match_stereo / OpenMVS.

Pose convention (VERIFIED 2026-07-10 on graveyard ref 1763232420_274928:
fixed-pose triangulation gives 0.45 px mean reprojection vs 0.35 px for
COLMAP's own free mapper on the same images):

  - *-INTERFRAME.corr-xml stores per-scene {x,y,z, az,tilt,roll} relative to
    the segment reference scene.  Rotation is
    Rotation(RotationOrder.YXZ, RotationConvention.FRAME_TRANSFORM,
             az, tilt, roll)                      (ErsCorrection.java)
    applied as a world->camera coordinate operator; as an active-rotation
    matrix product that is  (Ry(az) @ Rx(tilt) @ Rz(roll)).T
  - Elphel camera axes: X right, Y up, Z toward viewer (looking along -Z).
    COLMAP camera axes:  X right, Y down, Z forward.
    Flip F = diag(1,-1,-1).
  - COLMAP world->cam:  R = F @ (Ry(az) @ Rx(tilt) @ Rz(roll)).T
    translation:        t = -R @ C,   C = {x,y,z}
  - The COLMAP world equals the segment-reference camera frame, in
    "SfM-scale" metres (LY disparity-offset scale — true metric scale comes
    from the IMS, see project notes).
  - Camera: PINHOLE, f = focalLength/pixelSize (13.329114 mm / 12 um =
    1110.7595 px for LWIR16 Boson), c = (width/2, height/2).  Use the
    *-NADIR-MERGED-RECTILINEAR.tiff renders (common radial distortion
    removed as of commit 79c105c4).

Usage:
    python3 fopen_export_colmap_poses.py \
        --xml    <path>/<ref_ts>-INTERFRAME.corr-xml \
        --out    <workspace>/sparse_ourposes \
        [--db    <workspace>/colmap.db]   # match image ids/names (else 1..N)
        [--focal 1110.7595] [--width 640] [--height 512]
"""

import argparse
import os
import sqlite3
import xml.etree.ElementTree as ET

import numpy as np


def parse_interframe_xml(path):
    PREFIX = "EYESIS_DCT_AUX.scenes_"
    entries = {e.get("key"): e.text
               for e in ET.parse(path).getroot().findall("entry")}
    rows = []
    for key, value in entries.items():
        if not key.startswith(PREFIX):
            continue
        ts = key[len(PREFIX):]
        if ts.endswith("_dt") or ts.endswith("_d2t"):
            continue
        try:
            vals = [float(v) for v in value.split(",")]
        except (TypeError, ValueError):
            continue
        if len(vals) != 6:
            continue
        rows.append((ts, vals))
    rows.sort(key=lambda r: float(r[0].replace("_", ".", 1)))
    return rows


def elphel_atr_to_colmap_R(az, tilt, roll):
    """R (world->cam, COLMAP axes) from Elphel az/tilt/roll (radians)."""
    ca, sa = np.cos(az), np.sin(az)
    ct, st = np.cos(tilt), np.sin(tilt)
    cr, sr = np.cos(roll), np.sin(roll)
    Ry = np.array([[ca, 0, sa], [0, 1, 0], [-sa, 0, ca]])
    Rx = np.array([[1, 0, 0], [0, ct, -st], [0, st, ct]])
    Rz = np.array([[cr, -sr, 0], [sr, cr, 0], [0, 0, 1]])
    F = np.diag([1.0, -1.0, -1.0])
    return F @ (Ry @ Rx @ Rz).T


def R_to_quat(R):
    qw = 0.5 * np.sqrt(max(1e-12, 1.0 + np.trace(R)))
    return np.array([qw,
                     (R[2, 1] - R[1, 2]) / (4 * qw),
                     (R[0, 2] - R[2, 0]) / (4 * qw),
                     (R[1, 0] - R[0, 1]) / (4 * qw)])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--xml",   required=True)
    ap.add_argument("--out",   required=True)
    ap.add_argument("--db",    default=None,
                    help="colmap.db to map image names -> ids (names <ts>.png)")
    ap.add_argument("--focal", type=float, default=1110.7595)
    ap.add_argument("--width", type=int, default=640)
    ap.add_argument("--height", type=int, default=512)
    args = ap.parse_args()

    rows = parse_interframe_xml(args.xml)
    if not rows:
        raise SystemExit(f"no scene poses found in {args.xml}")

    name2id = None
    if args.db:
        db = sqlite3.connect(args.db)
        name2id = {n: i for i, n in
                   db.execute("SELECT image_id, name FROM images")}

    os.makedirs(args.out, exist_ok=True)
    with open(os.path.join(args.out, "cameras.txt"), "w") as f:
        f.write(f"1 PINHOLE {args.width} {args.height} "
                f"{args.focal} {args.focal} "
                f"{args.width / 2:.1f} {args.height / 2:.1f}\n")
    open(os.path.join(args.out, "points3D.txt"), "w").close()

    n_written = 0
    with open(os.path.join(args.out, "images.txt"), "w") as f:
        for i, (ts, (x, y, z, az, tilt, roll)) in enumerate(rows):
            name = f"{ts}.png"
            if name2id is not None:
                if name not in name2id:
                    continue
                image_id = name2id[name]
            else:
                image_id = i + 1
            R = elphel_atr_to_colmap_R(az, tilt, roll)
            t = -R @ np.array([x, y, z])
            q = R_to_quat(R)
            q /= np.linalg.norm(q)
            f.write(f"{image_id} {q[0]} {q[1]} {q[2]} {q[3]} "
                    f"{t[0]} {t[1]} {t[2]} 1 {name}\n\n")
            n_written += 1
    print(f"wrote {n_written} poses to {args.out}")


if __name__ == "__main__":
    main()
