sras-viewer/sras_viewer/dialogs.py

"""FFT option dialogs.

Angle alignment used to live here too, as ManualAlignmentDialog; it is now the
alignment wizard's first page (see align_wizard.py), which needed the same
mask-overlay editor plus the crop and export steps.
"""

from pathlib import Path

from PyQt6.QtWidgets import (
    QButtonGroup, QCheckBox, QDialog, QDialogButtonBox, QFileDialog,
    QGroupBox, QHBoxLayout, QLabel, QLineEdit, QPushButton, QRadioButton,
    QScrollArea, QSpinBox, QVBoxLayout, QWidget,
)

from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES

from .common import CH_LABELS, VELOCITY_MODE_IDX, _CSS_HINT, _CSS_WARN, _group, _make_dspin, _wrap_label


# ---------------------------------------------------------------------------
# FFT Options dialog
# ---------------------------------------------------------------------------

class FftOptionsDialog(QDialog):
    """Configure FFT zero-padding.

    Changes take effect only when the user clicks Apply.  Cancel discards
    all pending edits.  The live 'frequency resolution' label updates as
    the user adjusts the pad factor so they can see the trade-off before
    committing.
    """

    def __init__(self, parent=None, *,
                 current_pad_factor: int,
                 samples_per_frame: int | None,
                 sample_rate_hz: float | None,
                 grating_um: float):
        super().__init__(parent)
        self.setWindowTitle("FFT Options")
        self.setModal(True)
        self.setMinimumWidth(380)

        self._samples_per_frame = samples_per_frame
        self._sample_rate_hz = sample_rate_hz
        self._grating_um = grating_um

        layout = QVBoxLayout(self)

        # ---- Zero-padding ----------------------------------------------
        grp_zp = QGroupBox("Zero-Padding")
        zl = QVBoxLayout(grp_zp)

        pad_row = QHBoxLayout()
        pad_row.addWidget(QLabel("Pad factor:"))
        self._spin_pad = QSpinBox()
        self._spin_pad.setRange(1, 256)
        self._spin_pad.setValue(max(1, current_pad_factor))
        self._spin_pad.setToolTip(
            "Multiply the waveform length by this factor via zero-padding\n"
            "before computing the FFT.\n"
            "1 = no padding (natural length).\n"
            "Powers of 2 (2, 4, 8 …) give the best performance."
        )
        self._spin_pad.valueChanged.connect(self._update_info)
        pad_row.addWidget(self._spin_pad)
        zl.addLayout(pad_row)

        self._lbl_nfft = QLabel()
        self._lbl_freq_res = QLabel()
        self._lbl_vel_res = QLabel()
        for lbl in (self._lbl_nfft, self._lbl_freq_res, self._lbl_vel_res):
            lbl.setStyleSheet(_CSS_HINT)
            zl.addWidget(lbl)

        layout.addWidget(grp_zp)

        # ---- Buttons ---------------------------------------------------
        buttons = QDialogButtonBox()
        buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole
                          ).clicked.connect(self.accept)
        buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
                          ).clicked.connect(self.reject)
        layout.addWidget(buttons)

        self._update_info()

    def _update_info(self):
        spf = self._samples_per_frame
        sr = self._sample_rate_hz
        pad = self._spin_pad.value()

        if spf is None or sr is None:
            self._lbl_nfft.setText("Load a file to preview FFT parameters.")
            self._lbl_freq_res.setText("")
            self._lbl_vel_res.setText("")
            return

        n_fft = spf * pad
        freq_res_hz = sr / n_fft
        freq_res_mhz = freq_res_hz / 1e6
        # v (m/s) = freq (MHz) × grating (µm)
        vel_res_ms = freq_res_mhz * self._grating_um

        self._lbl_nfft.setText(f"FFT points: {spf} × {pad} = {n_fft:,}")
        self._lbl_freq_res.setText(
            f"Frequency bin: {freq_res_mhz:.4f} MHz  ({freq_res_hz / 1e3:.2f} kHz)")
        self._lbl_vel_res.setText(
            f"Velocity bin: {vel_res_ms:.3f} m/s  "
            f"(at grating = {self._grating_um:.2f} µm)")

    def get_pad_factor(self) -> int:
        return max(1, self._spin_pad.value())


# ---------------------------------------------------------------------------
# Row-Averaged FFT Options dialog
# ---------------------------------------------------------------------------

