augmenator/planner.py

import re

from augmenator.background_replace import parse_background_query
from augmenator.embedding_planner import select_keywords, warmup as warmup_embeddings
from augmenator.keyword_catalog import AugmentKeyword
from augmenator.spatial import infer_cover_count
from augmenator.spatial_triggers import allows_spatial_keyword

SPATIAL_OPS = {
    "cover_avoid_text",
    "cover_random",
    "cover_and_cutout",
    "cover_and_cutout_avoid_text",
    "add_cutout",
    "replace_background",
}

GENERATIVE_PATTERNS = (
    "put ",
    "place ",
    "insert ",
    "behind",
    "in front of",
    "table",
    "chair",
    "person",
    "object",
    "scene",
)

UNSUPPORTED_REASON = (
    "This instruction needs generative editing (adding or replacing objects/scenes). "
    "v1 supports color transforms, classical edits, procedural cutouts, text-aware covering, "
    "and background replacement — not object insertion."
)


def warmup():
    warmup_embeddings()


def _instruction_mentions_text(instruction: str) -> bool:
    lowered = instruction.lower()
    return any(
        phrase in lowered
        for phrase in ("not text", "avoid text", "except text", "not on text", "without text")
    )


def _infer_fill(instruction: str) -> str:
    lowered = instruction.lower()
    if any(word in lowered for word in ("blur", "soft", "soften")):
        return "blur"
    return "solid"


def _infer_cutout_style(instruction: str, explicit: str | None = None) -> str:
    if explicit in ("transparent", "solid"):
        return explicit
    lowered = instruction.lower()
    if any(w in lowered for w in ("transparent", "clear cutout", "cut a hole", "punch hole")):
        return "transparent"
    if any(w in lowered for w in ("solid", "opaque", "colored cutout", "red cutout", "blue cutout", "green cutout")):
        return "solid"
    return "solid"


def _infer_cutout_color(instruction: str) -> tuple[int, int, int]:
    lowered = instruction.lower()
    if "red" in lowered:
        return (220, 60, 60)
    if "blue" in lowered:
        return (60, 100, 220)
    if "green" in lowered:
        return (60, 180, 80)
    if "yellow" in lowered:
        return (230, 200, 50)
    return (255, 200, 80)


def _is_cutout_op(op: str) -> bool:
    return op in ("add_cutout", "cutout") or op.startswith("cover_and_cutout")


def _cutout_spatial_defaults(instruction: str) -> dict:
    return {
        "cutout_count": None,
        "varied_shapes": True,
    }


def _has_generative_intent(instruction: str) -> bool:
    lowered = instruction.lower()
    return any(pattern in lowered for pattern in GENERATIVE_PATTERNS)


def _build_spatial_entry(keyword: AugmentKeyword, instruction: str) -> dict:
    extra = keyword.spatial_extra or {}
    op = keyword.spatial_op or keyword.id
    avoid_text = extra.get("avoid_text", False) or _instruction_mentions_text(instruction)

    entry = {
        "op": op,
        "fill": _infer_fill(instruction),
        "cutout_style": extra.get("cutout_style") or _infer_cutout_style(instruction),
        "cutout_color": _infer_cutout_color(instruction),
        "avoid_text": avoid_text,
        "varied_shapes": extra.get("varied_shapes", _is_cutout_op(op)),
        "query": None,
    }
    if _is_cutout_op(op):
        entry.update(_cutout_spatial_defaults(instruction))
    if extra.get("avoid_text_cutouts"):
        entry["avoid_text_cutouts"] = True

    if op in ("cover_and_cutout", "cover_and_cutout_avoid_text"):
        entry["cover_count"] = infer_cover_count(instruction)
        entry["varied_shapes"] = True
        if op == "cover_and_cutout_avoid_text":
            entry["avoid_text"] = True
            entry["avoid_text_cutouts"] = True

    if op in ("add_cutout", "cutout"):
        entry["varied_shapes"] = True
        if avoid_text or _instruction_mentions_text(instruction):
            entry["avoid_text"] = True
            entry["avoid_text_cutouts"] = True

    if op == "cover_avoid_text":
        entry["avoid_text"] = True
    if op == "replace_background":
        parsed = parse_background_query(instruction)
        entry["query"] = None if parsed is False else parsed

    if avoid_text and op in ("add_cutout", "cover_random"):
        entry["avoid_text"] = True
        if op == "cover_random":
            entry["op"] = "cover_avoid_text"

    return entry


