plot_viewer.py

"""Interactive QtAgg matrix dialog and plot-to-PyMOL bridge."""

from __future__ import annotations

import math
from collections.abc import Callable
from typing import Any

import numpy as np

from . import mol_viewer
from .compat import QAction, QtCore, QtWidgets, SizePolicyExpanding
from .matrix_analysis import MatrixAnalysis
from .mol_viewer import CoordinateSnapshot
from .plots import MatrixKind, resolve_color_limits, resolve_colormap

CLICK_DRAG_THRESHOLD_PX = 4.0
MAX_INITIAL_CANVAS_DIMENSION = 900
SELECTION_COLOR = "#ff8c00"
ACTION_COLOR = "#ffd400"
SELECTION_HINT = (
    "Click/drag: select reference residues · "
    "AE only: Cmd/Ctrl-click aligns mobile and measures the column"
)


def _load_qt_backend() -> tuple[type, type]:
    try:
        from matplotlib.backends.backend_qtagg import (
            FigureCanvasQTAgg,
            NavigationToolbar2QT,
        )
    except Exception:
        from matplotlib.backends.backend_qt5agg import (
            FigureCanvasQTAgg,
            NavigationToolbar2QT,
        )
    return FigureCanvasQTAgg, NavigationToolbar2QT


class PlotDialog(QtWidgets.QDialog):
    """One interactive AE or DE matrix."""

    def __init__(
        self,
        figure: Any,
        analysis: MatrixAnalysis,
        kind: MatrixKind,
        reference_name: str,
        mobile_snapshot: CoordinateSnapshot,
        parent: QtWidgets.QWidget | None = None,
        *,
        on_close: Callable[[PlotDialog], None] | None = None,
        on_error: Callable[[str], None] | None = None,
    ) -> None:
        super().__init__(parent)
        self._figure = figure
        self._analysis = analysis
        self._kind = kind
        self._reference_name = reference_name
        self._mobile_snapshot = mobile_snapshot
        self._on_close = on_close
        self._on_error = on_error
        self._axis = figure.axes[0]
        self._press_event: Any | None = None
        self._highlight: Any | None = None
        self._event_ids: list[int] = []
        self._layout_pending = False
        self.setWindowTitle(
            "DistanceMatrix — "
            + ("Aligned error" if kind == "ae" else "Distance error")
        )

        FigureCanvas, NavigationToolbar = _load_qt_backend()
        self._canvas = FigureCanvas(figure)
        self._canvas.setParent(self)
        self._canvas.setSizePolicy(SizePolicyExpanding, SizePolicyExpanding)
        self._canvas.updateGeometry()
        self._toolbar = NavigationToolbar(self._canvas, self, coordinates=True)
        self._clear_action = QAction("Clear selection", self)
        self._clear_action.setToolTip("Clear the managed plot selection")
        self._clear_action.setEnabled(False)
        self._clear_action.triggered.connect(self.clear_selection)
        self._toolbar.addSeparator()
        self._toolbar.addAction(self._clear_action)

        layout = QtWidgets.QVBoxLayout(self)
        layout.addWidget(self._toolbar)
        layout.addWidget(self._canvas)
        self._selection_hint = QtWidgets.QLabel(SELECTION_HINT)
        self._selection_hint.setWordWrap(True)
        layout.addWidget(self._selection_hint)
        self._resize_to_figure_size()

        self._event_ids.extend(
            [
                self._canvas.mpl_connect("button_press_event", self._on_press),
                self._canvas.mpl_connect("button_release_event", self._on_release),
            ]
        )
        self._schedule_fit_figure_to_canvas()

    def _target_canvas_size(self) -> tuple[int, int]:
        """Return an initial canvas size that preserves the figure proportions."""

        try:
            width_in, height_in = self._figure.get_size_inches()
            dpi = float(self._figure.get_dpi())
        except Exception:
            return 900, 650
        width = max(1, int(round(float(width_in) * dpi)))
        height = max(1, int(round(float(height_in) * dpi)))
        max_width = MAX_INITIAL_CANVAS_DIMENSION
        max_height = MAX_INITIAL_CANVAS_DIMENSION
        try:
            screen = self.screen() or QtWidgets.QApplication.primaryScreen()
            if screen is not None:
                available = screen.availableGeometry()
                max_width = min(max_width, max(1, int(available.width() * 0.9)))
                max_height = min(max_height, max(1, int(available.height() * 0.85)))
        except Exception:
            pass
        scale = min(1.0, max_width / width, max_height / height)
        return max(1, int(round(width * scale))), max(1, int(round(height * scale)))

    def _resize_to_figure_size(self) -> None:
        """Include the toolbar and interaction hint around the target canvas."""

        canvas_width, canvas_height = self._target_canvas_size()
        self._canvas.resize(canvas_width, canvas_height)
        layout = self.layout()
        margins = layout.contentsMargins()
        spacing = max(0, int(layout.spacing()))
        dialog_width = canvas_width + margins.left() + margins.right()
        dialog_height = (
            canvas_height
            + self._toolbar.sizeHint().height()
            + self._selection_hint.sizeHint().height()
            + 2 * spacing
            + margins.top()
            + margins.bottom()
        )
        self.resize(dialog_width, dialog_height)

    def _schedule_fit_figure_to_canvas(self) -> None:
        """Apply tight layout after Qt completes the current resize/show step."""

        if self._layout_pending:
            return
        self._layout_pending = True
        QtCore.QTimer.singleShot(0, self._fit_figure_to_canvas)

    def _fit_figure_to_canvas(self) -> None:
        """Recompute layout using the renderer for the actual Qt canvas."""

        self._layout_pending = False
        modern_layout = False
        try:
            self._figure.set_layout_engine("tight")
            modern_layout = True
        except Exception:
            try:
                self._figure.set_tight_layout(True)
            except Exception:
                pass
        try:
            self._canvas.draw()
            if not modern_layout:
                self._figure.tight_layout()
        except Exception:
            pass
        self._canvas.draw_idle()

    def resizeEvent(self, event: Any) -> None:  # noqa: N802
        super().resizeEvent(event)
        self._schedule_fit_figure_to_canvas()

    def _toolbar_active(self) -> bool:
        mode = getattr(self._toolbar, "mode", "")
        name = getattr(mode, "name", None)
        if name is not None:
            return str(name).upper() != "NONE"
        return bool(str(mode))

    @staticmethod
    def _has_action_modifier(event: Any) -> bool:
        names = {
            str(value).lower() for value in (getattr(event, "modifiers", None) or ())
        }
        if names & {"cmd", "super", "meta", "ctrl", "control"}:
            return True
        key = str(getattr(event, "key", "") or "").lower()
        return bool(set(key.split("+")) & {"cmd", "super", "meta", "ctrl", "control"})

    def _on_press(self, event: Any) -> None:
        if (
            event.button != 1
            or event.inaxes is not self._axis
            or self._toolbar_active()
        ):
            self._press_event = None
            return
        self._press_event = event

    def _cell(self, x: float | None, y: float | None) -> tuple[int, int] | None:
        if x is None or y is None:
            return None
        count = len(self._analysis.pairs)
        col = int(np.floor(float(x) + 0.5))
        row = int(np.floor(float(y) + 0.5))
        if row < 0 or col < 0 or row >= count or col >= count:
            return None
        return row, col

    @staticmethod
    def _pixel_distance(start: Any, end: Any) -> float:
        return math.hypot(float(end.x) - float(start.x), float(end.y) - float(start.y))

    def _on_release(self, event: Any) -> None:
        start = self._press_event
        self._press_event = None
        if (
            start is None
            or event.button != 1
            or event.inaxes is not self._axis
            or self._toolbar_active()
        ):
            return
        if self._pixel_distance(start, event) < CLICK_DRAG_THRESHOLD_PX:
            cell = self._cell(event.xdata, event.ydata)
            if cell is None:
                return
            row, col = cell
            if self._kind == "ae" and (
                self._has_action_modifier(start) or self._has_action_modifier(event)
            ):
                self._apply_alignment_action(row, col)
                self._draw_highlight(row, col, row, col, color=ACTION_COLOR)
                return
            self._select_rectangle(row, col, row, col)
            return

        first = self._cell(start.xdata, start.ydata)
        last = self._cell(event.xdata, event.ydata)
        if first is None or last is None:
            return
        self._select_rectangle(first[0], first[1], last[0], last[1])

    def _report_error(self, message: str) -> None:
        if self._on_error is not None:
            self._on_error(message)
        else:
            print(f"DistanceMatrix plot action failed: {message}")

    def _apply_alignment_action(self, row: int, col: int) -> None:
        try:
            mol_viewer.align_and_measure(
                self._reference_name,
                self._mobile_snapshot,
                self._analysis.pairs[row],
                self._analysis.pairs[col],
                self._analysis.rotations[row],
                self._analysis.translations[row],
            )
        except Exception as exc:
            self._report_error(str(exc))

    def _select_rectangle(self, row0: int, col0: int, row1: int, col1: int) -> None:
        row_start, row_stop = sorted((row0, row1))
        col_start, col_stop = sorted((col0, col1))
        tokens = list(range(row_start, row_stop + 1))
        tokens.extend(range(col_start, col_stop + 1))
        tokens = list(dict.fromkeys(tokens))
        try:
            mol_viewer.select_reference_residues(
                self._reference_name, self._analysis.pairs, tokens
            )
        except Exception as exc:
            self._report_error(str(exc))
            return
        self._draw_highlight(
            row_start, col_start, row_stop, col_stop, color=SELECTION_COLOR
        )
        self._clear_action.setEnabled(True)

    def _remove_highlight(self) -> None:
        if self._highlight is None:
            return
        try:
            self._highlight.remove()
        except Exception:
            pass
        self._highlight = None

    def _draw_highlight(
        self,
        row_start: int,
        col_start: int,
        row_stop: int,
        col_stop: int,
        *,
        color: str,
    ) -> None:
        from matplotlib.patches import Rectangle

        self._remove_highlight()
        self._highlight = Rectangle(
            (col_start - 0.5, row_start - 0.5),
            col_stop - col_start + 1,
            row_stop - row_start + 1,
            fill=False,
            edgecolor=color,
            linewidth=2.0,
            zorder=10,
        )
        self._axis.add_patch(self._highlight)
        self._canvas.draw_idle()

    def clear_selection(self) -> None:
        try:
            mol_viewer.select_reference_residues(
                self._reference_name, self._analysis.pairs, []
            )
        except Exception as exc:
            self._report_error(str(exc))
        self._remove_highlight()
        self._clear_action.setEnabled(False)
        self._canvas.draw_idle()

    def update_color_settings(
        self,
        *,
        vmin: float | None,
        vmax: float | None,
        palette: str,
        reverse_palette: bool,
    ) -> None:
        """Update the existing image and colorbar without replacing the dialog."""

        matrix = (
            self._analysis.aligned_error
            if self._kind == "ae"
            else self._analysis.distance_error
        )
        resolved_vmin, resolved_vmax = resolve_color_limits(matrix, vmin, vmax)
        cmap = resolve_colormap(palette, reverse_palette)
        if not self._axis.images:
            raise ValueError("The matrix plot does not contain an image to update.")
        image = self._axis.images[0]
        image.set_cmap(cmap)
        image.set_clim(resolved_vmin, resolved_vmax)
        self._canvas.draw_idle()

    def closeEvent(self, event: Any) -> None:  # noqa: N802
        for event_id in self._event_ids:
            self._canvas.mpl_disconnect(event_id)
        self._event_ids.clear()
        from matplotlib import pyplot as plt

        plt.close(self._figure)
        if self._on_close is not None:
            self._on_close(self)
        super().closeEvent(event)