training_screen

def build(self) -> widgets.Widget:
    """Build and return the training screen layout.

    Returns:
        The root VBox widget for the training screen.
    """
    # Terminal output
    self.out = widgets.Output(
        layout=widgets.Layout(
            border="1px solid white",
            height="200px",
            background_color=BG_COLOR,
            overflow="auto",
        ),
        style={"color": "white"},
    )

    self._header = HeaderWidget(
        self.app, self.out, show_back_button=True, back_to="training_setup"
    )

    # ── Info row ──────────────────────────────────────────────────
    self._state_html = widgets.HTML(
        value=self._render_state_badge(),
        layout=widgets.Layout(flex="0 0 auto", padding="0 8px"),
    )
    self._config_html = widgets.HTML(
        value="",
        layout=widgets.Layout(flex="1 1 auto", padding="0 8px"),
    )
    self._loading_html = widgets.HTML(
        value="",
        layout=widgets.Layout(width="20px"),
    )
    self._countdown_html = widgets.HTML(
        value="",
        layout=widgets.Layout(width="auto", padding="0 4px"),
    )
    # Small, secondary caption next to the refresh spinner for non-tqdm phase
    # hints (Cutana orchestrator init, gallery cutout fetches, etc.).  Kept
    # OUT of the batch-progress ``status_label`` so a preview/catalogue load
    # during scoring can't clobber "Processing batches: X/Y".
    self._preview_hint_html = widgets.HTML(
        value="",
        layout=widgets.Layout(width="auto", padding="0 4px"),
    )
    info_row = HBox(
        [
            self._state_html,
            self._config_html,
            self._loading_html,
            self._preview_hint_html,
            self._countdown_html,
        ],
        layout=widgets.Layout(
            align_items="center",
            padding="2px 12px",
            background_color=ESA_BLUE_DEEP,
            gap="4px",
        ),
    )

    # ── Progress row ──────────────────────────────────────────────
    self.status_html = widgets.HTML(
        value=self._status_text("Ready"),
        layout=widgets.Layout(
            width="auto",
            min_width="140px",
            flex="0 0 auto",
            padding="0 8px 0 0",
        ),
    )

    self.progress = ProgressPanel()
    self.progress.progress_bar.layout.flex = "1 1 auto"
    self.progress.progress_bar.layout.height = "28px"

    self._stop_btn = widgets.Button(
        description="Stop",
        button_style="danger",
        disabled=True,
        layout=widgets.Layout(width="70px", height="28px"),
        style={"button_color": ESA_RED},
    )
    self._stop_btn.on_click(lambda _: self._on_stop())

    progress_row = HBox(
        [
            self.status_html,
            self.progress.spinner,
            self.progress.progress_bar,
            self.progress.status_label,
            self._stop_btn,
        ],
        layout=widgets.Layout(
            gap="8px",
            align_items="center",
            padding="4px 12px",
        ),
    )

    # ── Gallery (left panel) — uses LabelableThumbnailCell ────────
    def _cell_factory(**kwargs: object) -> LabelableThumbnailCell:
        return LabelableThumbnailCell(on_label_change=self._on_label_change, **kwargs)

    self.gallery = GalleryWidget(
        on_magnify=self._on_magnify,
        on_star=self._on_star,
        cell_factory=_cell_factory,
    )
    self.gallery.set_on_sort_change(self._on_sort_change)
    self.gallery.set_on_page_change(self._on_page_change)
    self.gallery.set_on_scale_change(self._on_scale_change)

    left_panel = VBox(
        [self.gallery.widget],
        layout=widgets.Layout(
            flex="1 1 auto",
            min_width="0",
            background_color=BG_COLOR,
        ),
    )

    # ── Right panel — stats, histogram, labels, controls ─────────
    self.stats_html = widgets.HTML(
        value=self._render_stats(0, 0.0, 0.0),
        layout=widgets.Layout(padding="4px 8px"),
    )

    self._histogram_image = widgets.Image(
        format="png",
        layout=widgets.Layout(
            width="300px",
            height="200px",
            object_fit="contain",
            background_color=BG_COLOR,
            display="none",
        ),
    )

    self._label_summary_html = widgets.HTML(
        value=self._render_label_summary(),
        layout=widgets.Layout(padding="4px 8px"),
    )

    # ── Normalisation config ──────────────────────────────────
    norm_cfg = self.cfg.normalisation
    self._norm_widget = NormalisationConfigWidget(
        n_channels=norm_cfg.n_output_channels,
        image_size=list(norm_cfg.image_size),
    )
    self._norm_widget.register_on_change(self._on_norm_change)

    # Training controls
    self._iter_slider = widgets.IntSlider(
        value=self.cfg.num_train_iter if hasattr(self.cfg, "num_train_iter") else 200,
        min=50,
        max=600,
        step=10,
        description="Iterations:",
        style={"description_width": "initial", "handle_color": "white"},
        layout=widgets.Layout(width="280px"),
    )

    # Mirror the setup screen's size-stratification toggle so it's clear it
    # applies to retrains too, and remains adjustable between rounds.  Gated
    # to Cutana sources in ``on_enter`` (only they carry a per-source size).
    # ``on_enter`` owns the value (re-syncs it from cfg each entry); this
    # initial value only matters for standalone construction.
    self._stratify_checkbox = widgets.Checkbox(
        value=bool(self.cfg.cutana_stratify_source_size),
        description="Stratify source size per tile (Cutana)",
        indent=False,
        disabled=True,
        style={"description_width": "initial"},
        layout=widgets.Layout(width="360px"),
    )
    # Write through to cfg on change so cfg is always the source of truth:
    # this lets ``on_enter`` re-sync from cfg without discarding a toggle the
    # user made on the training screen (it's not only pushed at retrain).
    self._stratify_checkbox.observe(self._on_stratify_toggle, names="value")

    self._retrain_warning_html = widgets.HTML(
        value="",
        layout=widgets.Layout(padding="0 8px"),
    )

    self._retrain_btn = widgets.Button(
        description="Retrain",
        button_style="primary",
        layout=widgets.Layout(width="120px", height="36px"),
        style={"font_size": "13px"},
    )
    self._retrain_btn.on_click(lambda _: self._on_retrain())

    self._save_labels_btn = widgets.Button(
        description="Save Labels",
        layout=widgets.Layout(width="120px", height="28px"),
    )
    self._save_labels_btn.on_click(lambda _: self._on_save_labels())

    controls = VBox(
        [
            self._iter_slider,
            self._stratify_checkbox,
            self._retrain_warning_html,
            HBox(
                [self._retrain_btn],
                layout=widgets.Layout(gap="8px", padding="4px 0"),
            ),
            HBox(
                [self._save_labels_btn],
                layout=widgets.Layout(gap="8px", padding="4px 0"),
            ),
        ],
        layout=widgets.Layout(padding="8px"),
    )

    right_panel = VBox(
        [
            widgets.HTML(
                value=(
                    f'<div style="color:{ESA_BLUE_BRIGHT}; font-size:14px;'
                    f' font-weight:bold; padding:4px 8px;">Statistics</div>'
                ),
            ),
            self.stats_html,
            self._histogram_image,
            self._label_summary_html,
            self._norm_widget.widget,
            controls,
        ],
        layout=widgets.Layout(
            width="360px",
            flex="0 0 auto",
            background_color=BG_COLOR,
            border_left=f"1px solid {ESA_BLUE_DEEP}",
            padding="0 8px 0 8px",
            overflow_x="hidden",
        ),
    )

    main_content = HBox(
        [left_panel, right_panel],
        layout=widgets.Layout(
            background_color=BG_COLOR,
            padding="8px 0",
            gap="0px",
        ),
    )

    # Action buttons for busy-state toggling
    self._action_buttons = [self._retrain_btn, self._save_labels_btn]

    root = VBox(
        [
            self._header.widget,
            info_row,
            progress_row,
            main_content,
            self.out,
        ],
        layout=widgets.Layout(background_color=BG_COLOR),
    )
    return root