def _correct_spatial_for_instruction(spatial: list[dict], instruction: str) -> list[dict]:
    lowered = instruction.lower()
    mentions_cover = any(
        phrase in lowered
        for phrase in ("cover", "hide part", "hide a", "hide section", "hide region")
    )
    mentions_cutout = any(
        phrase in lowered
        for phrase in ("cutout", "cut out", "hole", "sticker", "overlay", "punch")
    )
    compound_ops = {"cover_and_cutout", "cover_and_cutout_avoid_text"}
    ops = {s["op"] for s in spatial}

    if not (ops & compound_ops):
        return spatial

    entry = next(s for s in spatial if s["op"] in compound_ops)

    if mentions_cutout and not mentions_cover:
        avoid = entry.get("avoid_text") or _instruction_mentions_text(instruction)
        return [
            {
                "op": "add_cutout",
                "fill": entry.get("fill") or _infer_fill(instruction),
                "cutout_style": entry.get("cutout_style", "solid"),
                "cutout_color": entry.get("cutout_color", (255, 200, 80)),
                "avoid_text": avoid,
                "avoid_text_cutouts": avoid,
                "cutout_count": entry.get("cutout_count"),
                "varied_shapes": True,
                "query": None,
            }
        ]

    if mentions_cover and not mentions_cutout:
        op = "cover_avoid_text" if entry.get("avoid_text") else "cover_random"
        return [
            {
                "op": op,
                "fill": entry.get("fill") or _infer_fill(instruction),
                "cutout_style": entry.get("cutout_style", "solid"),
                "cutout_color": entry.get("cutout_color", (255, 200, 80)),
                "avoid_text": entry.get("avoid_text", False),
                "varied_shapes": False,
                "query": None,
            }
        ]

    return spatial


def _finalize_spatial(spatial: list[dict]) -> list[dict]:
    compound_ops = {"cover_and_cutout", "cover_and_cutout_avoid_text"}
    simple_cover = {"cover_random", "cover_avoid_text", "cover"}
    simple_cutout = {"add_cutout", "cutout"}
    ops = {s["op"] for s in spatial}

    if ops & compound_ops:
        return [
            s for s in spatial
            if s["op"] not in simple_cover and s["op"] not in simple_cutout
        ]

    if (ops & simple_cutout) and not (ops & simple_cover):
        return [s for s in spatial if s["op"] not in compound_ops]

    return spatial


def _consolidate_cover_cutout(spatial: list[dict], instruction: str) -> list[dict]:
    """Merge separate cover + cutout ops into a combined cover_and_cutout op."""
    has_cover = any(s["op"] in ("cover_random", "cover", "cover_avoid_text") for s in spatial)
    has_cutout = any(s["op"] == "add_cutout" for s in spatial)
    if not (has_cover and has_cutout):
        return spatial

    avoid_text = any(
        s.get("avoid_text") or s["op"] in ("cover_avoid_text", "cover_and_cutout_avoid_text")
        for s in spatial
    )
    cutout_style = "solid"
    cutout_color = (255, 200, 80)
    for s in spatial:
        if s["op"] == "add_cutout":
            cutout_style = s.get("cutout_style", cutout_style)
            cutout_color = s.get("cutout_color", cutout_color)

    combined_op = "cover_and_cutout_avoid_text" if avoid_text else "cover_and_cutout"
    combined = {
        "op": combined_op,
        "fill": _infer_fill(instruction),
        "cutout_style": cutout_style,
        "cutout_color": cutout_color,
        "avoid_text": avoid_text,
        "avoid_text_cutouts": avoid_text,
        "varied_shapes": True,
        "cover_count": infer_cover_count(instruction),
        "cutout_count": None,
        "query": None,
    }

    remaining = [
        s
        for s in spatial
        if s["op"] not in ("cover_random", "cover", "cover_avoid_text", "add_cutout", "cutout")
    ]
    if not any(s["op"] == combined_op for s in remaining):
        remaining.append(combined)
    return remaining


