adam/intelligence.py

"""Evidence-led summaries for ADAM's saved training and generation history.

This module intentionally stays independent of the Qt interface so its advice can
be tested and reused by a future remote surface.  It does not assess a model's
absolute quality: it identifies useful follow-up experiments from the evidence
ADAM has recorded locally.
"""
from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

from adam.experiment_tracker import ExperimentRun
from adam.generations import GenerationRecord
from adam.image_preferences import PreferenceProfile


@dataclass(frozen=True, slots=True)
class ModelIntelligence:
    key: str
    model_name: str
    architecture: str
    runs: tuple[ExperimentRun, ...]
    generations: tuple[GenerationRecord, ...]
    rated_images: int
    positive_ratings: int
    rejected_images: int
    state: str
    diagnosis: str
    recommendation: str
    recommended_epochs: int


def _normalized(value: str) -> str:
    return " ".join(str(value or "").casefold().split())


def _paths_match(left: str, right: str, cache: dict[str, Path] | None = None) -> bool:
    if not left or not right:
        return False
    try:
        cache = cache if cache is not None else {}
        for value in (left, right):
            if value not in cache:
                cache[value] = Path(value).expanduser().resolve()
        left_path, right_path = cache[left], cache[right]
        return left_path == right_path or left_path in right_path.parents or right_path in left_path.parents
    except OSError:
        return _normalized(left) == _normalized(right)


def _generation_matches(run: ExperimentRun, record: GenerationRecord, cache: dict[str, Path] | None = None) -> bool:
    if _paths_match(run.output_folder, record.model_path, cache):
        return True
    if any(_paths_match(path, record.model_path, cache) for path in run.checkpoint_paths):
        return True
    return _normalized(run.model_name) == _normalized(record.model_name)


def _ratings(records: tuple[GenerationRecord, ...], root: Path | None = None) -> tuple[int, int, int]:
    rated = positive = rejected = 0
    seen: set[str] = set()
    preferences_by_model: dict[tuple[str, str], PreferenceProfile] = {}
    image_keys: dict[Path, str] = {}
    for record in records:
        preferences = None
        if root:
            key = (record.provider_id, record.model_path)
            if key not in preferences_by_model:
                preferences_by_model[key] = PreferenceProfile(root, record.provider_id, record.model_name, record.model_path)
            preferences = preferences_by_model[key]
        for image in record.images:
            if image not in image_keys:
                image_keys[image] = str(image.expanduser().resolve())
            image_key = image_keys[image]
            if image_key in seen:
                continue
            seen.add(image_key)
            evaluation = record.image_evaluations.get(image_key, {})
            saved_rating = preferences.rating_for(image) if preferences else None
            rating = saved_rating.rating if saved_rating else str(evaluation.get("rating", "")).casefold()
            # Preference ratings are persisted separately today.  The evaluator
            # score still counts as review evidence when it is available.
            score = evaluation.get("score")
            if rating or isinstance(score, (int, float)):
                rated += 1
            if rating in {"favorite", "keep"} or (isinstance(score, (int, float)) and score >= 0.70):
                positive += 1
            if rating == "reject" or (isinstance(score, (int, float)) and score <= 0.35):
                rejected += 1
    return rated, positive, rejected


