#!/usr/bin/env python3
"""Live 16x16 thermal heatmap viewer.

Reads the firmware's RTT/defmt log stream and updates a heatmap window in
real time, one frame at a time. The firmware (see `src/main.rs`) prints one
line per frame that looks like:

    12.345 [INFO ] frame ambient_c=23.5 min_c=10.2 max_c=45.6 pixels=[23.1, ...] (tts_test ...)

Everything else on the line (timestamp, log level, source location) and any
other log lines (warnings, startup messages, ...) are ignored.

Usage:
    # live, straight from the device (uses the runner in .cargo/config.toml,
    # i.e. `probe-rs run --chip STM32U585CI`):
    cargo run --release | python3 live_heatmap.py

    # replay a previously captured log file instead of a live device:
    python3 live_heatmap.py --replay output.txt
"""

import argparse
import re
import select
import sys

import matplotlib.pyplot as plt
import numpy as np

GRID = 16

# Fixed color-scale range (\u00b0C). Kept constant across frames (rather than
# autoscaling to each frame's min/max) so colors stay comparable over time;
# override with --vmin/--vmax if your scene runs outside this range.
DEFAULT_VMIN_C = 0.0
DEFAULT_VMAX_C = 100.0

_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")
_FRAME_RE = re.compile(
    r"ambient_c=(?P<ambient>-?\d+(?:\.\d+)?).*?"
    r"min_c=(?P<min>-?\d+(?:\.\d+)?).*?"
    r"max_c=(?P<max>-?\d+(?:\.\d+)?).*?"
    r"pixels=\[(?P<pixels>[^\]]*)\]"
)


def parse_frame(line):
    """Return (ambient, min, max, 16x16 array), or None if `line` isn't a frame line."""
    line = _ANSI_RE.sub("", line)
    match = _FRAME_RE.search(line)
    if not match:
        return None

    pixels = [float(v) for v in match.group("pixels").split(",") if v.strip()]
    if len(pixels) != GRID * GRID:
        return None

    grid = np.array(pixels, dtype=float).reshape(GRID, GRID)
    ambient = float(match.group("ambient"))
    vmin = float(match.group("min"))
    vmax = float(match.group("max"))
    return ambient, vmin, vmax, grid


def read_latest_frame(stream):
    """Block until at least one frame line is available on `stream`, then
    drain any further lines that are *already* buffered and return only the
    most recently parsed frame.

    This is what keeps the viewer live: if frames arrive faster than they
    can be drawn, we skip straight to the newest one each time instead of
    rendering every buffered frame in order and steadily falling behind.

    Returns `None` once `stream` reaches EOF.
    """
    frame = None
    while True:
        line = stream.readline()
        if line == "":
            return frame  # EOF
        parsed = parse_frame(line)
        if parsed is not None:
            frame = parsed
        # non-blocking peek: is there already more data waiting?
        more_data_ready, _, _ = select.select([stream], [], [], 0)
        if not more_data_ready and frame is not None:
            return frame
        # otherwise keep draining (or, if nothing parsed yet, block again
        # on the next readline() until a new line arrives)


class LiveHeatmap:
    """A matplotlib window that redraws itself each time `update` is called.

    The color scale (`vmin`/`vmax`) is fixed for the lifetime of the window,
    so colors stay comparable from frame to frame; only the pixel data and
    the title (which still reports each frame's actual min/max) change.
    """

    def __init__(self, cmap, vmin, vmax, annotate=True):
        plt.ion()
        self.fig, ax = plt.subplots(figsize=(8, 7))
        self.image = ax.imshow(
            np.zeros((GRID, GRID)), cmap=cmap, interpolation="nearest", vmin=vmin, vmax=vmax
        )
        self.colorbar = self.fig.colorbar(self.image, ax=ax, label="\u00b0C")
        self.title = ax.set_title("waiting for first frame...")
        ax.set_xticks(range(GRID))
        ax.set_yticks(range(GRID))

        # midpoint of the fixed color scale, used to pick readable text color
        self.mid = (vmin + vmax) / 2
        self.annotations = None
        if annotate:
            self.annotations = [
                [
                    ax.text(col, row, "", ha="center", va="center", fontsize=6, color="white")
                    for col in range(GRID)
                ]
                for row in range(GRID)
            ]

        self.fig.tight_layout()
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
        self.frame_count = 0

    def update(self, ambient, frame_min, frame_max, grid):
        self.frame_count += 1
        self.image.set_data(grid)
        self.title.set_text(
            f"frame {self.frame_count}  ambient {ambient:.1f}\u00b0C  "
            f"range {frame_min:.1f}..{frame_max:.1f}\u00b0C"
        )
        if self.annotations is not None:
            for row in range(GRID):
                for col in range(GRID):
                    value = grid[row, col]
                    text = self.annotations[row][col]
                    text.set_text(f"{value:.0f}")
                    text.set_color("white" if value < self.mid else "black")
        self.fig.canvas.draw_idle()
        self.fig.canvas.flush_events()

    def is_open(self):
        return plt.fignum_exists(self.fig.number)


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--replay", help="replay a captured log file instead of reading stdin")
    parser.add_argument("--replay-fps", type=float, default=5.0, help="frame rate when replaying a file")
    parser.add_argument("--cmap", default="inferno", help="matplotlib colormap")
    parser.add_argument("--vmin", type=float, default=DEFAULT_VMIN_C, help="fixed color-scale minimum (\u00b0C)")
    parser.add_argument("--vmax", type=float, default=DEFAULT_VMAX_C, help="fixed color-scale maximum (\u00b0C)")
    parser.add_argument("--no-annot", action="store_true", help="do not print the value inside each cell")
    args = parser.parse_args()

    heatmap = LiveHeatmap(args.cmap, args.vmin, args.vmax, annotate=not args.no_annot)
    try:
        if args.replay:
            # deliberately paced, frame by frame, at --replay-fps
            with open(args.replay) as source:
                for line in source:
                    if not heatmap.is_open():
                        break
                    frame = parse_frame(line)
                    if frame is None:
                        continue
                    heatmap.update(*frame)
                    plt.pause(1.0 / args.replay_fps)
        else:
            # live: always show the newest frame, never a backlog
            while heatmap.is_open():
                frame = read_latest_frame(sys.stdin)
                if frame is None:
                    break  # EOF: the device stream ended
                heatmap.update(*frame)
                plt.pause(0.001)  # let the GUI event loop breathe
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()