#!/usr/bin/env python3
"""
session_map.py - geometry for plotting FireMapper sessions on a map.

Provides trigger points (GPS position of each frame) and oblique image
footprints projected onto flat ground.

Camera model (confirmed with the operator):
  * Cameras sit on a platform that rolls about the aircraft longitudinal (forward)
    axis by platform_angle_deg, sweeping every camera cross-track.
  * Mounting on the platform: thermal looks straight down (nadir); the RGB cameras
    are tilted toward the forward direction by their cam<NN> off-nadir angle
    (cam25 -> 25 deg from nadir, cam45 -> 45 deg).
  * The full chain is: aircraft attitude (yaw/pitch/roll, body->NED) @ platform roll
    @ camera mounting. The image corners are then ray-cast onto flat ground at a
    given elevation (footprints); the optical axis gives the true EXIF orientation.

If footprints / headings come out mirrored, flip SCAN_SIGN below.
"""

from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np

M_PER_DEG_LAT = 111320.0

SCAN_SIGN = +1.0    # sign of the platform_angle rotation about the aircraft roll (forward) axis

# Camera mounting on the rotating platform (operator spec):
#   thermal (FLIR A65): optical axis straight DOWN (nadir); mounted VERTICAL/portrait,
#                       so sensor width runs ALONG the flight track.
#   RGB cameras:        optical axis tilted toward the FORWARD direction by the cam<NN>
#                       off-nadir angle (cam25 = 25 deg from straight-down); mounted
#                       LANDSCAPE, so sensor width runs ACROSS the flight track.
# The platform then rolls about the forward axis (platform_angle_deg), sweeping every
# camera cross-track, and finally the aircraft attitude (yaw/pitch/roll) is applied.

# FLIR A65, 640x512, 25 deg lens -> ~25 x 20 deg field of view
THERMAL_FOV_DEFAULT = (25.0, 20.0)


# --------------------------------------------------------------------------- #
# rotations
# --------------------------------------------------------------------------- #
def _rot_x(a: float) -> np.ndarray:
    c, s = math.cos(a), math.sin(a)
    return np.array([[1, 0, 0], [0, c, -s], [0, s, c]])


