adam/image_preferences.py

from __future__ import annotations

from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
import hashlib
import json
from pathlib import Path
from typing import Any, Callable, Sequence

from adam.eve import EveVisionModel, classify_eve_embeddings


RATINGS = {"favorite", "keep", "unsure", "reject"}
POSITIVE_RATINGS = {"favorite", "keep"}
NEGATIVE_RATINGS = {"reject"}


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def preference_profile_id(provider_id: str, model_path: str) -> str:
    key = f"{provider_id}\n{str(Path(model_path).expanduser().resolve())}"
    return hashlib.sha1(key.encode("utf-8")).hexdigest()[:16]


def image_cache_id(path: str | Path) -> str:
    resolved = str(Path(path).expanduser().resolve())
    return hashlib.sha1(resolved.encode("utf-8")).hexdigest()


@dataclass(slots=True)
class GenerationRating:
    image_path: str
    rating: str
    provider_id: str
    model_name: str
    model_path: str
    seed: int = 0
    sampler: str = ""
    steps: int = 0
    resolution: str = ""
    generation_settings: dict[str, Any] = field(default_factory=dict)
    generation_created_at: str = ""
    rated_at: str = field(default_factory=_now)
    embedding: list[float] | None = None

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "GenerationRating":
        rating = str(payload.get("rating", "unsure")).casefold()
        return cls(
            image_path=str(Path(str(payload.get("image_path", ""))).expanduser().resolve()),
            rating=rating if rating in RATINGS else "unsure",
            provider_id=str(payload.get("provider_id", "")),
            model_name=str(payload.get("model_name", "")),
            model_path=str(payload.get("model_path", "")),
            seed=int(payload.get("seed", 0) or 0),
            sampler=str(payload.get("sampler", "")),
            steps=int(payload.get("steps", 0) or 0),
            resolution=str(payload.get("resolution", "")),
            generation_settings=dict(payload.get("generation_settings") or {}),
            generation_created_at=str(payload.get("generation_created_at", "")),
            rated_at=str(payload.get("rated_at") or _now()),
            embedding=[float(value) for value in payload["embedding"]]
            if isinstance(payload.get("embedding"), list)
            else None,
        )


@dataclass(frozen=True, slots=True)
class PreferenceScore:
    image_path: str
    score: float | None
    confidence: float
    category: str
    reason: str = ""


class PreferenceProfile:
    def __init__(self, root: Path, provider_id: str, model_name: str, model_path: str) -> None:
        self.root = root.resolve()
        self.provider_id = provider_id
        self.model_name = model_name
        self.model_path = str(Path(model_path).expanduser().resolve()) if model_path else ""
        self.id = preference_profile_id(provider_id, self.model_path)
        self.path = self.root / "data" / "generation_preferences" / f"{self.id}.json"
        self.keep_threshold = 0.70
        self.reject_threshold = 0.35
        self.ratings: dict[str, GenerationRating] = {}
        self.load()

    def load(self) -> None:
        try:
            payload = json.loads(self.path.read_text(encoding="utf-8"))
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            return
        self.model_name = str(payload.get("model_name") or self.model_name)
        self.provider_id = str(payload.get("provider_id") or self.provider_id)
        self.model_path = str(payload.get("model_path") or self.model_path)
        thresholds = payload.get("thresholds", {})
        if isinstance(thresholds, dict):
            self.keep_threshold = float(thresholds.get("keep", self.keep_threshold))
            self.reject_threshold = float(thresholds.get("reject", self.reject_threshold))
        ratings = payload.get("ratings", [])
        if isinstance(ratings, list):
            for item in ratings:
                if isinstance(item, dict):
                    rating = GenerationRating.from_dict(item)
                    self.ratings[rating.image_path] = rating

    def save(self) -> None:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        temporary = self.path.with_suffix(".tmp")
        temporary.write_text(
            json.dumps(
                {
                    "version": 1,
                    "profile_id": self.id,
                    "provider_id": self.provider_id,
                    "model_name": self.model_name,
                    "model_path": self.model_path,
                    "thresholds": {
                        "keep": self.keep_threshold,
                        "reject": self.reject_threshold,
                    },
                    "ratings": [asdict(item) for item in self.ratings.values()],
                    "updated_at": _now(),
                },
                indent=2,
            ),
            encoding="utf-8",
        )
        temporary.replace(self.path)

    def set_rating(
        self,
        image_path: str | Path,
        rating: str,
        *,
        seed: int = 0,
        sampler: str = "",
        steps: int = 0,
        resolution: str = "",
        generation_settings: dict[str, Any] | None = None,
        generation_created_at: str = "",
        embedding: Sequence[float] | None = None,
    ) -> GenerationRating:
        clean = rating.casefold().strip()
        if clean not in RATINGS:
            raise ValueError("Generation rating must be Favorite, Keep, Unsure, or Reject.")
        resolved = str(Path(image_path).expanduser().resolve())
        existing = self.ratings.get(resolved)
        record = GenerationRating(
            image_path=resolved,
            rating=clean,
            provider_id=self.provider_id,
            model_name=self.model_name,
            model_path=self.model_path,
            seed=int(seed),
            sampler=sampler,
            steps=int(steps),
            resolution=resolution,
            generation_settings=dict(generation_settings or {}),
            generation_created_at=generation_created_at,
            rated_at=_now(),
            embedding=[float(value) for value in embedding] if embedding is not None else (
                existing.embedding if existing else None
            ),
        )
        self.ratings[resolved] = record
        self.save()
        return record

    def rating_for(self, image_path: str | Path) -> GenerationRating | None:
        return self.ratings.get(str(Path(image_path).expanduser().resolve()))

    def examples(self) -> tuple[list[GenerationRating], list[GenerationRating]]:
        positive = [
            item for item in self.ratings.values()
            if item.rating in POSITIVE_RATINGS and Path(item.image_path).is_file()
        ]
        negative = [
            item for item in self.ratings.values()
            if item.rating in NEGATIVE_RATINGS and Path(item.image_path).is_file()
        ]
        return positive, negative

    def has_signal(self) -> bool:
        positive, _negative = self.examples()
        return bool(positive)


