deepivus/gui/annotation_editor.py

"""Interactive GUI for editing per-frame IVUS contour annotations."""

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from datetime import datetime
from glob import glob
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.lines import Line2D
from matplotlib.widgets import Button, Slider

from ..io.dicom import read_dicom


@dataclass
class FrameState:
    """In-memory annotation state for one frame."""

    frame: int
    lumen_x: list[float]
    lumen_y: list[float]
    plaque_x: list[float]
    plaque_y: list[float]
    lumen_confidence: float | None
    plaque_confidence: float | None
    bifurcation: bool = False


class AnnotationEditor:
    """Matplotlib-based editor for contour annotations."""

    def __init__(
        self,
        dicom_path: str,
        annotations_path: str,
        edits_path: str,
    ) -> None:
        self.dicom_path = dicom_path
        self.annotations_path = annotations_path
        self.edits_path = edits_path

        _, images = read_dicom(dicom_path)
        self.images = images

        self.frame_states = self._load_base_annotations(annotations_path)
        if len(self.frame_states) != images.shape[0]:
            raise ValueError(
                "Frame count mismatch between DICOM and annotations: "
                f"{images.shape[0]} vs {len(self.frame_states)}"
            )

        self._load_existing_edits(edits_path)
        self._ensure_edits_file_header()

        self.current_frame = 0
        self.drag_target: tuple[str, int] | None = None
        self.drag_threshold_px = 10.0
        self.frame_dirty = False

        self.fig = None
        self.ax_image = None
        self.image_artist = None
        self.lumen_line: Line2D | None = None
        self.plaque_line: Line2D | None = None
        self.lumen_points = None
        self.plaque_points = None

        self.slider: Slider | None = None
        self.status_text = None
        self.is_updating_slider = False

    @staticmethod
    def _safe_float(value: Any) -> float | None:
        if value is None:
            return None
        try:
            out = float(value)
        except (TypeError, ValueError):
            return None
        if np.isnan(out):
            return None
        return out

    def _load_base_annotations(self, path: str) -> list[FrameState]:
        states: list[FrameState] = []
        with open(path, "r", encoding="utf-8") as fp:
            for raw in fp:
                line = raw.strip()
                if not line:
                    continue
                rec = json.loads(line)
                if rec.get("record_type") != "frame":
                    continue
                states.append(
                    FrameState(
                        frame=int(rec["frame"]),
                        lumen_x=[float(v) for v in rec.get("lumen", {}).get("x", [])],
                        lumen_y=[float(v) for v in rec.get("lumen", {}).get("y", [])],
                        plaque_x=[float(v) for v in rec.get("plaque", {}).get("x", [])],
                        plaque_y=[float(v) for v in rec.get("plaque", {}).get("y", [])],
                        lumen_confidence=self._safe_float(rec.get("lumen_confidence")),
                        plaque_confidence=self._safe_float(rec.get("plaque_confidence")),
                        bifurcation=bool(rec.get("bifurcation", False)),
                    )
                )

        states.sort(key=lambda s: s.frame)
        expected = list(range(len(states)))
        actual = [s.frame for s in states]
        if actual != expected:
            raise ValueError(
                "Base annotations must contain contiguous frame records in order "
                f"(got first frames: {actual[:10]})"
            )
        return states

    def _load_existing_edits(self, path: str) -> None:
        if not os.path.exists(path):
            return

        latest: dict[int, dict[str, Any]] = {}
        with open(path, "r", encoding="utf-8") as fp:
            for raw in fp:
                line = raw.strip()
                if not line:
                    continue
                rec = json.loads(line)
                if rec.get("record_type") != "frame_edit":
                    continue
                frame_idx = int(rec["frame"])
                latest[frame_idx] = rec

        for frame_idx, rec in latest.items():
            if frame_idx < 0 or frame_idx >= len(self.frame_states):
                continue
            state = self.frame_states[frame_idx]
            state.lumen_x = [float(v) for v in rec.get("lumen", {}).get("x", state.lumen_x)]
            state.lumen_y = [float(v) for v in rec.get("lumen", {}).get("y", state.lumen_y)]
            state.plaque_x = [float(v) for v in rec.get("plaque", {}).get("x", state.plaque_x)]
            state.plaque_y = [float(v) for v in rec.get("plaque", {}).get("y", state.plaque_y)]
            state.bifurcation = bool(rec.get("bifurcation", state.bifurcation))

    def _ensure_edits_file_header(self) -> None:
        os.makedirs(os.path.dirname(os.path.abspath(self.edits_path)), exist_ok=True)
        if os.path.exists(self.edits_path):
            return
        meta = {
            "record_type": "meta",
            "created_at": datetime.now().isoformat(timespec="seconds"),
            "dicom_path": os.path.abspath(self.dicom_path),
            "base_annotations_path": os.path.abspath(self.annotations_path),
            "format": "append_only_frame_edits",
        }
        with open(self.edits_path, "w", encoding="utf-8") as fp:
            fp.write(json.dumps(meta) + "\n")

    def _append_frame_edit(self, frame_idx: int, reason: str) -> None:
        state = self.frame_states[frame_idx]
        rec = {
            "record_type": "frame_edit",
            "saved_at": datetime.now().isoformat(timespec="seconds"),
            "reason": reason,
            "frame": frame_idx,
            "bifurcation": state.bifurcation,
            "lumen": {"x": state.lumen_x, "y": state.lumen_y},
            "plaque": {"x": state.plaque_x, "y": state.plaque_y},
            "lumen_confidence": state.lumen_confidence,
            "plaque_confidence": state.plaque_confidence,
        }

        with open(self.edits_path, "a", encoding="utf-8") as fp:
            fp.write(json.dumps(rec) + "\n")
            fp.flush()
            os.fsync(fp.fileno())

        self.frame_dirty = False

    def _build_ui(self) -> None:
        self.fig = plt.figure(figsize=(13, 9))
        self.fig.canvas.manager.set_window_title("DeepIVUS Annotation Editor")

        self.ax_image = self.fig.add_axes([0.05, 0.22, 0.9, 0.74])
        self.ax_image.set_title("DeepIVUS Annotation Editor", fontsize=14, weight="bold")
        self.ax_image.set_axis_off()

        slider_ax = self.fig.add_axes([0.12, 0.13, 0.76, 0.035])
        self.slider = Slider(
            ax=slider_ax,
            label="Frame",
            valmin=0,
            valmax=len(self.frame_states) - 1,
            valinit=0,
            valstep=1,
            color="#1f77b4",
        )

        prev_ax = self.fig.add_axes([0.12, 0.05, 0.1, 0.055])
        next_ax = self.fig.add_axes([0.24, 0.05, 0.1, 0.055])
        save_ax = self.fig.add_axes([0.42, 0.05, 0.14, 0.055])
        bif_ax = self.fig.add_axes([0.72, 0.05, 0.16, 0.055])

        self.prev_button = Button(prev_ax, "Prev Frame", color="#E0E0E0", hovercolor="#D0D0D0")
        self.next_button = Button(next_ax, "Next Frame", color="#E0E0E0", hovercolor="#D0D0D0")
        self.save_button = Button(save_ax, "Save Frame", color="#D6F5D6", hovercolor="#BFF0BF")
        self.bif_button = Button(bif_ax, "Bifurcation: No", color="#F5D6D6", hovercolor="#F0BFBF")

        self.status_text = self.fig.text(0.05, 0.18, "", fontsize=10)

        self.slider.on_changed(self._on_slider)
        self.prev_button.on_clicked(self._on_prev)
        self.next_button.on_clicked(self._on_next)
        self.save_button.on_clicked(self._on_save)
        self.bif_button.on_clicked(self._on_toggle_bifurcation)

        self.fig.canvas.mpl_connect("button_press_event", self._on_press)
        self.fig.canvas.mpl_connect("motion_notify_event", self._on_motion)
        self.fig.canvas.mpl_connect("button_release_event", self._on_release)
        self.fig.canvas.mpl_connect("key_press_event", self._on_key)
        self.fig.canvas.mpl_connect("close_event", self._on_close)

    def _state(self) -> FrameState:
        return self.frame_states[self.current_frame]

    def _render_frame(self) -> None:
        state = self._state()
        image = self.images[self.current_frame]

        if self.image_artist is None:
            self.image_artist = self.ax_image.imshow(image, cmap="gray")
            self.lumen_line = self.ax_image.plot([], [], color="#1db954", lw=2)[0]
            self.plaque_line = self.ax_image.plot([], [], color="#ff5a5a", lw=2)[0]
            self.lumen_points = self.ax_image.scatter([], [], c="#1db954", s=28, edgecolors="black", linewidths=0.4)
            self.plaque_points = self.ax_image.scatter([], [], c="#ff5a5a", s=28, edgecolors="black", linewidths=0.4)
        else:
            self.image_artist.set_data(image)

        lumen_x, lumen_y = state.lumen_x, state.lumen_y
        plaque_x, plaque_y = state.plaque_x, state.plaque_y

        self.lumen_line.set_data(lumen_x + lumen_x[:1], lumen_y + lumen_y[:1])
        self.plaque_line.set_data(plaque_x + plaque_x[:1], plaque_y + plaque_y[:1])

        lumen_offsets = np.c_[lumen_x, lumen_y] if lumen_x and lumen_y else np.empty((0, 2))
        plaque_offsets = np.c_[plaque_x, plaque_y] if plaque_x and plaque_y else np.empty((0, 2))
        self.lumen_points.set_offsets(lumen_offsets)
        self.plaque_points.set_offsets(plaque_offsets)

        bif_text = "Yes" if state.bifurcation else "No"
        bif_color = "#D6F5D6" if state.bifurcation else "#F5D6D6"
        self.bif_button.label.set_text(f"Bifurcation: {bif_text}")
        self.bif_button.ax.set_facecolor(bif_color)

        self.status_text.set_text(
            f"Frame {self.current_frame + 1}/{len(self.frame_states)}    "
            f"Lumen pts: {len(lumen_x)}    Plaque pts: {len(plaque_x)}    "
            f"Autosave file: {os.path.basename(self.edits_path)}"
        )
        self.fig.canvas.draw_idle()

    def _set_frame(self, frame_idx: int) -> None:
        frame_idx = int(np.clip(frame_idx, 0, len(self.frame_states) - 1))
        if frame_idx == self.current_frame:
            return

        if self.frame_dirty:
            self._append_frame_edit(self.current_frame, reason="frame_change")

        self.current_frame = frame_idx
        self.is_updating_slider = True
        self.slider.set_val(frame_idx)
        self.is_updating_slider = False
        self._render_frame()

    def _nearest_point(self, x: float, y: float) -> tuple[str, int] | None:
        state = self._state()

        def best_idx(xs: list[float], ys: list[float]) -> tuple[int, float] | None:
            if not xs:
                return None
            pts = np.column_stack((np.asarray(xs), np.asarray(ys)))
            dist = np.linalg.norm(pts - np.asarray([x, y]), axis=1)
            idx = int(np.argmin(dist))
            return idx, float(dist[idx])

        lumen = best_idx(state.lumen_x, state.lumen_y)
        plaque = best_idx(state.plaque_x, state.plaque_y)

        choice: tuple[str, int, float] | None = None
        if lumen is not None:
            choice = ("lumen", lumen[0], lumen[1])
        if plaque is not None and (choice is None or plaque[1] < choice[2]):
            choice = ("plaque", plaque[0], plaque[1])

        if choice is None or choice[2] > self.drag_threshold_px:
            return None
        return choice[0], choice[1]

    def _on_slider(self, val: float) -> None:
        if self.is_updating_slider:
            return
        self._set_frame(int(val))

    def _on_prev(self, _event: Any) -> None:
        self._set_frame(self.current_frame - 1)

    def _on_next(self, _event: Any) -> None:
        self._set_frame(self.current_frame + 1)

    def _on_save(self, _event: Any) -> None:
        self._append_frame_edit(self.current_frame, reason="manual_save")
        self._render_frame()

    def _on_toggle_bifurcation(self, _event: Any) -> None:
        state = self._state()
        state.bifurcation = not state.bifurcation
        self.frame_dirty = True
        self._append_frame_edit(self.current_frame, reason="bifurcation_toggle")
        self._render_frame()

    def _on_key(self, event: Any) -> None:
        if event.key in {"left", "a"}:
            self._set_frame(self.current_frame - 1)
        elif event.key in {"right", "d"}:
            self._set_frame(self.current_frame + 1)
        elif event.key == "s":
            self._on_save(event)
        elif event.key == "b":
            self._on_toggle_bifurcation(event)

    def _on_press(self, event: Any) -> None:
        if event.inaxes != self.ax_image or event.button != 1:
            return
        if event.xdata is None or event.ydata is None:
            return

        hit = self._nearest_point(float(event.xdata), float(event.ydata))
        self.drag_target = hit

    def _on_motion(self, event: Any) -> None:
        if self.drag_target is None:
            return
        if event.inaxes != self.ax_image or event.xdata is None or event.ydata is None:
            return

        state = self._state()
        contour_name, idx = self.drag_target

        h, w = self.images[self.current_frame].shape
        x = float(np.clip(event.xdata, 0, w - 1))
        y = float(np.clip(event.ydata, 0, h - 1))

        if contour_name == "lumen":
            state.lumen_x[idx] = x
            state.lumen_y[idx] = y
        else:
            state.plaque_x[idx] = x
            state.plaque_y[idx] = y

        self.frame_dirty = True
        self._render_frame()

    def _on_release(self, _event: Any) -> None:
        if self.drag_target is None:
            return
        self.drag_target = None
        if self.frame_dirty:
            self._append_frame_edit(self.current_frame, reason="point_drag")
            self._render_frame()

    def _on_close(self, _event: Any) -> None:
        if self.frame_dirty:
            self._append_frame_edit(self.current_frame, reason="window_close")

    def show(self) -> None:
        self._build_ui()
        self._render_frame()
        plt.show()