def on_enter(self) -> None:
    """Set up terminal output and auto-start training on first entry."""
    BackendInterface.set_terminal_output(self.out)
    self._log_sink.attach(self.out)
    self._update_config_summary()

    # Sync extension column headers for Cutana sources BEFORE restoring
    # widget values, so the channel-combination matrix has the right
    # column count when update_from_config fills in cfg values.  Use
    # sync_cutana_bands (not apply_cutana_bands) so the user's
    # n_output_channels choice from the setup screen is preserved.
    is_cutana = (
        bool(self.cfg.data_dir)
        and _auto_detect_source_type(self.cfg.data_dir) == DataSourceType.CUTANA
    )
    if is_cutana:
        filter_names = detect_cutana_filter_names(self.cfg.data_dir, self.cfg)
        if filter_names:
            self._norm_widget.sync_cutana_bands(filter_names)
        else:
            self._norm_widget.show_cutout_zoom(True)

    # Size stratification only applies to Cutana sources; disable the toggle
    # otherwise so a retrain can't carry it into an image/Zarr run.  Re-sync the
    # value from cfg on every entry: screens are cached, so a training screen
    # built before the setup checkbox was ticked would otherwise keep a stale
    # unchecked box and, on ``_start_training``, clobber cfg back to False —
    # silently dropping the user's setup choice.
    self._stratify_checkbox.disabled = not is_cutana
    self._stratify_checkbox.value = is_cutana and bool(self.cfg.cutana_stratify_source_size)

    # Sync normalisation widget from cfg (setup screen may have changed it)
    self._norm_widget.update_from_config(dict(self.cfg.normalisation))
    self._norm_snapshot = self._norm_widget.get_normalisation_config()
    self._retrain_warning_html.value = ""

    # Reopen DB if returning from detail screen
    if self._db_path and BackendInterface._score_db is None:
        BackendInterface.open_prediction_monitor(self._db_path, self.cfg.data_dir or "")

    # If we have existing results, refresh gallery
    if self._db_path:
        count = BackendInterface.get_prediction_count()
        if count > 0:
            self._last_db_count = count
            self._refresh_gallery()
            if count >= 10:
                self._update_histogram()

    # on_leave stops both poll threads so a hidden screen doesn't keep
    # refreshing.  Returning from the image-detail screen therefore lands
    # with polling dead: if scoring or training is still in flight the
    # gallery/progress would freeze until the next button click.  Restart
    # whichever work is still running so the screen stays live — the same
    # guard PredictionScreen.on_enter uses for its own poll loop.
    if (
        self._prediction_thread is not None
        and self._prediction_thread.is_alive()
        and not self._polling
    ):
        self._start_polling()
    # Restart the training poll whenever a subprocess handle is still set
    # and its poll thread is gone.  ``_on_training_complete`` clears
    # ``self._process`` once it has dispatched, so a non-None handle means
    # completion hasn't been handled yet — including the case where the
    # subprocess *exited while the detail screen was open*.  Gating on
    # ``poll() is None`` would skip that exited-while-away case, leaving
    # the screen stuck in the training phase because the loop that calls
    # ``_on_training_complete`` never runs again.  The restarted loop
    # detects the already-exited process and dispatches completion itself.
    if self._process is not None and self._training_poll_thread is None:
        self._start_training_poll()

    # Auto-start training on first entry (no results yet, browsing state)
    if self._state == TrainingState.BROWSING and not self._db_path:
        self._start_training()

def on_leave(self) -> None:
    """Stop polling when navigating away. Keep DB open for detail screen."""
    self._stop_polling()
    self._stop_training_poll()
    self._uninstall_progress_tap()
    self._log_sink.detach()