#!/usr/bin/env python3
"""
fopen_dem_match.py — register per-segment terrain models (*_SZXY.tiff) against
a USGS 3DEP DEM and report per-segment vertical offset, horizontal shift and
SCALE. Absolute calibration of the LY disparity offset and an IMS-error trace,
using public LiDAR as ground truth.

Inputs per segment (from a *-ground_planes.csv index row): Latitude, Longitude,
Altitude (MSL-ish, treated as approximate), qenu[0..3]; and
<linked_root>/<ts>/v88/<ts>_SZXY.tiff with slices:
  1: S  per-tile confidence 0..1
  2: Z  distance below camera (positive down), metres (system scale)
  3: X  per-tile world X (Elphel reference-camera frame)
  4: Y  per-tile world Y
Camera-frame point: p = (X, Y, -Z).

Frame chain (validated 2026-07-11 on segment 1763232305_820138 and the pond at
1763233023_625897): ENU = qR(qenu) @ M @ p, with the fixed axis adapter
M = [[0,0,1],[-1,0,0],[0,-1,0]]  (perm (2,0,1), signs (1,-1,-1)) found by
minimizing DEM residuals over all 24 proper axis mappings.

Fit per segment: vertical offset dz always; horizontal (dx,dy) and scale s by
coarse-to-fine search, only meaningful where the DEM has relief (reported).

CAVEATS (pond segment 1763233023_625897, 2026-07-11): --fit-scale requires
STABLE relief. The robust trim protects dz (0.17 m rms at s=1) but lets the
scale search trim away changed relief (a 2020 pond, hydro-flattened to its
water surface, dry ~3 m in 2026) and then s is degenerate with dz on the flat
remainder - the fitted s=0.82/rms=0.05 there is meaningless. Use scale only
where dem_relief is high AND the relief is temporally stable; otherwise trust
s=1.0 rows. dz doubles as an IMS-altitude-error trace (observed -43 m to
+19 m between passes of the same flight).

Usage:
  python3 fopen_dem_match.py --ground-planes <index>/v88/<ts>-ground_planes.csv \
      --linked-root <...>/linked_..._v88 \
      --dem USGS_1M_*.tif --dem-origin 659993.99997 3860006.00005 \
      [--out table.tsv] [--seg <ts substring>] [--fit-scale]
"""
import argparse, csv, os, sys
import numpy as np
from PIL import Image
Image.MAX_IMAGE_PIXELS = 300_000_000

M_ADAPT = np.array([[0., 0., 1.], [-1., 0., 0.], [0., -1., 0.]])