def _default_annotations_path(dicom_path: str, output_root: str) -> str:
    stem = os.path.splitext(os.path.basename(dicom_path))[0]
    pattern = os.path.join(output_root, "*", f"{stem}_contours.jsonl")
    matches = [p for p in glob(pattern) if os.path.isfile(p)]
    if matches:
        matches.sort(key=lambda p: os.path.getmtime(p), reverse=True)
        return matches[0]

    raise FileNotFoundError(
        "No pipeline contour JSONL found for "
        f"'{stem}'. Expected files like: {pattern}"
    )


def _default_edits_path(annotations_path: str) -> str:
    if annotations_path.endswith("_contours.jsonl"):
        return annotations_path.replace("_contours.jsonl", "_edited_annotations.jsonl")
    root, _ = os.path.splitext(annotations_path)
    return root + "_edited_annotations.jsonl"


def launch_annotation_editor(
    dicom_path: str,
    annotations_path: str | None = None,
    edits_path: str | None = None,
    output_root: str = "output",
) -> None:
    """Launch interactive annotation GUI.

    If ``annotations_path`` is not provided, the latest contours JSONL is
    selected from ``output/<timestamp>/``.
    """
    if annotations_path is None:
        annotations_path = _default_annotations_path(dicom_path, output_root)

    if edits_path is None:
        edits_path = _default_edits_path(annotations_path)

    editor = AnnotationEditor(
        dicom_path=dicom_path,
        annotations_path=annotations_path,
        edits_path=edits_path,
    )
    editor.show()