adam/ui/model_intelligence.py

from __future__ import annotations

from pathlib import Path

from PySide6.QtCore import QSize, Qt, QTimer, QUrl, Signal
from PySide6.QtGui import QDesktopServices, QIcon, QImageReader, QPixmap
from PySide6.QtWidgets import (
    QAbstractItemView, QFrame, QGridLayout, QHBoxLayout, QHeaderView, QLabel, QListWidget,
    QListWidgetItem, QPlainTextEdit, QPushButton, QScrollArea, QTableWidget,
    QTableWidgetItem, QVBoxLayout, QWidget,
)

from adam.experiment_tracker import ExperimentStore
from adam.generations import load_generation_history
from adam.intelligence import ModelIntelligence, build_model_intelligence, recommended_training_request
from adam.ui.theme import COLORS
from adam.ui.background import BackgroundRead


def _card() -> QFrame:
    frame = QFrame()
    frame.setProperty("card", True)
    return frame


def _title(text: str) -> QLabel:
    label = QLabel(text)
    label.setObjectName("CardTitle")
    return label


def _thumbnail(path: Path, width: int, height: int) -> QPixmap:
    reader = QImageReader(str(path))
    reader.setAutoTransform(True)
    size = reader.size()
    if size.isValid():
        size.scale(QSize(width, height), Qt.KeepAspectRatio)
        reader.setScaledSize(size)
    image = reader.read()
    return QPixmap.fromImage(image) if not image.isNull() else QPixmap()


