# /// script
# requires-python = ">=3.10"
# dependencies = ["pyserial", "numpy", "matplotlib"]
# ///
"""Download a recording from CrashRecorder, convert it to CSV and plot it.

    uv run tools/crash_dump.py                      download, save and plot
    uv run tools/crash_dump.py --png                save the plot as PNG instead of opening a window
    uv run tools/crash_dump.py --from-bin FILE.bin  re-parse and plot a saved recording

Keep the record format in sync with RecordTag in CrashRecorder/CrashRecorder.ino.
"""

import argparse
import csv
import struct
import sys
import time
from datetime import datetime
from pathlib import Path

import numpy as np

TAG_HEADER = 0x01
TAG_IMU = 0x02
TAG_MAG = 0x03
TAG_INPUTS = 0x04
TAG_MARKER = 0x05
TAG_OVERFLOW = 0x06
TAG_END = 0xFF

PAYLOAD_LENGTHS = {
    TAG_HEADER: 12,
    TAG_IMU: 12,
    TAG_MAG: 6,
    TAG_INPUTS: 1,
    TAG_MARKER: 1,
    TAG_OVERFLOW: 0,
}

INPUT_BITS = {"pedal": 0, "forward": 1, "reverse": 2, "mode": 3, "left": 4, "right": 5}

MAG_UT_PER_LSB = 0.15  # 16-bit mode, without the per-chip sensitivity adjustment

# from the tilt test on 2026-10-21: subtracting these brings |a| to 1g in every pose within ~2mg,
# cross-checked against gyro-integrated tilt angles
ACCEL_OFFSET = np.array([0.490, -0.277, -0.004])  # g, sensor x y z