class RowAverageFftOptionsDialog(QDialog):
    """Configure the same-row, distance-weighted neighbor averaging applied
    to each pixel's CH1 waveform before 'Batch Compute Row-Averaged FFT and
    Store' re-runs the FFT peak search — a same-row SNR cleanup pass, never
    mixing across rows/Y (see sras_compute._row_average_waveforms).

    Unlike the plain FFT batch action (which stores unmasked and defers
    masking to display time), the DC threshold here is required up front:
    it decides which same-row neighbors are eligible to contribute to a
    pixel's average, so it can't be deferred.

    Changes take effect only when the user clicks Apply. Cancel discards
    all pending edits.
    """

    def __init__(self, parent=None, *,
                 current_n: int,
                 current_threshold_mv: float,
                 pixel_x_mm: float | None):
        super().__init__(parent)
        self.setWindowTitle("Row-Averaged FFT Options")
        self.setModal(True)
        self.setMinimumWidth(380)

        self._pixel_x_mm = pixel_x_mm

        layout = QVBoxLayout(self)

        # ---- Neighbor window ---------------------------------------------
        grp_window = QGroupBox("Same-Row Neighbor Window")
        wl = QVBoxLayout(grp_window)

        n_row = QHBoxLayout()
        n_row.addWidget(QLabel("Neighbor half-width (n):"))
        self._spin_n = QSpinBox()
        self._spin_n.setRange(1, 50)
        self._spin_n.setValue(max(1, current_n))
        self._spin_n.setToolTip(
            "Each pixel's CH1 waveform is averaged with up to n same-row\n"
            "neighbors on each side, distance-weighted (Gaussian) and\n"
            "counting only neighbors that already pass the DC threshold\n"
            "below. Never mixes across rows/Y.")
        self._spin_n.valueChanged.connect(self._update_info)
        n_row.addWidget(self._spin_n)
        wl.addLayout(n_row)

        self._lbl_width = QLabel()
        self._lbl_width.setStyleSheet(_CSS_HINT)
        wl.addWidget(self._lbl_width)

        layout.addWidget(grp_window)

        # ---- DC threshold ------------------------------------------------
        grp_thr = QGroupBox("Neighbor Validity")
        tl = QVBoxLayout(grp_thr)
        thr_row = QHBoxLayout()
        thr_row.addWidget(QLabel("DC threshold:"))
        self._spin_threshold = _make_dspin(-500.0, 500.0, 3, suffix=" mV",
                                           value=current_threshold_mv, step=0.025)
        self._spin_threshold.setToolTip(
            "A same-row neighbor only contributes to a pixel's average if\n"
            "its own CH4 signal is at or above this threshold -- the same\n"
            "test used for RF mask display. A pixel below threshold stays\n"
            "masked, exactly as today; it is never rescued by its neighbors.")
        thr_row.addWidget(self._spin_threshold)
        tl.addLayout(thr_row)
        layout.addWidget(grp_thr)

        # ---- Buttons -----------------------------------------------------
        buttons = QDialogButtonBox()
        buttons.addButton("Apply", QDialogButtonBox.ButtonRole.AcceptRole
                          ).clicked.connect(self.accept)
        buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
                          ).clicked.connect(self.reject)
        layout.addWidget(buttons)

        self._update_info()

    def _update_info(self):
        n = self._spin_n.value()
        if self._pixel_x_mm is None:
            self._lbl_width.setText("Load a file to preview the window's physical width.")
            return
        width_um = 2 * n * self._pixel_x_mm * 1e3
        self._lbl_width.setText(
            f"Window: ±{n} px = {width_um:.2f} µm full width "
            f"(pixel pitch {self._pixel_x_mm * 1e3:.3g} µm)")

    def get_half_width(self) -> int:
        return self._spin_n.value()

    def get_threshold_mv(self) -> float:
        return self._spin_threshold.value()


# ---------------------------------------------------------------------------
# Export Fused ROI dialog
# ---------------------------------------------------------------------------

