adam/training_assistant.py

from __future__ import annotations

import json
import re
import shutil
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from adam.models import ExecutionPlan
from adam.orion import apply_orion_review


DEFAULT_PRESETS: dict[str, dict[str, Any]] = {
    "Character LoRA": {
        "trainer": "lora", "epochs": 100, "image_count": 60,
        "description": "A balanced starting point for a character or person.",
    },
    "Style LoRA": {
        "trainer": "lora", "epochs": 80, "image_count": 80,
        "description": "A broader image set for learning a visual style.",
    },
    "DDPM Test Run": {
        "trainer": "ddpm", "epochs": 25, "image_count": 40,
        "description": "A short run to verify the dataset and training setup.",
    },
    "DDPM Full Run": {
        "trainer": "ddpm", "epochs": 100, "image_count": 100,
        "description": "A practical default for a full DDPM experiment.",
    },
    "Flow Test Run": {
        "trainer": "flow", "epochs": 25, "image_count": 40,
        "description": "A short Flow Matching setup check.",
    },
}


def parse_model_batch_names(text: str) -> list[str]:
    """Return unique, user-ordered model subjects from a pasted line list."""
    names: list[str] = []
    seen: set[str] = set()
    for raw in text.splitlines():
        name = re.sub(r"^\s*(?:[-*•]|\d+[.)])\s*", "", raw).strip()
        key = re.sub(r"\s+", " ", name).casefold()
        if name and key not in seen:
            names.append(re.sub(r"\s+", " ", name))
            seen.add(key)
    return names


def build_dataset_collection_request(
    subject: str,
    *,
    image_count: int = 100,
    collection_mode: str = "target",
) -> str:
    """Build the dataset-only first phase used by a saved model batch."""
    subject = subject.strip()
    if not subject:
        raise ValueError("Dataset collection requires a subject.")
    if collection_mode == "all_available":
        return (
            f"Collect a dataset of {subject} with as many available images as Bing "
            "returns (up to 5,000)."
        )
    return f"Collect a dataset of {image_count} images of {subject}."


@dataclass(slots=True)
class PreflightItem:
    level: str
    message: str


@dataclass(slots=True)
class DatasetMatch:
    status: str
    dataset_name: str = ""
    score: float = 0.0


_DATASET_NAME_NOISE = {
    "dataset", "datasets", "image", "images", "picture", "pictures",
    "photo", "photos", "collection", "collected",
}


def _dataset_name_tokens(value: object) -> set[str]:
    words = re.findall(r"[a-z0-9]+", str(value).casefold())
    return {
        word[:-1] if word.endswith("s") and len(word) > 3 else word
        for word in words if word not in _DATASET_NAME_NOISE
    }


def suggest_existing_dataset(state: dict[str, object], datasets: list[Any]) -> DatasetMatch:
    """Safely match one batch model to a registered dataset by its human name."""
    queries = [
        _dataset_name_tokens(state.get("model_name", "")),
        _dataset_name_tokens(state.get("subject", "")),
    ]
    queries = [query for query in queries if query]
    if not queries:
        return DatasetMatch("unmatched")
    scored: list[tuple[float, Any]] = []
    for asset in datasets:
        name = str(getattr(asset, "name", ""))
        path = Path(str(getattr(asset, "path", "")))
        tokens = _dataset_name_tokens(name)
        if not name or not tokens or not path.is_dir():
            continue
        score = 0.0
        for query in queries:
            overlap = len(query & tokens) / len(query)
            extra_penalty = min(0.20, len(tokens - query) * 0.08)
            score = max(score, overlap - extra_penalty)
        if score >= 0.80:
            scored.append((score, asset))
    if not scored:
        return DatasetMatch("unmatched")
    scored.sort(key=lambda item: (-item[0], len(str(getattr(item[1], "name", "")))))
    best_score, best = scored[0]
    if len(scored) > 1 and best_score - scored[1][0] < 0.10:
        return DatasetMatch("ambiguous", score=best_score)
    return DatasetMatch("matched", str(getattr(best, "name", "")), best_score)


def combine_training_plans(plans: list[ExecutionPlan]) -> ExecutionPlan:
    """Combine independently validated model plans into one sequential job."""
    usable = [plan for plan in plans if plan.steps]
    if not usable:
        raise ValueError("A training batch needs at least one actionable model plan.")
    if len(usable) == 1:
        return usable[0]
    summaries = [
        f"{index}. {plan.project_name}: {plan.summary.splitlines()[0]}"
        for index, plan in enumerate(usable, 1)
    ]
    reasons = [plan.confirmation_reason for plan in usable if plan.confirmation_reason]
    has_training = any(
        step.tool_id.endswith("_trainer")
        for plan in usable for step in plan.steps
    )
    batch_kind = "training" if has_training else "dataset collection"
    return ExecutionPlan(
        request="\n\n".join(plan.request for plan in usable),
        summary=(
            f"Sequential {batch_kind} batch with {len(usable)} items. ADAM will finish "
            "each item before starting the next; a failed step stops the batch.\n\n"
            + "\n".join(summaries)
        ),
        steps=[step for plan in usable for step in plan.steps],
        requires_confirmation=any(plan.requires_confirmation for plan in usable),
        confirmation_reason="; ".join(dict.fromkeys(reasons)) or (
            "This batch contains multiple model workflows. Review every model and its "
            "output path before starting."
        ),
        project_name=(
            f"Training batch ({len(usable)} models)" if has_training
            else f"Dataset collection batch ({len(usable)} datasets)"
        ),
    )