def _matches_to_plan(
    matches: list[tuple[AugmentKeyword, float]],
    instruction: str,
) -> tuple[list[str], list[dict], list[tuple[str, float]]]:
    tags: list[str] = []
    spatial: list[dict] = []
    scores: list[tuple[str, float]] = []

    for keyword, score in matches:
        scores.append((keyword.id, score))
        if keyword.kind == "tag":
            if keyword.id not in tags:
                tags.append(keyword.id)
        elif keyword.kind == "spatial":
            entry = _build_spatial_entry(keyword, instruction)
            dup_key = (entry.get("op"), entry.get("cutout_style"), entry.get("query"))
            if not any(
                (e.get("op"), e.get("cutout_style"), e.get("query")) == dup_key for e in spatial
            ):
                spatial.append(entry)

    spatial = _consolidate_cover_cutout(spatial, instruction)
    spatial = _finalize_spatial(spatial)
    spatial = _correct_spatial_for_instruction(spatial, instruction)
    return tags, spatial, scores


def plan_from_instruction(instruction: str) -> dict:
    instruction = instruction.strip()

    if not instruction:
        return {
            "tags": [],
            "spatial": [],
            "scores": [],
            "supported": False,
            "reason": "Please enter an instruction.",
        }

    matches = select_keywords(instruction)
    matches = [
        (keyword, score)
        for keyword, score in matches
        if keyword.kind != "spatial" or allows_spatial_keyword(keyword.id, instruction)
    ]
    tags, spatial, scores = _matches_to_plan(matches, instruction)

    if not tags and not spatial and _has_generative_intent(instruction):
        return {
            "tags": [],
            "spatial": [],
            "scores": scores,
            "supported": False,
            "reason": UNSUPPORTED_REASON,
        }

    if tags or spatial:
        normalized_spatial = []
        for entry in spatial:
            op = entry.get("op", "")
            if op in SPATIAL_OPS:
                normalized_spatial.append(
                    {
                        "op": op,
                        "fill": entry.get("fill") or _infer_fill(instruction),
                        "cutout_style": entry.get("cutout_style", "solid"),
                        "cutout_color": entry.get("cutout_color", (255, 200, 80)),
                        "avoid_text": entry.get("avoid_text", op in ("cover_avoid_text", "cover_and_cutout_avoid_text")),
                        "avoid_text_cutouts": entry.get(
                            "avoid_text_cutouts",
                            entry.get("avoid_text", False) or op == "cover_and_cutout_avoid_text",
                        ),
                        "varied_shapes": entry.get(
                            "varied_shapes",
                            op in ("cover_and_cutout", "cover_and_cutout_avoid_text", "add_cutout", "cutout"),
                        ),
                        "cover_count": entry.get("cover_count"),
                        "cutout_count": entry.get("cutout_count"),
                        "query": entry.get("query"),
                    }
                )

        return {
            "tags": tags,
            "spatial": normalized_spatial,
            "scores": scores,
            "supported": True,
            "reason": None,
        }

    return {
        "tags": [],
        "spatial": [],
        "scores": scores,
        "supported": False,
        "reason": (
            f'Could not map "{instruction}" to a supported transform. '
            "Try describing an effect more clearly, or use one of the example prompts below."
        ),
    }