# sensor mounted on its side: +x up, +y left, +z backward
# vehicle frame accel: forward, left, up (up reads +1g at rest)
# vehicle frame rotation, aviation style: roll + right side down, pitch + nose up, yaw + turning right
def to_vehicle(accel: np.ndarray, gyro: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    corrected = accel - ACCEL_OFFSET
    vehicle_accel = np.column_stack([-corrected[:, 2], corrected[:, 1], corrected[:, 0]])
    vehicle_rotation = np.column_stack([-gyro[:, 2], -gyro[:, 1], -gyro[:, 0]])

    return vehicle_accel, vehicle_rotation

DEFAULT_OUT = Path(__file__).resolve().parent.parent / "recordings"


def download(port: str) -> bytes:
    import serial

    with serial.Serial(port, timeout=1) as connection:
        time.sleep(0.2)
        connection.reset_input_buffer()

        connection.write(b"i")
        while True:
            line = connection.readline().decode(errors="replace").strip()
            if not line:
                sys.exit("no response from the board, is CrashRecorder running?")
            if line == "INFO END":
                break
            print(f"  {line}")

        connection.write(b"d")
        while True:
            line = connection.readline().decode(errors="replace").strip()
            if not line:
                sys.exit("no response to the dump command")
            if line == "DUMP UNAVAILABLE":
                sys.exit("board is busy, stop the recording first")
            if line.startswith("DUMP "):
                length = int(line.split()[1])
                break

        print(f"downloading {length} bytes")
        connection.timeout = 10
        data = connection.read(length)

        if len(data) != length:
            sys.exit(f"download incomplete, got {len(data)} of {length} bytes")

        return data


def parse(data: bytes) -> dict:
    header = None
    imu = []  # raw int16 rows: ax ay az gx gy gz
    mag_at = []  # (imu index, mx, my, mz)
    inputs_at = []  # (imu index, raw levels)
    markers = []  # (imu index, marker)
    overflows = []  # imu index

    position = 0
    while position < len(data):
        tag = data[position]
        if tag == TAG_END:
            break

        length = PAYLOAD_LENGTHS.get(tag)
        if length is None:
            print(f"warning: unknown tag 0x{tag:02X} at byte {position}, stopping", file=sys.stderr)
            break

        payload = data[position + 1 : position + 1 + length]
        if len(payload) < length:
            break  # cut off by power loss or storage full
        position += 1 + length

        if tag == TAG_HEADER:
            if payload[:3] != b"QLR":
                sys.exit("not a CrashRecorder recording")
            version, who_am_i, mag_ok, rate, accel_lsb, gyro_lsb_x10 = struct.unpack("<BBBHHH", payload[3:])
            header = {
                "version": version,
                "who_am_i": who_am_i,
                "mag_ok": bool(mag_ok),
                "rate": rate,
                "accel_lsb_per_g": accel_lsb,
                "gyro_lsb_per_dps": gyro_lsb_x10 / 10,
            }
        elif tag == TAG_IMU:
            imu.append(struct.unpack(">6h", payload))
        elif tag == TAG_MAG:
            mag_at.append((len(imu), *struct.unpack("<3h", payload)))
        elif tag == TAG_INPUTS:
            inputs_at.append((len(imu), payload[0]))
        elif tag == TAG_MARKER:
            markers.append((len(imu), payload[0]))
        elif tag == TAG_OVERFLOW:
            overflows.append(len(imu))

    if header is None:
        sys.exit("recording has no header")

    raw = np.array(imu, dtype=np.float64).reshape(-1, 6)
    count = len(raw)

    # mag and inputs are stored when they change, hold them until the next record
    mag = np.full((count, 3), np.nan)
    for (index, *values), following in zip(mag_at, mag_at[1:] + [(count,)]):
        mag[index : following[0]] = np.array(values) * MAG_UT_PER_LSB

    inputs = np.zeros((count, len(INPUT_BITS)), dtype=np.int8)
    for (index, levels), following in zip(inputs_at, inputs_at[1:] + [(count,)]):
        for column, bit in enumerate(INPUT_BITS.values()):
            inputs[index : following[0], column] = 0 if levels >> bit & 1 else 1  # low is active

    accel = raw[:, :3] / header["accel_lsb_per_g"]
    gyro = raw[:, 3:] / header["gyro_lsb_per_dps"]
    vehicle_accel, vehicle_rotation = to_vehicle(accel, gyro)

    return {
        "header": header,
        "t": np.arange(count) / header["rate"],
        "accel": accel,
        "accel_magnitude": np.linalg.norm(accel, axis=1),
        "gyro": gyro,
        "vehicle_accel": vehicle_accel,
        "vehicle_magnitude": np.linalg.norm(vehicle_accel, axis=1),
        "horizontal": np.linalg.norm(vehicle_accel[:, :2], axis=1),
        "vehicle_rotation": vehicle_rotation,
        "mag": mag,
        "inputs": inputs,
        "markers": markers,
        "overflows": overflows,
    }


def write_csv(recording: dict, path: Path) -> None:
    marker_column = np.zeros(len(recording["t"]), dtype=np.int8)
    for index, marker in recording["markers"]:
        marker_column[min(index, len(marker_column) - 1)] = marker

    overflow_column = np.zeros(len(recording["t"]), dtype=np.int8)
    for index in recording["overflows"]:
        overflow_column[min(index, len(overflow_column) - 1)] = 1

    with path.open("w", newline="") as file:
        writer = csv.writer(file)
        writer.writerow(
            [
                "t",
                "a_forward", "a_left", "a_up", "a_vehicle_magnitude", "a_horizontal",
                "roll_rate", "pitch_rate", "yaw_rate",
                "ax", "ay", "az", "a_magnitude", "gx", "gy", "gz",
                "mx", "my", "mz",
                *INPUT_BITS, "marker", "overflow",
            ]
        )
        for i, t in enumerate(recording["t"]):
            writer.writerow(
                [
                    f"{t:.3f}",
                    *(f"{v:.4f}" for v in recording["vehicle_accel"][i]),
                    f"{recording['vehicle_magnitude'][i]:.4f}",
                    f"{recording['horizontal'][i]:.4f}",
                    *(f"{v:.2f}" for v in recording["vehicle_rotation"][i]),
                    *(f"{v:.4f}" for v in recording["accel"][i]),
                    f"{recording['accel_magnitude'][i]:.4f}",
                    *(f"{v:.2f}" for v in recording["gyro"][i]),
                    *("" if np.isnan(v) else f"{v:.2f}" for v in recording["mag"][i]),
                    *recording["inputs"][i],
                    marker_column[i],
                    overflow_column[i],
                ]
            )


def plot(recording: dict, title: str, png: Path | None) -> None:
    import matplotlib

    if png:
        matplotlib.use("Agg")

    import matplotlib.pyplot as plt

    t = recording["t"]
    figure, axes = plt.subplots(4, 1, sharex=True, figsize=(16, 11), height_ratios=[3, 2, 1, 1])
    figure.suptitle(title)

    accel_axis, gyro_axis, mag_axis, inputs_axis = axes

    for column, name in enumerate(["forward", "left", "up"]):
        accel_axis.plot(t, recording["vehicle_accel"][:, column], linewidth=0.7, label=name)
    accel_axis.plot(t, recording["horizontal"], linewidth=0.7, color="black", alpha=0.6, label="horizontal")
    accel_axis.set_ylabel("acceleration (g)")

    for column, name in enumerate(["roll (+ right down)", "pitch (+ nose up)", "yaw (+ right)"]):
        gyro_axis.plot(t, recording["vehicle_rotation"][:, column], linewidth=0.7, label=name)
    gyro_axis.set_ylabel("rotation (°/s)")

    for column, name in enumerate("xyz"):
        mag_axis.plot(t, recording["mag"][:, column], linewidth=0.7, label=f"m{name}")
    mag_axis.set_ylabel("magnetic (µT)")

    # one row per input, filled while active
    for column, name in enumerate(INPUT_BITS):
        inputs_axis.fill_between(
            t, column + 0.1, column + 0.1 + recording["inputs"][:, column] * 0.8, step="post", linewidth=0, label=name
        )
    inputs_axis.set_yticks(np.arange(len(INPUT_BITS)) + 0.5, list(INPUT_BITS))
    inputs_axis.set_ylim(0, len(INPUT_BITS))
    inputs_axis.set_xlabel("time (s)")

    for axis in axes:
        for index, marker in recording["markers"]:
            axis.axvline(index / recording["header"]["rate"], color="red" if marker == 1 else "orange", linestyle="--", linewidth=1)
        for index in recording["overflows"]:
            axis.axvline(index / recording["header"]["rate"], color="grey", linestyle=":", linewidth=1)
        axis.grid(alpha=0.3)
        axis.legend(loc="upper right", fontsize="small", ncol=7)

    figure.tight_layout()

    if png:
        figure.savefig(png, dpi=120)
        print(f"plot saved to {png}")
    else:
        plt.show()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--port", default="/dev/ttyACM0")
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT, help="directory for downloaded recordings")
    parser.add_argument("--from-bin", type=Path, help="parse a saved recording instead of downloading")
    parser.add_argument("--png", action="store_true", help="save the plot as PNG instead of opening a window")
    parser.add_argument("--no-plot", action="store_true")
    arguments = parser.parse_args()

    if arguments.from_bin:
        bin_path = arguments.from_bin
        data = bin_path.read_bytes()
    else:
        data = download(arguments.port)
        arguments.out.mkdir(parents=True, exist_ok=True)
        bin_path = arguments.out / f"crash-{datetime.now():%Y%m%d-%H%M%S}.bin"
        bin_path.write_bytes(data)
        print(f"recording saved to {bin_path}")

    recording = parse(data)
    header = recording["header"]
    duration = len(recording["t"]) / header["rate"]

    print(
        f"{len(recording['t'])} samples ({duration:.1f} s at {header['rate']} Hz), "
        f"WHO_AM_I 0x{header['who_am_i']:02X}, {len(recording['markers'])} markers, "
        f"{len(recording['overflows'])} fifo overflows"
    )

    csv_path = bin_path.with_suffix(".csv")
    write_csv(recording, csv_path)
    print(f"csv saved to {csv_path}")

    if not arguments.no_plot:
        plot(recording, bin_path.stem, bin_path.with_suffix(".png") if arguments.png else None)


if __name__ == "__main__":
    main()