def estimate_plan(plan: Any) -> list[PreflightItem]:
    """Add deliberately conservative, clearly labelled planning estimates."""
    estimates: list[PreflightItem] = []
    for step in plan.steps:
        if not step.tool_id.endswith("_trainer"):
            continue
        trainer = step.tool_id.removesuffix("_trainer")
        epochs = max(1, int(step.arguments.get("epochs", 1) or 1))
        dataset = Path(str(step.arguments.get("dataset_dir", ""))).expanduser()
        image_count = 0
        if dataset.is_dir():
            try:
                image_count = sum(
                    1
                    for path in dataset.rglob("*")
                    if path.is_file()
                    and path.suffix.casefold()
                    in {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
                )
            except OSError:
                image_count = 0
        image_count = image_count or 60
        workload = epochs * image_count
        seconds_per_image_epoch = {
            "lora": 0.12,
            "ddpm": 0.07,
            "flow": 0.10,
        }.get(trainer, 0.10)
        center_minutes = max(1, int(workload * seconds_per_image_epoch / 60))
        low = max(1, center_minutes // 2)
        high = max(low + 1, center_minutes * 3)
        checkpoint_gb = {
            "lora": 0.25,
            "ddpm": 1.0,
            "flow": 1.0,
        }.get(trainer, 0.75)
        checkpoint_count = max(1, min(20, epochs // 25 + 1))
        disk_gb = checkpoint_gb * checkpoint_count
        typical_vram = {"lora": 8, "ddpm": 6, "flow": 8}.get(trainer, 8)
        estimates.extend(
            [
                PreflightItem(
                    "estimate",
                    f"Estimated workload: {workload:,} image-epochs "
                    f"({epochs:,} epochs × about {image_count:,} images)",
                ),
                PreflightItem(
                    "estimate",
                    f"Rough duration: {low}–{high} minutes; model size, resolution, "
                    "batch size, and GPU can change this substantially",
                ),
                PreflightItem(
                    "estimate",
                    f"Suggested capacity: about {typical_vram} GB VRAM and "
                    f"{disk_gb:.1f} GB free for checkpoints",
                ),
            ]
        )
    return estimates


def presets_from_config(config: Any) -> dict[str, dict[str, Any]]:
    presets = {name: dict(values) for name, values in DEFAULT_PRESETS.items()}
    stored = config.get("training_presets", {})
    if isinstance(stored, dict):
        for name, values in stored.items():
            if isinstance(name, str) and isinstance(values, dict):
                presets[name] = dict(values)
    return presets


def build_training_request(
    *,
    trainer: str,
    subject: str,
    dataset_name: str,
    create_dataset: bool,
    epochs: int,
    image_count: int,
    model_name: str,
    collection_mode: str = "target",
    training_options: dict[str, Any] | None = None,
) -> str:
    subject = subject.strip()
    dataset_name = dataset_name.strip()
    model_name = model_name.strip() or subject or dataset_name
    trainer_label = {"lora": "LoRA", "ddpm": "DDPM", "flow": "Flow Matching"}[trainer]
    if create_dataset:
        collection_phrase = (
            "as many available images as Bing returns (up to 5,000)"
            if collection_mode == "all_available"
            else f"up to {image_count} images"
        )
        if trainer == "lora":
            request = (
                f"Create and train a LoRA of {subject} for {epochs} epochs "
                f"using {collection_phrase}. Name the model {model_name}."
            )
        elif trainer == "ddpm":
            request = (
                f"Grab a dataset of {subject} off the internet with {collection_phrase}, "
                f"name the model {model_name}, train it on a DDPM for {epochs} epochs, "
                "and save it to the DDPM output."
            )
        else:
            request = (
            f"Collect a dataset of {collection_phrase} of {subject}. Then train the "
            f"{subject} dataset with Flow Matching for {epochs} epochs and name the model {model_name}."
            )
    else:
        request = (
            f"From the {dataset_name} dataset, train a {trainer_label} model for {epochs} epochs. "
            f"Name the model {model_name}."
        )
    if training_options:
        request += " [ADAM_TRAINING_OPTIONS:" + json.dumps(training_options, sort_keys=True) + "]"
    return request


def build_fine_tune_request(
    *,
    model_name: str,
    trainer: str,
    epochs: int,
    dataset_mode: str = "original",
    dataset_name: str = "",
    new_subject: str = "",
    image_count: int = 60,
    training_options: dict[str, Any] | None = None,
) -> str:
    """Build the explicit continuation request used by the Fine-Tune assistant."""
    labels = {"lora": "LoRA", "ddpm": "DDPM", "flow": "Flow Matching"}
    if trainer not in labels:
        raise ValueError("Fine-tuning requires a supported trainer.")
    if not model_name.strip():
        raise ValueError("Fine-tuning requires a model name.")
    if epochs < 1:
        raise ValueError("Fine-tuning requires at least one additional epoch.")
    if dataset_mode not in {"original", "existing", "new"}:
        raise ValueError("Fine-tuning requires a valid dataset choice.")
    payload = {
        "model_name": model_name.strip(),
        "trainer": trainer,
        "epochs": epochs,
        "dataset_mode": dataset_mode,
        "dataset_name": dataset_name.strip(),
        "new_subject": new_subject.strip(),
        "image_count": max(10, min(int(image_count), 5000)),
        "training_options": dict(training_options or {}),
    }
    return (
        f"Fine-tune {model_name.strip()} for {epochs} epochs with {labels[trainer]}. "
        "[ADAM_FINE_TUNE:" + json.dumps(payload, sort_keys=True) + "]"
    )


def inspect_plan(plan: Any, config: Any) -> list[PreflightItem]:
    items: list[PreflightItem] = []
    folders = config.get("tool_folders", {})
    folders = folders if isinstance(folders, dict) else {}
    checked_tools: set[str] = set()
    checked_paths: set[str] = set()

    for step in plan.steps:
        if step.tool_id.endswith("_trainer") or step.tool_id == "dataset_collector":
            if step.tool_id not in checked_tools:
                configured = Path(str(folders.get(step.tool_id, ""))).expanduser()
                if configured.is_dir():
                    items.append(PreflightItem("ready", f"{step.title}: connected"))
                else:
                    items.append(PreflightItem("warning", f"{step.title}: program folder is not connected"))
                checked_tools.add(step.tool_id)
        dataset = str(step.arguments.get("dataset_dir", ""))
        if dataset and dataset not in checked_paths:
            if Path(dataset).is_dir():
                image_count = sum(
                    1 for path in Path(dataset).iterdir()
                    if path.suffix.casefold() in {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
                )
                detail = f"{image_count} images found" if image_count else "folder found; no top-level images detected"
                items.append(PreflightItem("ready" if image_count else "warning", f"Dataset: {detail}"))
            elif not any(
                prior.tool_id == "dataset_collector"
                and prior.arguments.get("output_dir") == dataset
                for prior in plan.steps
            ):
                items.append(PreflightItem("warning", "Dataset folder does not exist yet"))
            checked_paths.add(dataset)
        base_model = str(step.arguments.get("base_model", ""))
        if step.tool_id == "lora_trainer":
            items.append(
                PreflightItem(
                    "ready" if base_model and Path(base_model).is_file() else "warning",
                    "LoRA base model is available" if base_model and Path(base_model).is_file()
                    else "LoRA base model still needs to be selected",
                )
            )
        output = str(step.arguments.get("output_dir", ""))
        if output:
            probe = Path(output)
            while not probe.exists() and probe.parent != probe:
                probe = probe.parent
            try:
                free_gb = shutil.disk_usage(probe).free / (1024 ** 3)
                items.append(
                    PreflightItem(
                        "ready" if free_gb >= 10 else "warning",
                        f"Output drive has {free_gb:.1f} GB free",
                    )
                )
            except OSError:
                items.append(PreflightItem("warning", "Output drive space could not be checked"))
    return items


def append_preflight_summary(plan: Any, config: Any) -> None:
    if not plan.steps:
        return
    if "Pre-flight:" not in plan.summary:
        items = inspect_plan(plan, config) + estimate_plan(plan)
        if items:
            lines = [
                (
                    "Ready"
                    if item.level == "ready"
                    else "Estimate"
                    if item.level == "estimate"
                    else "Check"
                )
                + f": {item.message}"
                for item in items
            ]
            plan.summary += "\n\nPre-flight:\n" + "\n".join(f"• {line}" for line in lines)
    if not getattr(plan, "orion_review", None) and "ORION —" not in plan.summary:
        apply_orion_review(plan)


def completion_recommendation(plan: Any) -> str:
    tools = {step.tool_id for step in plan.steps}
    if "lora_trainer" in tools or "ddpm_trainer" in tools or "flow_trainer" in tools:
        return (
            "Recommended next step: generate a few preview images and compare them with "
            "the training dataset. If the subject is weak, improve the dataset before adding epochs."
        )
    if "dataset_collector" in tools:
        return (
            "Recommended next step: review the images and captions, remove weak or duplicate "
            "examples, then open the Model Creation Assistant to start a short test run."
        )
    return ""