alnn_rewrite/main.py

import os
import threading
import gc
from typing import Any, Dict

# Must be set before first CUDA context initialization.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import toml
from textual import work
from textual.app import App, ComposeResult
from textual.containers import Horizontal, Vertical
from textual.widgets import (
    Button,
    Footer,
    Header,
    Input,
    Label,
    ProgressBar,
    RichLog,
    Rule,
)

from tasks import PIPELINE_TASKS, SCENARIOS
from util.config_manager import handle_directory_config
from util.progress import ProgressTracker
from util.screens import OverwriteConfirmScreen
from util.ui_logger import PipelineLogger


class CNNHarnessApp(App):
    CSS_PATH = "style.tcss"

    def __init__(self):
        super().__init__()
        self.stop_event = threading.Event()
        self.pause_event = threading.Event()
        self.current_config: Dict[str, Any] = {}
        self.ui_logger = PipelineLogger(self, widget_id="#live_log")

    # ==========================================
    # UI Layout
    # ==========================================
    def compose(self) -> ComposeResult:
        yield Header()

        with Vertical(id="main_container"):
            # --- TOP SECTION: Two-Column Configuration ---
            with Horizontal(id="config_section"):
                # Left Column: Path entry and buttons
                with Vertical(id="config_left_col"):
                    yield Label("Working Directory:")
                    yield Input(placeholder="./outputs/experiment_1", id="work_dir")
                    yield Button("Load / Init Dir", variant="primary", id="load_btn")
                    yield Button("Reload Config", variant="default", id="reload_btn")

                # Right Column: Read-only parameters display
                with Vertical(id="config_right_col"):
                    yield Label("Current Parameters:")
                    yield RichLog(id="config_display", highlight=True, markup=True)

            yield Rule()

            # --- BOTTOM SECTION: Controls, Progress, and Logs ---
            with Horizontal(id="bottom_area"):
                # --- Pipeline Controls ---
                with Vertical(id="control_panel"):
                    yield Button("Start Pipeline", variant="success", id="btn_start")

                    with Horizontal(id="active_controls"):
                        yield Button("Pause", variant="warning", id="btn_pause")
                        yield Button("Stop", variant="error", id="btn_stop")

                    # --- Progress Trackers ---
                    with Vertical(id="progress_area"):
                        for i in range(5):
                            with Horizontal(
                                id=f"progress_row_{i}", classes="progress_row"
                            ):
                                yield Label(
                                    id=f"progress_title_{i}", classes="progress_title"
                                )
                                yield ProgressBar(
                                    id=f"progress_bar_{i}",
                                    classes="progress_bar",
                                    show_eta=True,
                                )
                                yield Label(
                                    id=f"progress_stats_{i}", classes="progress_stats"
                                )

                # --- Live Log Output ---
                with Vertical(id="log_panel"):
                    with Horizontal(id="log_status_row"):
                        yield Label("STATUS: READY", id="pipeline_status_label")
                    with Horizontal(id="log_header_row"):
                        yield Label("Live Logs:", classes="section_label")
                    yield RichLog(id="live_log", highlight=True, markup=True)

        yield Footer()

    def on_mount(self) -> None:
        self.query_one("#active_controls").display = False
        self.ui_logger.info("Application initialized and ready.")

    # ==========================================
    # Helper Methods
    # ==========================================
    def _set_pipeline_status(self, text: str) -> None:
        try:
            self.query_one("#pipeline_status_label", Label).update(text)
        except Exception:
            pass

    def _update_config_display(self) -> None:
        config_log = self.query_one("#config_display", RichLog)
        config_log.clear()

        if not self.current_config:
            config_log.write("[italic]No configuration loaded.[/italic]")
            return

        toml_string = toml.dumps(self.current_config)
        config_log.write(toml_string)

    def _toggle_config_inputs(self, disabled: bool) -> None:
        self.query_one("#work_dir", Input).disabled = disabled
        self.query_one("#load_btn", Button).disabled = disabled
        self.query_one("#reload_btn", Button).disabled = disabled

    # ==========================================
    # Event Handlers
    # ==========================================
    def on_button_pressed(self, event: Button.Pressed) -> None:

        if event.button.id in ("load_btn", "reload_btn"):
            work_dir = self.query_one("#work_dir", Input).value
            success, message, config_data = handle_directory_config(work_dir)

            if success:
                self.current_config = config_data
                self.current_config["work_dir"] = work_dir
                self._update_config_display()
                self.ui_logger.info(message)
            else:
                self.ui_logger.error(message)

        elif event.button.id == "btn_start":
            if not self.current_config:
                self.ui_logger.error(
                    "No configuration loaded! Please load a directory first."
                )
                return

            scenario_id = self.current_config.get("scenario")
            if not scenario_id or scenario_id not in SCENARIOS:
                self.ui_logger.error(
                    f"Invalid or missing scenario '{scenario_id}' in config.toml!"
                )
                return

            work_dir = self.current_config["work_dir"]

            # Check if the working directory exists and is non-empty (except for the config.toml file)
            if os.path.exists(work_dir) and any(
                f for f in os.listdir(work_dir) if f != "config.toml"
            ):

                def check_overwrite_callback(proceed: bool) -> None:
                    if proceed:
                        self._execute_pipeline()
                    else:
                        self.ui_logger.error(
                            "Pipeline start cancelled by user (folder non-empty)."
                        )

                self.app.push_screen(
                    OverwriteConfirmScreen(work_dir),
                    check_overwrite_callback,  # pyright: ignore
                )
            else:
                self._execute_pipeline()

        elif event.button.id == "btn_pause":
            if self.pause_event.is_set():
                self.pause_event.clear()
                event.button.label = "Pause"
                event.button.variant = "warning"
                self._set_pipeline_status("[bold green]PIPELINE RUNNING[/bold green]")
                self.ui_logger.info("Pipeline Resumed.")
            else:
                self.pause_event.set()
                event.button.label = "Resume"
                event.button.variant = "success"
                self._set_pipeline_status("[bold yellow]PIPELINE PAUSED[/bold yellow]")
                self.ui_logger.info(
                    "Pipeline Paused. Waiting for current operation to yield..."
                )

        elif event.button.id == "btn_stop":
            self._set_pipeline_status("[bold red]STOPPING PIPELINE...[/bold red]")
            self.ui_logger.error("Stop requested! Terminating gracefully...")
            event.button.disabled = True
            self.stop_event.set()

    # ==========================================
    # Pipeline Execution
    # ==========================================
    def _execute_pipeline(self) -> None:
        self._toggle_config_inputs(disabled=True)
        self.query_one("#btn_start").display = False
        self.query_one("#active_controls").display = True

        self._set_pipeline_status("[bold green]PIPELINE RUNNING[/bold green]")

        self.stop_event.clear()
        self.pause_event.clear()

        btn_pause = self.query_one("#btn_pause", Button)
        btn_pause.label = "Pause"
        btn_pause.variant = "warning"

        self.run_background_pipeline(self.current_config.copy())

    @work(exclusive=True, thread=True)
    def run_background_pipeline(self, config: Dict[str, Any]) -> None:
        scenario_id = config["scenario"]
        steps = SCENARIOS[scenario_id]["tasks"]
        scenario_label = SCENARIOS[scenario_id]["label"]

        root_tracker = ProgressTracker(self, level=0)
        pipeline_state: Dict[str, Any] = {
            "events": {"stop": self.stop_event, "pause": self.pause_event}
        }

        try:
            self.ui_logger.info(f"Initiating Scenario: {scenario_label}")
            total_tasks = len(steps)

            root_tracker.reset()
            root_tracker.set_title("Pipeline Status")
            root_tracker.update(total=total_tasks, advance=0)

            for i, task_id in enumerate(steps):
                if self.stop_event.is_set():
                    break

                task_name = PIPELINE_TASKS[task_id]["task_name"]
                task_func = PIPELINE_TASKS[task_id]["task_func"]

                task_logger = self.ui_logger.get_task_logger(task_name)
                self.ui_logger.info(f"Starting Task {i + 1}/{total_tasks}: {task_name}")

                task_tracker = root_tracker.get_sub_tracker(task_name)

                task_func(task_tracker, task_logger, config, pipeline_state)

                if not self.stop_event.is_set():
                    task_logger.info("Task completed.")
                    root_tracker.update(advance=1)

            if self.stop_event.is_set():
                self.call_from_thread(
                    self._set_pipeline_status, "[bold red]PIPELINE STOPPED[/bold red]"
                )
                self.ui_logger.error("Pipeline stopped by user.")
                root_tracker.reset()
            else:
                self.call_from_thread(
                    self._set_pipeline_status,
                    "[bold blue]PIPELINE COMPLETE[/bold blue]",
                )
                self.ui_logger.info("All tasks finished successfully!")

        except InterruptedError as e:
            self.ui_logger.error(f"STOPPED: {str(e)}")
            self.call_from_thread(
                self._set_pipeline_status, "[bold red]PIPELINE STOPPED[/bold red]"
            )
            root_tracker.reset()

        except Exception as e:
            self.ui_logger.file_logger.error(
                f"Pipeline failed: {str(e)}", exc_info=True
            )
            self.ui_logger.error(f"ERROR: {str(e)}")
            self.call_from_thread(
                self._set_pipeline_status, "[bold red]PIPELINE FAILED[/bold red]"
            )

        finally:
            pipeline_state.clear()
            gc.collect()
            self.call_from_thread(self._reset_ui)

    def _reset_ui(self) -> None:
        self.query_one("#btn_stop", Button).disabled = False
        self.query_one("#btn_start").display = True
        self.query_one("#active_controls").display = False
        self._toggle_config_inputs(disabled=False)


if __name__ == "__main__":
    app = CNNHarnessApp()
    app.run()