Trixy7/trixy_core/tui/views/trainer_training_view.py

# -*- coding: utf-8 -*-
"""
Trainer SubView: Training (F6).

Steuerung und Visualisierung des Trainings:
Start/Pause/Resume/Stop Buttons, Metriken-Zeile und LossChart.
"""

from __future__ import annotations

from textual.app import ComposeResult
from textual.containers import Horizontal, Vertical
from textual.widgets import Static, Button

from trixy_core.tui.views.base import MainView
from trixy_core.tui.widgets.loss_chart import LossChart


class TrainerTrainingView(MainView):
    """F6 im Trainer-Kontext: Training-Steuerung und Visualisierung."""

    TITLE = "Training"
    VIEW_ID = "trainer_training"

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
        self._trainer_data: dict = {}
        self._is_training: bool = False
        self._is_paused: bool = False

    def set_trainer(self, trainer_data: dict) -> None:
        """Setzt die Trainer-Basisdaten."""
        self._trainer_data = trainer_data
        self._is_training = False
        self._is_paused = False

    def compose(self) -> ComposeResult:
        # Steuerungsbuttons
        with Horizontal(id="trainer-controls"):
            yield Button("Start", id="trainer-btn-start", variant="success")
            yield Button("Pause", id="trainer-btn-pause", variant="warning", disabled=True)
            yield Button("Resume", id="trainer-btn-resume", variant="primary", disabled=True)
            yield Button("Stop", id="trainer-btn-stop", variant="error", disabled=True)

        # Metriken-Leiste
        yield Static("", id="trainer-metrics")

        # Status
        yield Static("", id="trainer-training-status")

        # Loss-Chart
        with Vertical(classes="trainer-section"):
            yield Static("[bold]Loss-Verlauf[/bold]", classes="trainer-section-header")
            yield LossChart(id="trainer-chart")

    async def on_activate(self) -> None:
        """Prueft ob ein Training laeuft."""
        await self._poll_progress()

    async def refresh_data(self) -> None:
        """Pollt Progress waehrend Training (alle 2s durch _periodic_refresh)."""
        if self._is_training:
            await self._poll_progress()

    async def _poll_progress(self) -> None:
        """Holt den aktuellen Trainings-Fortschritt."""
        if not self.connection:
            return

        try:
            progress = await self.connection.request_training_progress()
        except Exception:
            return

        if progress is None:
            return

        state = progress.state if hasattr(progress, "state") else "idle"
        self._is_training = state in ("training", "paused", "starting")
        self._is_paused = state == "paused"

        self._update_buttons(state)
        self._update_metrics(progress)
        self._update_chart(progress)

    def _update_buttons(self, state: str) -> None:
        """Aktualisiert Button-Zustaende."""
        start_btn = self.query_one("#trainer-btn-start", Button)
        pause_btn = self.query_one("#trainer-btn-pause", Button)
        resume_btn = self.query_one("#trainer-btn-resume", Button)
        stop_btn = self.query_one("#trainer-btn-stop", Button)

        is_active = state in ("training", "paused", "starting")

        start_btn.disabled = is_active
        pause_btn.disabled = state != "training"
        resume_btn.disabled = state != "paused"
        stop_btn.disabled = not is_active

    def _update_metrics(self, progress: object) -> None:
        """Aktualisiert die Metriken-Leiste."""
        metrics = self.query_one("#trainer-metrics", Static)
        state = getattr(progress, "state", "idle")
        epoch = getattr(progress, "epoch", 0)
        total = getattr(progress, "total_epochs", 0)
        loss = getattr(progress, "loss", 0.0)
        val_loss = getattr(progress, "val_loss", 0.0)
        accuracy = getattr(progress, "accuracy", 0.0)
        lr = getattr(progress, "learning_rate", 0.0)
        eta = getattr(progress, "eta_seconds", 0.0)
        vram = getattr(progress, "vram_used_mb", 0.0)

        # Status-Farbe
        state_colors = {
            "idle": "dim",
            "starting": "yellow",
            "training": "green",
            "paused": "yellow",
            "stopping": "yellow",
            "completed": "cyan",
            "failed": "red",
        }
        color = state_colors.get(state, "dim")

        # ETA formatieren
        eta_str = ""
        if eta > 0:
            mins, secs = divmod(int(eta), 60)
            hours, mins = divmod(mins, 60)
            if hours > 0:
                eta_str = f"{hours}:{mins:02d}:{secs:02d}"
            else:
                eta_str = f"{mins:02d}:{secs:02d}"

        parts = [
            f"[{color}]{state.upper()}[/{color}]",
            f"Epoche: {epoch}/{total}" if total > 0 else "",
            f"Loss: {loss:.4f}" if loss > 0 else "",
            f"Val: {val_loss:.4f}" if val_loss > 0 else "",
            f"Acc: {accuracy:.1f}%" if accuracy > 0 else "",
            f"LR: {lr:.1e}" if lr > 0 else "",
            f"ETA: {eta_str}" if eta_str else "",
            f"VRAM: {vram:.1f} GB" if vram > 0 else "",
        ]
        metrics.update(" | ".join(p for p in parts if p))

        # Status-Nachricht
        status = self.query_one("#trainer-training-status", Static)
        message = getattr(progress, "message", "")
        error = getattr(progress, "error", "")
        if error:
            status.update(f"[red]{error}[/red]")
        elif message:
            status.update(f"[dim]{message}[/dim]")

    def _update_chart(self, progress: object) -> None:
        """Aktualisiert den Loss-Chart."""
        chart = self.query_one("#trainer-chart", LossChart)
        loss_history = getattr(progress, "loss_history", [])
        val_loss_history = getattr(progress, "val_loss_history", [])

        if loss_history:
            chart.update_data(loss_history, val_loss_history or None)

    async def on_button_pressed(self, event: Button.Pressed) -> None:
        """Verarbeitet Button-Klicks."""
        btn_id = event.button.id
        if not self.connection:
            return

        trainer_id = self._trainer_data.get("trainer_id", "")
        status = self.query_one("#trainer-training-status", Static)

        if btn_id == "trainer-btn-start":
            if not trainer_id:
                status.update("[red]Kein Trainer ausgewaehlt[/red]")
                return

            status.update("[yellow]Training wird gestartet...[/yellow]")
            try:
                result = await self.connection.start_training(trainer_id)
                if result and result.success:
                    self._is_training = True
                    self._update_buttons("training")
                    status.update("[green]Training gestartet[/green]")
                else:
                    error = result.error if result else "Unbekannter Fehler"
                    status.update(f"[red]{error}[/red]")
            except Exception as e:
                status.update(f"[red]Fehler: {e}[/red]")

        elif btn_id == "trainer-btn-pause":
            try:
                result = await self.connection.pause_training()
                if result and result.success:
                    self._is_paused = True
                    self._update_buttons("paused")
                    status.update("[yellow]Training pausiert[/yellow]")
            except Exception as e:
                status.update(f"[red]Fehler: {e}[/red]")

        elif btn_id == "trainer-btn-resume":
            try:
                result = await self.connection.resume_training()
                if result and result.success:
                    self._is_paused = False
                    self._update_buttons("training")
                    status.update("[green]Training fortgesetzt[/green]")
            except Exception as e:
                status.update(f"[red]Fehler: {e}[/red]")

        elif btn_id == "trainer-btn-stop":
            status.update("[yellow]Training wird gestoppt...[/yellow]")
            try:
                result = await self.connection.stop_training()
                if result and result.success:
                    self._is_training = False
                    self._is_paused = False
                    self._update_buttons("idle")
                    status.update("[cyan]Training gestoppt[/cyan]")
            except Exception as e:
                status.update(f"[red]Fehler: {e}[/red]")