alnn_rewrite/util/progress.py

import logging

from textual.app import App
from textual.containers import Horizontal
from textual.widgets import Label, ProgressBar

logger = logging.getLogger(__name__)

# Number of stacked progress-bar rows (levels 0 .. MAX_LEVELS-1). Deep enough for
# the k-fold nesting: root(0) -> run_kfold(1) -> fold task(2) -> train/eval(3) ->
# epochs/passes(4) -> batches(5). main.py builds exactly this many rows.
MAX_LEVELS = 6


class ProgressTracker:
    """Manages hierarchical progress bars for the UI thread."""

    def __init__(self, app: App, level: int = 0):
        self.app = app
        self.level = level

    def get_sub_tracker(self, title: str = "") -> "ProgressTracker":
        """Returns a sub-tracker for child processes."""
        if self.level >= MAX_LEVELS - 1:
            return self  # Max depth reached
        sub = ProgressTracker(self.app, self.level + 1)
        if title:
            sub.set_title(title)
        return sub

    def set_title(self, title: str) -> None:
        def _update_title():
            try:
                self.app.query_one(f"#progress_title_{self.level}", Label).update(title)
            except Exception:
                pass

        self.app.call_from_thread(_update_title)

    def update(self, total: int | None = None, advance: int = 0) -> None:
        def _do_update():
            try:
                bar = self.app.query_one(f"#progress_bar_{self.level}", ProgressBar)
                row = self.app.query_one(f"#progress_row_{self.level}", Horizontal)
                stats = self.app.query_one(f"#progress_stats_{self.level}", Label)

                if total is not None:
                    bar.update(total=total)
                if advance > 0:
                    bar.advance(advance)

                # Format Completed / Total label
                if bar.total is not None and bar.total > 0:
                    stats.update(f"{int(bar.progress)} / {int(bar.total)}")
                else:
                    stats.update("")

                # Visibility logic: Visible only when value/total is non-zero
                is_active = (bar.progress > 0) or (
                    bar.total is not None and bar.total > 0
                )
                row.display = is_active

                # Cascading clear logic: Clear sub-bars if we advance or reset
                if advance > 0 or total == 0:
                    self._clear_subs()
            except Exception as e:
                logger.error(f"Failed to update progress: {e}")

        self.app.call_from_thread(_do_update)

    def reset(self) -> None:
        """Clears and hides this progress bar and cascades to all children."""

        def _do_reset():
            try:
                bar = self.app.query_one(f"#progress_bar_{self.level}", ProgressBar)
                row = self.app.query_one(f"#progress_row_{self.level}", Horizontal)
                title = self.app.query_one(f"#progress_title_{self.level}", Label)
                stats = self.app.query_one(f"#progress_stats_{self.level}", Label)

                bar.progress = 0
                bar.total = None
                title.update("")
                stats.update("")
                row.display = False

                self._clear_subs()
            except Exception:
                pass

                pass
        self.app.call_from_thread(_do_reset)

    def _clear_subs(self) -> None:
        """Clears and hides all child progress bars (must run in UI thread)."""
        for l in range(self.level + 1, MAX_LEVELS):
            try:
                bar = self.app.query_one(f"#progress_bar_{l}", ProgressBar)
                row = self.app.query_one(f"#progress_row_{l}", Horizontal)
                title = self.app.query_one(f"#progress_title_{l}", Label)
                stats = self.app.query_one(f"#progress_stats_{l}", Label)

                bar.progress = 0
                bar.total = None
                title.update("")
                stats.update("")
                row.display = False
            except Exception:
                pass