def _body_to_ned(yaw: float, pitch: float, roll: float) -> np.ndarray:
    """Aerospace 3-2-1 (yaw, pitch, roll) body->NED rotation; angles in radians."""
    cy, sy = math.cos(yaw), math.sin(yaw)
    cp, sp = math.cos(pitch), math.sin(pitch)
    cr, sr = math.cos(roll), math.sin(roll)
    return np.array([
        [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
        [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
        [-sp,     cp * sr,                cp * cr],
    ])


# --------------------------------------------------------------------------- #
# sessions & frames
# --------------------------------------------------------------------------- #
def is_session_dir(p: Path) -> bool:
    """A capture session = a folder holding thermal/ or cam*/ or a manifest.json.
    Name-independent (works whether it's 'session_...' or anything else), and
    excludes our own derived outputs."""
    if not p.is_dir() or p.name.endswith(("_exif", "_stretched")) or p.name == "__pycache__":
        return False
    return ((p / "manifest.json").exists() or (p / "thermal").is_dir()
            or any(p.glob("cam*")))


def list_sessions(root) -> list[Path]:
    """Session folders directly under root, or one level deeper (e.g. grouped in a
    parent folder like 'Streifen/'). Sorted by name."""
    root = Path(root)
    found: list[Path] = []
    for p in sorted(root.iterdir()):
        if not p.is_dir():
            continue
        if is_session_dir(p):
            found.append(p)
        elif not p.name.endswith(("_exif", "_stretched")) and p.name != "__pycache__":
            found += [q for q in sorted(p.iterdir()) if is_session_dir(q)]
    return found


def off_nadir_from_name(folder: str) -> float:
    digits = "".join(ch for ch in folder if ch.isdigit())
    return float(digits) if digits else 0.0


def _read_frame(json_path: Path, kind: str, off_nadir: float) -> dict | None:
    try:
        j = json.loads(json_path.read_text(encoding="utf-8"))
    except Exception:  # noqa: BLE001
        return None
    pos = (j.get("gps") if kind == "thermal" else (j.get("position") or j.get("gps"))) or {}
    if pos.get("lat") is None or pos.get("lon") is None:
        return None
    imu = j.get("imu") or {}
    plat = j.get("platform") or {}
    img = j.get("image") or {}
    lens = j.get("lens") or {}
    return {
        "kind": kind, "off_nadir": off_nadir, "session": json_path.parent.parent.name,
        "lat": pos["lat"], "lon": pos["lon"], "alt": pos.get("alt"),
        "yaw": imu.get("yaw"), "pitch": imu.get("pitch"), "roll": imu.get("roll"),
        "platform_angle": plat.get("platform_angle_deg"),
        "width": img.get("width"), "height": img.get("height"),
        "focal_mm": lens.get("focal_length_mm"), "pixel_um": lens.get("pixel_size_um"),
    }


def iter_session_frames(session, include_cam=True, include_thermal=True,
                        thermal_off_nadir=0.0) -> list[dict]:
    """All plottable frames in a session (those with GPS), tagged by kind."""
    session = Path(session)
    frames: list[dict] = []
    for sub in sorted(p for p in session.iterdir() if p.is_dir()):
        if sub.name.startswith("cam") and include_cam:
            off = off_nadir_from_name(sub.name)
            for jf in sorted(sub.glob("*.json")):
                fr = _read_frame(jf, "cam", off)
                if fr:
                    frames.append(fr)
        elif sub.name == "thermal" and include_thermal:
            for jf in sorted(sub.glob("*.json")):
                fr = _read_frame(jf, "thermal", thermal_off_nadir)
                if fr:
                    frames.append(fr)
    return frames


# --------------------------------------------------------------------------- #
# footprint projection
# --------------------------------------------------------------------------- #
def lens_fov(focal_mm, pixel_um, width, height):
    """(hfov, vfov) in radians from focal length / pixel pitch / pixel counts."""
    if not (focal_mm and pixel_um and width and height):
        return None
    sw = pixel_um / 1000.0 * width   # mm
    sh = pixel_um / 1000.0 * height
    return (2 * math.atan((sw / 2) / focal_mm), 2 * math.atan((sh / 2) / focal_mm))


def camera_axes_ned(frame: dict, sensor: bool = True):
    """
    Camera frame as NED unit vectors (image_right, image_down, optical_axis) for the
    TRUE camera pointing, built as:
        aircraft attitude (body->NED)  @  platform roll about forward axis  @  mounting

    The optical axis is identical either way. With sensor=True the right/down axes
    follow the real sensor mounting (RGB width ALONG track = portrait) - used for
    footprints. With sensor=False a canonical landscape frame is used (right =
    cross-track) so the orientation 'roll' is the camera bank, not the 90 deg sensor
    rotation.
    """
    R = _body_to_ned(math.radians(frame["yaw"]), math.radians(frame["pitch"]),
                     math.radians(frame["roll"])) \
        @ _rot_x(math.radians(SCAN_SIGN * (frame["platform_angle"] or 0.0)))
    if frame["kind"] == "thermal":
        # nadir, VERTICAL/portrait mount -> sensor width runs ALONG track
        optical = np.array([0.0, 0.0, 1.0])
        if sensor:
            right = np.array([-1.0, 0.0, 0.0])    # aft (along-track) = width
            down_im = np.array([0.0, -1.0, 0.0])  # port (cross-track) = height
        else:                                     # canonical landscape (for bank/roll only)
            right = np.array([0.0, 1.0, 0.0])
            down_im = np.array([-1.0, 0.0, 0.0])
    else:
        # RGB tilted forward by off-nadir; LANDSCAPE mount -> width ACROSS track.
        # (landscape == canonical here, so sensor and orientation frames coincide)
        a = math.radians(frame["off_nadir"])
        optical = np.array([math.sin(a), 0.0, math.cos(a)])    # forward & down
        right = np.array([0.0, 1.0, 0.0])                      # cross-track = width
        down_im = np.array([-math.cos(a), 0.0, math.sin(a)])   # along-track = height
    return R @ right, R @ down_im, R @ optical


def camera_orientation(frame: dict):
    """
    (heading_deg, pitch_deg, roll_deg) of the true camera optical axis, or None.
      heading: compass azimuth of the optical axis (0-360, from true north)
      pitch  : elevation of the optical axis (0 = horizon, -90 = straight down)
      roll   : bank about the optical axis (0 = image bottom points to ground-down)
    """
    if any(frame.get(k) is None for k in ("yaw", "pitch", "roll")):
        return None
    right, down_im, optical = camera_axes_ned(frame, sensor=False)   # canonical for bank
    if math.hypot(optical[0], optical[1]) < 1e-6:      # ~nadir: heading from image-up
        ref = -down_im
        heading = math.degrees(math.atan2(ref[1], ref[0])) % 360.0
    else:
        heading = math.degrees(math.atan2(optical[1], optical[0])) % 360.0
    pitch = math.degrees(math.asin(max(-1.0, min(1.0, -float(optical[2])))))
    wd = np.array([0.0, 0.0, 1.0])                     # world down
    roll = math.degrees(math.atan2(float(wd @ right), float(wd @ down_im)))
    return heading, pitch, roll


def metashape_ypr(frame: dict):
    """
    (yaw, pitch, roll) of the camera in the Pix4D / Agisoft Metashape convention,
    or None. This is the attitude photogrammetry software expects in the XMP
    ``Camera:Yaw/Pitch/Roll`` tags, and it differs from ``camera_orientation``:

      * pitch = 0 means the optical axis looks straight DOWN (nadir); pitch grows
        toward +90 as the camera tilts to look FORWARD. So our nadir thermal -> 0,
        cam25 -> ~25, cam45 -> ~45 (the cam<NN> off-nadir angle), as Metashape wants.
      * yaw    = azimuth of the image-TOP direction (0 = top points north), 0-360.
      * roll   = rotation about the optical axis.

    These are the angles psi/theta/phi of the body->NED rotation
    C^n_b = Rz(yaw) Ry(pitch) Rx(roll)  (NED nav frame; right-hand rule), recovered
    from the true camera axes so aircraft attitude and the platform sweep are all
    folded in. Pix4D's default "image top = flight direction, camera looks down"
    body frame maps to our axes as:  x_b (front) = image-up = -down_im,
    y_b (right) = image-right,  z_b (bottom) = optical axis.
    Ref: Pix4D "Yaw, Pitch, Roll and Omega, Phi, Kappa angles".
    """
    if any(frame.get(k) is None for k in ("yaw", "pitch", "roll")):
        return None
    right, down_im, optical = camera_axes_ned(frame, sensor=True)  # true stored-image axes
    xb = -down_im                       # body front  (image top), col 0 of C^n_b
    # C^n_b = [xb | right | optical]; extract from Rz(psi)Ry(theta)Rx(phi):
    #   xb[2] = -sin(theta);  yaw = atan2(xb[1], xb[0]);  roll = atan2(right[2], optical[2])
    pitch = math.degrees(math.asin(max(-1.0, min(1.0, -float(xb[2])))))
    yaw = math.degrees(math.atan2(float(xb[1]), float(xb[0]))) % 360.0
    roll = math.degrees(math.atan2(float(right[2]), float(optical[2])))
    return yaw, pitch, roll


def footprint(frame: dict, ground_elev: float,
              thermal_fov_deg: tuple[float, float] | None = None) -> list[tuple[float, float]] | None:
    """Four (lat, lon) ground corners of the image, or None if not projectable."""
    if any(frame[k] is None for k in ("alt", "yaw", "pitch", "roll")):
        return None
    H = frame["alt"] - ground_elev
    if H <= 1.0:
        return None

    if frame["kind"] == "thermal":
        if not thermal_fov_deg:
            return None
        hfov, vfov = math.radians(thermal_fov_deg[0]), math.radians(thermal_fov_deg[1])
    else:
        fov = lens_fov(frame["focal_mm"], frame["pixel_um"], frame["width"], frame["height"])
        if not fov:
            return None
        hfov, vfov = fov

    right, down_im, optical = camera_axes_ned(frame)   # width axis, height axis, boresight
    th, tv = math.tan(hfov / 2), math.tan(vfov / 2)
    coslat = math.cos(math.radians(frame["lat"]))
    corners = []
    for sx, sy in ((-1, -1), (1, -1), (1, 1), (-1, 1)):
        d = optical + sx * th * right + sy * tv * down_im   # ray toward an image corner
        if d[2] <= 1e-3:          # ray not pointing at the ground
            return None
        t = H / d[2]
        dlat = (t * d[0]) / M_PER_DEG_LAT
        dlon = (t * d[1]) / (M_PER_DEG_LAT * coslat)
        corners.append((frame["lat"] + dlat, frame["lon"] + dlon))
    return corners


def bounds(frames: list[dict]):
    """(min_lat, min_lon, max_lat, max_lon) over frame trigger points."""
    lats = [f["lat"] for f in frames]
    lons = [f["lon"] for f in frames]
    return min(lats), min(lons), max(lats), max(lons)