class FusedRoiExportDialog(QDialog):
    """Choose a value type and which angles to fuse for Export Fused ROI.

    Every angle in the file is listed (a live AlignmentResult's per_angle
    always covers every angle, and the raw-shared-grid path needs no per-
    angle transform at all). Each checkbox is disabled — and auto-unchecked
    — whenever *availability_fn(angle_idx, ch_idx)* is False for the
    currently selected value type; switching the value-type radio
    re-evaluates every checkbox live, since availability is per (angle,
    value type) rather than just per angle — e.g. DC may be ready
    everywhere while FFT is ready nowhere.
    """

    _VALUE_MODES = (CH1_IDX, CH3_IDX, CH4_IDX, VELOCITY_MODE_IDX)

    def __init__(self, parent=None, *,
                 angles: list[tuple[int, float]],
                 availability_fn,
                 default_ch_idx: int,
                 out_dir: str,
                 stem: str,
                 grid_note: str):
        super().__init__(parent)
        self.setWindowTitle("Export Fused ROI")
        self.setModal(True)
        self.setMinimumWidth(420)

        self._availability_fn = availability_fn
        self._out_dir = out_dir
        self._stem = stem
        self._path_user_chosen = False

        layout = QVBoxLayout(self)
        layout.addWidget(_wrap_label(grid_note, _CSS_HINT))

        # ---- Value type --------------------------------------------------
        grp_val, vl = _group("Value to Export")
        self._val_group = QButtonGroup(self)
        self._val_buttons: dict[int, QRadioButton] = {}
        for ch_idx, label in zip(self._VALUE_MODES, CH_LABELS):
            rb = QRadioButton(label)
            self._val_group.addButton(rb, id=ch_idx)
            self._val_buttons[ch_idx] = rb
            vl.addWidget(rb)
        self._val_buttons[default_ch_idx].setChecked(True)
        self._val_group.idClicked.connect(self._on_value_type_changed)
        layout.addWidget(grp_val)

        # ---- Angles --------------------------------------------------
        grp_ang, al = _group("Angles to Include")
        sel_row = QHBoxLayout()
        btn_all = QPushButton("Select All Available")
        btn_none = QPushButton("Select None")
        btn_all.clicked.connect(self._on_select_all_available)
        btn_none.clicked.connect(self._on_select_none)
        sel_row.addWidget(btn_all)
        sel_row.addWidget(btn_none)
        al.addLayout(sel_row)

        scroll_inner = QWidget()
        scroll_layout = QVBoxLayout(scroll_inner)
        self._angle_checks: dict[int, QCheckBox] = {}
        for angle_idx, angle_deg in angles:
            cb = QCheckBox(f"{angle_deg:.1f}°  (angle {angle_idx})")
            self._angle_checks[angle_idx] = cb
            cb.toggled.connect(self._update_accept_enabled)
            scroll_layout.addWidget(cb)
        scroll = QScrollArea()
        scroll.setWidget(scroll_inner)
        scroll.setWidgetResizable(True)
        scroll.setMaximumHeight(220)
        al.addWidget(scroll)

        self._lbl_none_available = _wrap_label("", _CSS_WARN)
        al.addWidget(self._lbl_none_available)
        layout.addWidget(grp_ang)

        # ---- Output path --------------------------------------------------
        grp_out, ol = _group("Output File")
        path_row = QHBoxLayout()
        self._edit_path = QLineEdit()
        self._edit_path.setReadOnly(True)
        path_row.addWidget(self._edit_path, 1)
        btn_browse = QPushButton("Browse…")
        btn_browse.clicked.connect(self._on_browse)
        path_row.addWidget(btn_browse)
        ol.addLayout(path_row)
        layout.addWidget(grp_out)

        # ---- Buttons -----------------------------------------------------
        buttons = QDialogButtonBox()
        self._btn_export = buttons.addButton(
            "Export", QDialogButtonBox.ButtonRole.AcceptRole)
        self._btn_export.clicked.connect(self.accept)
        buttons.addButton("Cancel", QDialogButtonBox.ButtonRole.RejectRole
                          ).clicked.connect(self.reject)
        layout.addWidget(buttons)

        self._refresh_default_path()
        self._apply_availability()

    # ---- internals -----------------------------------------------------

    def _current_ch_idx(self) -> int:
        return self._val_group.checkedId()

    def _apply_availability(self):
        ch_idx = self._current_ch_idx()
        n_ok = 0
        for angle_idx, cb in self._angle_checks.items():
            ok = self._availability_fn(angle_idx, ch_idx)
            cb.setEnabled(ok)
            if ok:
                n_ok += 1
                cb.setToolTip("")
            else:
                cb.setChecked(False)
                cb.setToolTip(
                    f"No cached/stored {CH_LABELS[ch_idx]} data for this "
                    "angle yet — view it in the main window (or run "
                    "Batch Compute) first.")
        self._lbl_none_available.setText(
            "" if n_ok else "No angle has this value type ready yet.")
        self._update_accept_enabled()

    def _on_value_type_changed(self, _id: int):
        self._apply_availability()
        if not self._path_user_chosen:
            self._refresh_default_path()

    def _on_select_all_available(self):
        for cb in self._angle_checks.values():
            if cb.isEnabled():
                cb.setChecked(True)

    def _on_select_none(self):
        for cb in self._angle_checks.values():
            cb.setChecked(False)

    def _refresh_default_path(self):
        ch_idx = self._current_ch_idx()
        name = f"{self._stem}_fused_roi_{CH_NAMES[ch_idx]}.csv"
        self._edit_path.setText(str(Path(self._out_dir) / name))
        self._update_accept_enabled()

    def _on_browse(self):
        path, _ = QFileDialog.getSaveFileName(
            self, "Export Fused ROI as CSV", self._edit_path.text(),
            "CSV files (*.csv);;All files (*)")
        if path:
            self._edit_path.setText(path)
            self._path_user_chosen = True
            self._update_accept_enabled()

    def _update_accept_enabled(self):
        any_checked = any(cb.isChecked() for cb in self._angle_checks.values())
        self._btn_export.setEnabled(any_checked and bool(self._edit_path.text()))

    # ---- getters ---------------------------------------------------------

    def get_ch_idx(self) -> int:
        return self._current_ch_idx()

    def get_selected_angles(self) -> list[int]:
        return sorted(a for a, cb in self._angle_checks.items() if cb.isChecked())

    def get_output_path(self) -> str:
        return self._edit_path.text()