agents/critic.py

"""
Critic — evaluates generated images for style fidelity + technical quality.

Calls the AMD Qwen API with a vision-capable prompt if VLM is available;
otherwise returns a heuristic score based on simple image statistics
(entropy, saturation, edge density).
"""
from __future__ import annotations

import logging
import math
from dataclasses import dataclass
from typing import Dict

import numpy as np
from PIL import Image

from .base import AgentClient, AgentResponse
from config.styles import StyleSpec

log = logging.getLogger(__name__)

SYSTEM_PROMPT = """You are a critic evaluating whether a generated image faithfully represents a chosen Indian heritage art style.

Score on three axes (1-10):
- style_fidelity: how recognizable the heritage style is
- composition: balance, framing, focal point
- technical_quality: sharpness, color harmony, artifact absence

Output STRICT JSON:
{"style_fidelity": <int>, "composition": <int>, "technical_quality": <int>,
 "overall": <float>,
 "should_regenerate": <bool>,
 "feedback": "<one sentence>"}"""


@dataclass
class CritiqueResult:
    style_fidelity: int
    composition: int
    technical_quality: int
    overall: float
    should_regenerate: bool
    feedback: str
    source: str  # "amd_agent_vlm" | "amd_agent" | "heuristic"


class Critic:
    def __init__(self, client: AgentClient | None = None) -> None:
        self.client = client or AgentClient(temperature=0.4, max_tokens=400)

    def evaluate(self, image: Image.Image, style: StyleSpec, prompt: str) -> CritiqueResult:
        """Evaluate a generated image. Falls back to image-statistics heuristic."""
        # Heuristic path (always computed as fallback / baseline)
        heur = self._heuristic_eval(image, style)

        # If agent API is available, attempt a text-only critic call
        # (full VLM call would require an OpenAI-compatible image endpoint,
        # which the AMD API may not expose in this hackathon window.)
        if self.client.enabled:
            try:
                resp: AgentResponse = self.client.chat(
                    system_prompt=SYSTEM_PROMPT,
                    user_prompt=(
                        f"Style: {style.id} ({style.display_name}).\n"
                        f"Original prompt: {prompt!r}.\n"
                        f"Heuristic stats: overall={heur.overall:.2f}.\n"
                        f"Style keywords: {', '.join(style.prompt_tags[:4])}.\n"
                        "Return JSON critique."
                    ),
                )
                if resp.ok:
                    return self._parse_agent(resp.content, heur)
            except Exception as exc:
                log.warning("Critic agent call failed: %s", exc)

        return heur

    @staticmethod
    def _parse_agent(raw: str, fallback: CritiqueResult) -> CritiqueResult:
        import json, re
        try:
            data = json.loads(raw)
        except Exception:
            m = re.search(r"\{[^{}]*\}", raw, re.DOTALL)
            if not m:
                return fallback
            try:
                data = json.loads(m.group(0))
            except Exception:
                return fallback
        return CritiqueResult(
            style_fidelity=int(data.get("style_fidelity", fallback.style_fidelity)),
            composition=int(data.get("composition", fallback.composition)),
            technical_quality=int(data.get("technical_quality", fallback.technical_quality)),
            overall=float(data.get("overall", fallback.overall)),
            should_regenerate=bool(data.get("should_regenerate", False)),
            feedback=str(data.get("feedback", fallback.feedback)),
            source="amd_agent",
        )

    @staticmethod
    def _heuristic_eval(image: Image.Image, style: StyleSpec) -> CritiqueResult:
        """Image-statistics-based critique. Robust but shallow."""
        arr = np.array(image.convert("RGB")).astype(np.float32) / 255.0

        # Saturation
        maxc = arr.max(axis=2)
        minc = arr.min(axis=2)
        sat = (maxc - minc).mean()

        # Edge density (proxy for detail / linework)
        gray = arr.mean(axis=2)
        gx = np.abs(np.diff(gray, axis=1)).mean()
        gy = np.abs(np.diff(gray, axis=0)).mean()
        edge_density = (gx + gy) / 2

        # Entropy of intensity histogram
        hist, _ = np.histogram(gray, bins=32, range=(0, 1))
        hist = hist / max(hist.sum(), 1)
        entropy = -np.sum(hist * np.log(hist + 1e-12))

        # Score mapping (calibrated empirically)
        tech = min(10, int(entropy * 6 + edge_density * 30))
        comp = min(10, int(5 + (entropy - 3.0) * 1.5)) if entropy > 3 else 6
        # Style fidelity: high saturation + high edge density favors traditional styles
        style_score = min(10, int(sat * 12 + edge_density * 25))
        overall = round((tech + comp + style_score) / 3, 2)

        return CritiqueResult(
            style_fidelity=max(1, min(10, style_score)),
            composition=max(1, min(10, comp)),
            technical_quality=max(1, min(10, tech)),
            overall=overall,
            should_regenerate=overall < 5.0,
            feedback=f"Heuristic: sat={sat:.2f}, edge={edge_density:.3f}, entropy={entropy:.2f}",
            source="heuristic",
        )