class ModelIntelligencePage(QWidget):
    """A visual, evidence-first bridge between experiment history and generations."""

    plan_requested = Signal(str)
    open_experiments_requested = Signal()

    def __init__(self, root: Path, store: ExperimentStore) -> None:
        super().__init__()
        self.root = root.resolve()
        self.store = store
        self.profiles: list[ModelIntelligence] = []
        self.current_key = ""
        root_layout = QVBoxLayout(self)
        root_layout.setContentsMargins(24, 20, 24, 17)
        root_layout.setSpacing(12)

        header = QHBoxLayout()
        head = QWidget()
        head_layout = QVBoxLayout(head)
        head_layout.setContentsMargins(0, 0, 0, 0)
        page_title = QLabel("Model Intelligence")
        page_title.setObjectName("PageTitle")
        detail = QLabel("Turn your local training and generation history into evidence-led follow-up experiments.")
        detail.setProperty("muted", True)
        detail.setWordWrap(True)
        head_layout.addWidget(page_title)
        head_layout.addWidget(detail)
        header.addWidget(head, 1)
        self.refresh_button = QPushButton("Refresh intelligence")
        self.refresh_button.clicked.connect(self.refresh)
        header.addWidget(self.refresh_button, 0, Qt.AlignTop)
        root_layout.addLayout(header)

        summary = QHBoxLayout()
        self.metric_models = self._metric(summary, "MODELS TRACKED")
        self.metric_runs = self._metric(summary, "TRAINING RUNS")
        self.metric_generations = self._metric(summary, "TEST GENERATIONS")
        self.metric_reviews = self._metric(summary, "REVIEW SIGNAL")
        root_layout.addLayout(summary)

        split = QHBoxLayout()
        model_card = _card()
        model_card.setMinimumWidth(265)
        model_layout = QVBoxLayout(model_card)
        model_layout.addWidget(_title("MODEL LIBRARY"))
        model_hint = QLabel("Choose a model to see its evolution, evidence, and next experiment.")
        model_hint.setProperty("muted", True)
        model_hint.setWordWrap(True)
        model_layout.addWidget(model_hint)
        self.models = QListWidget()
        self.models.setSelectionMode(QAbstractItemView.SingleSelection)
        self.models.itemSelectionChanged.connect(self._select_profile)
        model_layout.addWidget(self.models, 1)
        split.addWidget(model_card, 1)

        detail_scroll = QScrollArea()
        detail_scroll.setWidgetResizable(True)
        detail_scroll.setFrameShape(QFrame.NoFrame)
        detail_page = QWidget()
        detail_layout = QVBoxLayout(detail_page)
        detail_layout.setContentsMargins(0, 0, 0, 0)
        detail_layout.setSpacing(12)
        detail_scroll.setWidget(detail_page)

        hero = _card()
        hero_layout = QVBoxLayout(hero)
        hero_layout.setContentsMargins(18, 16, 18, 16)
        self.model_title = QLabel("Select a model")
        self.model_title.setStyleSheet("font-size: 20px; font-weight: 650;")
        self.state = QLabel("Waiting for local experiment records")
        self.state.setProperty("muted", True)
        self.diagnosis = QLabel("ADAM will show a diagnosis only when it has local evidence to support one.")
        self.diagnosis.setWordWrap(True)
        self.recommendation = QLabel("")
        self.recommendation.setWordWrap(True)
        self.recommendation.setStyleSheet(f"color: {COLORS['blue_2']};")
        hero_layout.addWidget(self.model_title)
        hero_layout.addWidget(self.state)
        hero_layout.addWidget(self.diagnosis)
        hero_layout.addWidget(self.recommendation)
        actions = QHBoxLayout()
        self.create_follow_up = QPushButton("Create recommended experiment")
        self.create_follow_up.setProperty("primary", True)
        self.create_follow_up.setEnabled(False)
        self.create_follow_up.clicked.connect(self._create_follow_up)
        self.open_experiments = QPushButton("Open Experiment Tracker")
        self.open_experiments.clicked.connect(self.open_experiments_requested.emit)
        actions.addWidget(self.create_follow_up)
        actions.addWidget(self.open_experiments)
        actions.addStretch(1)
        hero_layout.addLayout(actions)
        detail_layout.addWidget(hero)

        evolution = _card()
        evolution_layout = QVBoxLayout(evolution)
        evolution_layout.setContentsMargins(18, 16, 18, 16)
        evolution_layout.addWidget(_title("EVOLUTION TIMELINE"))
        self.timeline = QTableWidget(0, 7)
        self.timeline.setHorizontalHeaderLabels(["RUN", "STATUS", "EPOCHS", "LOSS", "TIME", "QUALITY", "DATASET"])
        self.timeline.setEditTriggers(QAbstractItemView.NoEditTriggers)
        self.timeline.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.timeline.verticalHeader().hide()
        header_view = self.timeline.horizontalHeader()
        header_view.setStretchLastSection(True)
        for column in range(6):
            header_view.setSectionResizeMode(column, QHeaderView.ResizeToContents)
        self.timeline.setMinimumHeight(180)
        evolution_layout.addWidget(self.timeline)
        detail_layout.addWidget(evolution)

        gallery_card = _card()
        gallery_layout = QVBoxLayout(gallery_card)
        gallery_layout.setContentsMargins(18, 16, 18, 16)
        gallery_layout.addWidget(_title("VISUAL EVIDENCE"))
        self.gallery_summary = QLabel("Training previews and generated images will appear here.")
        self.gallery_summary.setProperty("muted", True)
        self.gallery = QListWidget()
        self.gallery.setViewMode(QListWidget.IconMode)
        self.gallery.setIconSize(QSize(160, 116))
        self.gallery.setGridSize(QSize(184, 158))
        self.gallery.setResizeMode(QListWidget.Adjust)
        self.gallery.setMaximumHeight(220)
        self.gallery.itemDoubleClicked.connect(self._open_image)
        gallery_layout.addWidget(self.gallery_summary)
        gallery_layout.addWidget(self.gallery)
        detail_layout.addWidget(gallery_card)
        detail_layout.addStretch(1)
        split.addWidget(detail_scroll, 3)
        root_layout.addLayout(split, 1)

        # Loading generation metadata can be substantial.  MainWindow creates
        # every workspace at startup, so defer that work until this tab opens.
        self._show_profile(None)
        self._reader = BackgroundRead(self)
        self._reader.ready.connect(self._loaded)
        self._reader.failed.connect(self._load_failed)

    def _metric(self, layout: QHBoxLayout, label: str) -> QLabel:
        card = _card()
        card_layout = QVBoxLayout(card)
        card_layout.setContentsMargins(14, 10, 14, 10)
        heading = QLabel(label)
        heading.setObjectName("CardTitle")
        value = QLabel("0")
        value.setStyleSheet("font-size: 22px; font-weight: 700;")
        card_layout.addWidget(heading)
        card_layout.addWidget(value)
        layout.addWidget(card, 1)
        return value

    def refresh(self) -> None:
        self._request_refresh(force=True)

    def refresh_if_stale(self) -> None:
        self._request_refresh(force=False)

    def _request_refresh(self, *, force: bool) -> None:
        root, store = self.root, self.store

        def read():
            runs = store.list_runs(limit=1000)
            records = load_generation_history(root, limit=250)
            return runs, build_model_intelligence(runs, records, root=root)

        if not self.profiles:
            self.state.setText("Loading local training and generation history…")
        self._reader.request(read, force=force)

    def _load_failed(self, error: str) -> None:
        self.state.setText(f"Could not refresh intelligence: {error}")

    def _loaded(self, result: object) -> None:
        selected = self.current_key
        runs, self.profiles = result
        self.metric_models.setText(str(len(self.profiles)))
        self.metric_runs.setText(str(len(runs)))
        self.metric_generations.setText(str(sum(len(profile.generations) for profile in self.profiles)))
        reviewed = sum(profile.rated_images for profile in self.profiles)
        self.metric_reviews.setText(f"{reviewed} reviewed" if reviewed else "Needs ratings")
        self.models.blockSignals(True)
        self.models.clear()
        for profile in self.profiles:
            item = QListWidgetItem(f"{profile.model_name}\n{profile.architecture.upper()} · {len(profile.runs)} run(s) · {profile.state}")
            item.setData(Qt.UserRole, profile.key)
            self.models.addItem(item)
        self.models.blockSignals(False)
        if self.models.count():
            index = next((index for index in range(self.models.count()) if self.models.item(index).data(Qt.UserRole) == selected), 0)
            self.models.setCurrentRow(index)
        else:
            self._show_profile(None)

    def _select_profile(self) -> None:
        item = self.models.currentItem()
        key = str(item.data(Qt.UserRole) or "") if item else ""
        profile = next((entry for entry in self.profiles if entry.key == key), None)
        self._show_profile(profile)

    def _show_profile(self, profile: ModelIntelligence | None) -> None:
        self.gallery.clear()
        self.timeline.setRowCount(0)
        self.current_key = profile.key if profile else ""
        if profile is None:
            self.model_title.setText("No training records yet")
            self.state.setText("Train a model in ADAM and it will appear here automatically.")
            self.diagnosis.setText("This page only uses local evidence ADAM has actually recorded.")
            self.recommendation.setText("")
            self.gallery_summary.setText("No visual evidence available.")
            self.create_follow_up.setEnabled(False)
            return
        self.model_title.setText(profile.model_name)
        self.state.setText(f"{profile.architecture.upper()} · {profile.state}")
        state_color = COLORS["green"] if profile.state == "Promising" else COLORS["orange"] if "Needs" not in profile.state else COLORS["blue_2"]
        self.state.setStyleSheet(f"color: {state_color}; font-weight: 700;")
        self.diagnosis.setText("Evidence: " + profile.diagnosis)
        self.recommendation.setText("Training Coach: " + profile.recommendation)
        self.create_follow_up.setEnabled(bool(profile.runs and profile.runs[0].dataset_path))
        self.timeline.setRowCount(len(profile.runs))
        for row, run in enumerate(profile.runs):
            values = [
                run.id, run.status, str(run.epochs or "—"),
                "—" if run.final_loss is None else f"{run.final_loss:.5f}",
                self._duration(run.training_time_seconds),
                "—" if run.quality_score is None else f"{run.quality_score}/100",
                run.dataset_name or Path(run.dataset_path).name or "—",
            ]
            for column, value in enumerate(values):
                item = QTableWidgetItem(value)
                if column in {1, 2, 3, 4, 5}:
                    item.setTextAlignment(Qt.AlignCenter)
                self.timeline.setItem(row, column, item)
        paths: list[tuple[Path, str]] = []
        for run in profile.runs:
            paths.extend((Path(path), f"Training preview · {run.id}") for path in run.preview_images)
        for record in profile.generations:
            paths.extend((path, f"Generation · seed {record.seed + index}") for index, path in enumerate(record.images))
        deduplicated: dict[str, tuple[Path, str]] = {}
        for path, label in paths:
            key = str(path)
            if key not in deduplicated and path.is_file():
                deduplicated[key] = (path, label)
            if len(deduplicated) >= 24:
                break
        for path, label in list(deduplicated.values())[:24]:
            pixmap = _thumbnail(path, 160, 116)
            item = QListWidgetItem(QIcon(pixmap), label)
            item.setData(Qt.UserRole, str(path))
            self.gallery.addItem(item)
        self.gallery_summary.setText(
            f"{len(profile.runs)} training run(s) · {len(profile.generations)} generation batch(es) · "
            f"{profile.rated_images} reviewed image signal(s). Double-click an image to open it."
        )

    def _create_follow_up(self) -> None:
        profile = next((entry for entry in self.profiles if entry.key == self.current_key), None)
        if profile:
            self.plan_requested.emit(recommended_training_request(profile))

    def _open_image(self, item: QListWidgetItem) -> None:
        path = str(item.data(Qt.UserRole) or "")
        if path:
            QDesktopServices.openUrl(QUrl.fromLocalFile(path))

    @staticmethod
    def _duration(seconds: int) -> str:
        if not seconds:
            return "—"
        if seconds >= 3600:
            return f"{seconds // 3600}h {(seconds % 3600) // 60}m"
        return f"{seconds // 60}m {seconds % 60}s"