def qR(w, x, y, z):
    return np.array([
        [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
        [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
        [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)]])


def utm17(lat, lon):
    a = 6378137.0; f = 1 / 298.257223563; k0 = 0.9996
    e2 = f * (2 - f); ep2 = e2 / (1 - e2); lam0 = np.radians(-81.0)
    phi = np.radians(lat); lam = np.radians(lon)
    N = a / np.sqrt(1 - e2 * np.sin(phi) ** 2)
    T = np.tan(phi) ** 2; C = ep2 * np.cos(phi) ** 2
    A = (lam - lam0) * np.cos(phi)
    Mm = a * ((1 - e2 / 4 - 3 * e2 * e2 / 64 - 5 * e2 ** 3 / 256) * phi
              - (3 * e2 / 8 + 3 * e2 * e2 / 32 + 45 * e2 ** 3 / 1024) * np.sin(2 * phi)
              + (15 * e2 * e2 / 256 + 45 * e2 ** 3 / 1024) * np.sin(4 * phi)
              - (35 * e2 ** 3 / 3072) * np.sin(6 * phi))
    E = k0 * N * (A + (1 - T + C) * A ** 3 / 6
                  + (5 - 18 * T + T * T + 72 * C - 58 * ep2) * A ** 5 / 120) + 500000
    Nn = k0 * (Mm + N * np.tan(phi) * (A * A / 2 + (5 - T + 9 * C + 4 * C * C) * A ** 4 / 24
               + (61 - 58 * T + T * T + 600 * C - 330 * ep2) * A ** 6 / 720))
    return E, Nn


def read_szxy(path):
    im = Image.open(path)
    sl = []
    for i in range(im.n_frames):
        im.seek(i)
        sl.append(np.array(im, dtype=np.float32))
    S, Z, X, Y = sl[:4]
    m = np.isfinite(Z) & np.isfinite(X) & np.isfinite(Y) & np.isfinite(S) & (S > 0.1)
    P = np.stack([X[m], Y[m], -Z[m]], 1).astype(np.float64)
    return P, S[m].astype(np.float64)


def sample(dem, px, py):
    ok = (px > 1) & (px < dem.shape[1] - 2) & (py > 1) & (py < dem.shape[0] - 2)
    z = np.full(len(px), np.nan)
    if ok.sum():
        x0 = np.floor(px[ok]).astype(int); y0 = np.floor(py[ok]).astype(int)
        fx = px[ok] - x0; fy = py[ok] - y0
        z[ok] = (dem[y0, x0] * (1 - fx) * (1 - fy) + dem[y0, x0 + 1] * fx * (1 - fy)
                 + dem[y0 + 1, x0] * (1 - fx) * fy + dem[y0 + 1, x0 + 1] * fx * fy)
    z[z < -100] = np.nan
    return z


def fit_segment(P, W, R, E0, N0, alt, dem, ox, oy, fit_scale):
    def resid(s, dx, dy):
        enu = (R @ (s * P.T)).T
        px = E0 + enu[:, 0] + dx - ox
        py = oy - (N0 + enu[:, 1] + dy)
        zd = sample(dem, px, py)
        g = np.isfinite(zd)
        if g.sum() < 50:
            return None, None
        r = (alt + enu[g, 2]) - zd[g]
        w = W[g].copy()
        # robust: two rounds of trimming the worst 25% |residual| - protects
        # against hydro-flattened water (DEM = 2020 water surface, not bottom)
        # and vegetation in the terrain band
        for _ in range(2):
            dz = np.average(r, weights=w)
            rr = np.abs(r - dz)
            cut = np.percentile(rr[w > 0], 75)
            w[rr > cut] = 0.0
        dz = np.average(r, weights=w)
        rr = r - dz
        return np.sqrt(np.average(rr ** 2, weights=w)), dz
    rms0, dz0 = resid(1.0, 0, 0)
    best = (rms0, 1.0, 0.0, 0.0, dz0)
    if fit_scale and rms0 is not None:
        for s in np.arange(0.85, 1.40001, 0.05):
            for dx in range(-12, 13, 4):
                for dy in range(-12, 13, 4):
                    r, dz = resid(s, dx, dy)
                    if r is not None and r < best[0]:
                        best = (r, s, dx, dy, dz)
        s0, dx0, dy0 = best[1:4]
        for s in np.arange(s0 - 0.05, s0 + 0.0501, 0.01):
            for dx in np.arange(dx0 - 3, dx0 + 3.1, 1.0):
                for dy in np.arange(dy0 - 3, dy0 + 3.1, 1.0):
                    r, dz = resid(s, dx, dy)
                    if r is not None and r < best[0]:
                        best = (r, s, dx, dy, dz)
    return rms0, dz0, best


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--ground-planes', required=True, nargs='+')
    ap.add_argument('--linked-root', required=True)
    ap.add_argument('--dem', required=True)
    ap.add_argument('--dem-origin', required=True, nargs=2, type=float,
                    help='UTM E,N of DEM pixel (0,0) corner (GeoTIFF tiepoint)')
    ap.add_argument('--out', default=None)
    ap.add_argument('--seg', default=None)
    ap.add_argument('--fit-scale', action='store_true')
    a = ap.parse_args()
    dem = np.array(Image.open(a.dem), dtype=np.float32)
    ox, oy = a.dem_origin
    rows_out = []
    hdr = ('timestamp ntiles dem_relief rms_nofit dz_nofit '
           'scale dx dy dz rms_fit').split()
    print('\t'.join(hdr))
    for gp in a.ground_planes:
        for r in csv.DictReader(open(gp), delimiter='\t'):
            ts = r['timestamp']
            if a.seg and a.seg not in ts:
                continue
            szxy = os.path.join(a.linked_root, ts, 'v88', f'{ts}_SZXY.tiff')
            if not os.path.exists(szxy):
                continue
            P, W = read_szxy(szxy)
            lat, lon, alt = (float(r['Latitude']), float(r['Longitude']),
                             float(r['Altitude']))
            q = np.array([float(r[f'qenu[{i}]']) for i in range(4)])
            R = qR(*q) @ M_ADAPT
            E0, N0 = utm17(lat, lon)
            # DEM relief within footprint (nominal placement)
            enu = (R @ P.T).T
            zd = sample(dem, E0 + enu[:, 0] - ox, oy - (N0 + enu[:, 1]))
            g = np.isfinite(zd)
            relief = zd[g].std() if g.sum() > 50 else float('nan')
            rms0, dz0, best = fit_segment(P, W, R, E0, N0, alt, dem, ox, oy,
                                          a.fit_scale)
            line = [ts, str(len(P)), f'{relief:.2f}',
                    f'{rms0:.2f}' if rms0 else 'nan',
                    f'{dz0:+.2f}' if dz0 else 'nan',
                    f'{best[1]:.2f}', f'{best[2]:+.0f}', f'{best[3]:+.0f}',
                    f'{best[4]:+.2f}' if best[4] else 'nan',
                    f'{best[0]:.2f}' if best[0] else 'nan']
            print('\t'.join(line))
            rows_out.append(line)
    if a.out:
        with open(a.out, 'w') as f:
            f.write('\t'.join(hdr) + '\n')
            for l in rows_out:
                f.write('\t'.join(l) + '\n')


if __name__ == '__main__':
    main()