def _diagnose(
    runs: tuple[ExperimentRun, ...], generations: tuple[GenerationRecord, ...], rated: int, positive: int, rejected: int,
) -> tuple[str, str, str, int]:
    latest = runs[0]
    finished = [run for run in runs if run.status.casefold() == "finished"]
    epoch_budget = max(1, latest.epochs)
    previous = runs[1] if len(runs) > 1 else None

    if not finished:
        return (
            "Needs a completed run",
            "ADAM has not recorded a finished training run for this model yet, so it cannot judge training behavior.",
            "Finish one run and generate a small, fixed-prompt test batch before changing several settings at once.",
            epoch_budget,
        )
    if rated == 0:
        return (
            "Needs visual review",
            "Training history exists, but there are no scored test generations linked to this model. Loss alone cannot tell ADAM which checkpoint you prefer.",
            "Generate 4–8 images with one repeatable prompt and seed, then rate the results in Generations before starting a follow-up.",
            epoch_budget,
        )
    if rejected > positive and rated >= 3:
        return (
            "Review data or settings",
            f"{rejected} of {rated} reviewed generated images were rejected or scored low. More epochs by themselves are unlikely to be the best first change.",
            "Review dataset variety, captions, and the fixed-prompt gallery. Keep the epoch count similar for the next controlled test, changing only one training setting.",
            epoch_budget,
        )
    if previous and latest.final_loss is not None and previous.final_loss is not None:
        loss_change = latest.final_loss - previous.final_loss
        if abs(loss_change) <= max(0.0001, abs(previous.final_loss) * 0.03):
            return (
                "Likely plateau",
                "The last two recorded losses changed very little. That is a plateau signal, not proof that the model has stopped improving visually.",
                "Run a shorter follow-up (about 25% fewer epochs) with a lower learning rate or improved data; compare it using the same evaluation prompt and seed.",
                max(1, round(epoch_budget * 0.75)),
            )
    if positive >= max(2, rejected * 2):
        return (
            "Promising",
            f"{positive} reviewed generated images look positive versus {rejected} rejected. The model has enough signal for a focused continuation test.",
            "Preserve this run as a baseline. Try a modest continuation of about 25% more epochs, then compare the same prompt-and-seed gallery before committing further.",
            max(epoch_budget + 1, round(epoch_budget * 1.25)),
        )
    return (
        "Gather one more comparison",
        "ADAM has mixed review evidence. A single outcome can be affected by prompt choice, seed, or dataset coverage.",
        "Make another small fixed-prompt generation batch, rate it, and change only one setting in the next run so the result is interpretable.",
        epoch_budget,
    )


def build_model_intelligence(
    runs: list[ExperimentRun], generations: list[GenerationRecord], *, root: Path | None = None,
) -> list[ModelIntelligence]:
    """Group local records into newest-first, model-centered intelligence cards."""
    grouped: dict[tuple[str, str], list[ExperimentRun]] = {}
    for run in runs:
        key = (_normalized(run.model_name), _normalized(run.model_architecture))
        if key[0]:
            grouped.setdefault(key, []).append(run)

    profiles: list[ModelIntelligence] = []
    path_cache: dict[str, Path] = {}
    for (name_key, architecture_key), raw_runs in grouped.items():
        model_runs = tuple(sorted(raw_runs, key=lambda run: run.timestamp, reverse=True))
        model_generations = tuple(
            record for record in generations
            if any(_generation_matches(run, record, path_cache) for run in model_runs)
        )
        rated, positive, rejected = _ratings(model_generations, root)
        state, diagnosis, recommendation, epochs = _diagnose(
            model_runs, model_generations, rated, positive, rejected
        )
        profiles.append(ModelIntelligence(
            key=f"{architecture_key}:{name_key}",
            model_name=model_runs[0].model_name,
            architecture=model_runs[0].model_architecture,
            runs=model_runs,
            generations=model_generations,
            rated_images=rated,
            positive_ratings=positive,
            rejected_images=rejected,
            state=state,
            diagnosis=diagnosis,
            recommendation=recommendation,
            recommended_epochs=epochs,
        ))
    return sorted(profiles, key=lambda profile: profile.runs[0].timestamp, reverse=True)


def recommended_training_request(profile: ModelIntelligence) -> str:
    """Create an approval-aware follow-up request using the latest run as a baseline."""
    run = profile.runs[0]
    options = {
        key: value for key, value in run.settings.items()
        if key not in {"dataset_dir", "model_name", "epochs", "output_dir", "resume_from"}
    }
    return (
        f"From the {run.dataset_name or run.dataset_path} dataset, train a "
        f"{run.model_architecture.upper()} model for {profile.recommended_epochs} epochs. "
        f"Name the model {run.model_name} Follow-up. "
        "[ADAM_TRAINING_OPTIONS:" + json.dumps(options, sort_keys=True) + "] "
        "[ADAM_TRAINER:" + run.model_architecture + "]"
    )