class ImageEmbeddingCache:
    def __init__(self, root: Path, model_id: str) -> None:
        self.root = root.resolve()
        self.model_id = model_id
        self.folder = self.root / "data" / "image_embeddings" / hashlib.sha1(model_id.encode("utf-8")).hexdigest()[:12]

    def get(self, path: str | Path) -> list[float] | None:
        cache_path = self.folder / f"{image_cache_id(path)}.json"
        try:
            payload = json.loads(cache_path.read_text(encoding="utf-8"))
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            return None
        source = Path(path).expanduser().resolve()
        try:
            stat = source.stat()
        except OSError:
            return None
        if payload.get("path") != str(source) or payload.get("mtime") != stat.st_mtime:
            return None
        vector = payload.get("embedding")
        return [float(value) for value in vector] if isinstance(vector, list) else None

    def set(self, path: str | Path, embedding: Sequence[float]) -> None:
        source = Path(path).expanduser().resolve()
        try:
            stat = source.stat()
        except OSError:
            return
        self.folder.mkdir(parents=True, exist_ok=True)
        cache_path = self.folder / f"{image_cache_id(source)}.json"
        temporary = cache_path.with_suffix(".tmp")
        temporary.write_text(
            json.dumps(
                {
                    "path": str(source),
                    "mtime": stat.st_mtime,
                    "model_id": self.model_id,
                    "embedding": [float(value) for value in embedding],
                }
            ),
            encoding="utf-8",
        )
        temporary.replace(cache_path)


class GenerationPreferenceEvaluator:
    """Shared EVE-backed scorer for generated images."""

    def __init__(
        self,
        root: Path,
        vision: EveVisionModel | None = None,
        *,
        embedder: Callable[[Sequence[str | Path]], list[list[float]]] | None = None,
    ) -> None:
        self.root = root.resolve()
        self.vision = vision or EveVisionModel(prefer_gpu=False)
        self.embedder = embedder
        self.cache = ImageEmbeddingCache(self.root, self.vision.model_id)

    def _embedding(self, path: str | Path) -> list[float]:
        cached = self.cache.get(path)
        if cached is not None:
            return cached
        vectors = self.embedder([path]) if self.embedder else self.vision.embed([path])
        vector = [float(value) for value in vectors[0]]
        self.cache.set(path, vector)
        return vector

    def score(
        self,
        profile: PreferenceProfile,
        image_paths: Sequence[str | Path],
        *,
        keep_threshold: float | None = None,
        reject_threshold: float | None = None,
    ) -> list[PreferenceScore]:
        positive, negative = profile.examples()
        if not positive:
            return [
                PreferenceScore(str(Path(path).expanduser().resolve()), None, 0.0, "Needs Review", "No preference examples yet")
                for path in image_paths
            ]
        positive_vectors = [item.embedding or self._embedding(item.image_path) for item in positive]
        negative_vectors = [item.embedding or self._embedding(item.image_path) for item in negative]
        image_vectors = [self._embedding(path) for path in image_paths]
        keep = max(0.001, min(1.0, float(keep_threshold if keep_threshold is not None else profile.keep_threshold)))
        reject = max(0.0, min(float(reject_threshold if reject_threshold is not None else profile.reject_threshold), keep - 0.001))
        results = classify_eve_embeddings(
            image_paths,
            image_vectors,
            positive_vectors,
            negative_vectors,
            keep_threshold=keep,
            reject_threshold=reject,
        )
        categories = {"keep": "Strong Keep", "reject": "Likely Reject", "unreviewed": "Needs Review"}
        return [
            PreferenceScore(result.path, result.match_score, result.decision_confidence, categories[result.suggestion])
            for result in results
        ]


def score_generated_images(
    root: Path,
    *,
    provider_id: str,
    model_name: str,
    model_path: str,
    image_paths: Sequence[str | Path],
    keep_threshold: float | None = None,
    reject_threshold: float | None = None,
) -> list[PreferenceScore]:
    profile = PreferenceProfile(root, provider_id, model_name, model_path)
    evaluator = GenerationPreferenceEvaluator(root)
    return evaluator.score(
        profile,
        image_paths,
        keep_threshold=keep_threshold,
        reject_threshold=reject_threshold,
    )