src/semantic_coherence.py

import random
from src.prompt_parser import ParsedPrompt
from src.prompt_rewriter import get_tag_categories
from src.synonym_data import _find_synonym_group
from src.tag_searcher import get_cooccurrence_tags

# Categories considered "core" (should be protected from replacement)
_SPLIT_CORE_CATEGORIES = {
    "identity", "feature", "body", "hair", "eye", "face"
}


def split_core_decorative(tags: list[str]) -> tuple[list[str], list[str]]:
    """Split user tags into core (subject/subject descriptors) and decorative (add-on elements).
    Returns: (core_tags, decorative_tags)"""
    core: list[str] = []
    decorative: list[str] = []
    for tag in tags:
        tl = tag.lower().strip()
        cats = get_tag_categories(tl)
        if any(c in _SPLIT_CORE_CATEGORIES for c in cats):
            core.append(tag)
        else:
            decorative.append(tag)
    return core, decorative



# Intent definitions: which categories vote for which intent
_INTENT_CATEGORIES = {
    "portrait": {"expression", "hair", "body", "framing", "lighting", "makeup", "accessory"},
    "action": {"pose", "weapon", "effects", "special_fx", "vehicle", "composition"},
    "environment": {"background", "architecture", "weather", "season", "atmosphere", "colors", "lighting"},
    "horror": {"demon", "atmosphere", "effects", "special_fx", "lighting", "colors", "expression"},
    "romantic": {"expression", "bloom", "lighting", "colors", "atmosphere", "accessory", "clothing"},
    "fantasy": {"demon", "angelic", "magic", "weapon", "architecture", "clothing", "effects", "colors"},
}

# Category -> intent alignment: +3 for matching, -2 for opposing
_INTENT_ALIGNMENT = {
    "portrait": {"expression", "hair", "body", "framing"},
    "action": {"pose", "weapon", "effects", "special_fx", "vehicle"},
    "environment": {"background", "architecture", "atmosphere", "weather"},
    "horror": {"demon", "atmosphere", "effects"},
    "romantic": {"expression", "colors", "atmosphere"},
    "fantasy": {"demon", "angelic", "magic", "weapon"},
}

_INTENT_OPPOSITION = {
    "portrait": {"vehicle", "weapon", "animal", "furry", "architecture"},
    "action": {"expression", "hair", "makeup", "food", "season"},
    "environment": {"demon", "angelic", "magic", "body", "hair"},
    "horror": {"season", "food", "hair", "makeup", "romantic"},
    "romantic": {"demon", "weapon", "horror", "gore", "blood"},
    "fantasy": {"vehicle", "modern", "technology", "realistic", "photorealistic"},
}


def detect_intent(parsed: ParsedPrompt) -> str:
    """Определяет тип сцены по тегам пользователя: portrait/action/environment/horror/romantic/fantasy/general."""
    scores: dict[str, float] = {k: 0.0 for k in _INTENT_CATEGORIES}
    all_tags = [(parsed.subject or "").lower().strip()] if parsed.subject else []
    all_tags.extend(t.lower().strip() for t in (parsed.general_tags or []))
    all_tags.extend(t.lower().strip() for t in (parsed.character or "").split(",") if parsed.character)
    all_tags = [t for t in all_tags if t]

    for tag in all_tags:
        cats = get_tag_categories(tag)
        for cat in cats:
            for intent, aligned in _INTENT_ALIGNMENT.items():
                if cat in aligned:
                    scores[intent] += 2.0
            for intent, opposed in _INTENT_OPPOSITION.items():
                if cat in opposed:
                    scores[intent] -= 1.0

    best = max(scores, key=scores.get)
    return best if scores[best] >= 1.0 else "general"


_HARMFUL_PAIRS: set[tuple[str, str]] = {
    # time-of-day & sky
    ("day", "night"), ("noon", "midnight"), ("sunrise", "sunset"),
    ("golden hour", "night"),
    # weather & sky
    ("sunny", "rainy"), ("sunny", "snowy"), ("rainy", "snowy"),
    ("clear sky", "cloudy sky"), ("clear sky", "starry sky"),
    ("cloudy sky", "starry sky"), ("rainbow", "monochrome"),
    # palette
    ("warm colors", "cool colors"), ("vibrant colors", "muted colors"),
    ("vibrant colors", "monochrome"), ("pastel colors", "dark colors"),
    ("pastel colors", "neon palette"), ("bloom", "monochrome"),
    ("colorful", "monochrome"),
    # lighting
    ("moonlight", "sunlight"), ("moonlight", "daylight"),
    ("soft lighting", "harsh lighting"), ("soft lighting", "hard lighting"),
    ("backlighting", "front lighting"),
    # setting
    ("indoors", "outdoors"), ("underwater", "space"), ("cityscape", "countryside"),
    ("bedroom", "outdoors"), ("kitchen", "space"),
    # framing & camera
    ("portrait", "wide shot"), ("close-up", "wide shot"),
    ("cowboy shot", "extreme close-up"), ("full body", "extreme close-up"),
    ("from above", "from below"), ("from side", "from behind"),
    # gaze
    ("looking at viewer", "looking away"), ("looking at viewer", "looking down"),
    ("looking at viewer", "eyes closed"),
    # expression clashes
    ("innocent", "seductive smile"), ("crying", "laughing"),
    ("angry", "smile"), ("serious", "laughing"),
    # style clashes
    ("chibi", "photorealistic"), ("chibi", "realistic"),
    ("flat color", "photorealistic"), ("sketch", "photorealistic"),
    ("lineart", "full color"), ("monochrome", "full color"),
    ("watercolor (medium)", "photorealistic"),
    # season clashes
    ("winter", "summer"), ("spring", "autumn"),
    # posture
    ("standing", "sitting"), ("standing", "lying"), ("sitting", "lying"),
    ("standing", "kneeling"),
    # subject-type clashes
    ("solo", "multiple girls"), ("solo", "multiple boys"),
    ("solo focus", "group focus"),
}

# Pre-symmetrized pairs so callers never need to check both orderings at runtime.
_HARMFUL_PAIRS_SYM: set[tuple[str, str]] = _HARMFUL_PAIRS | {
    (b, a) for a, b in _HARMFUL_PAIRS
}


# Central place for all selection-scoring weights (tweak here, not inline).
SCORING = {
    "theme_aligned": 3.0,     # candidate category aligns with detected intent
    "theme_opposed": -2.0,    # candidate category fights the detected intent
    "complement_cap": 6.0,    # max complement penalty
    "harmony_conflict": -5.0, # a selected tag is a hard conflict
    "cooc_cap": 5.0,          # max co-occurrence bonus
    "context_score_cap": 5.0, # cap for subject/character/series string match
    "synonym_sub_bonus": 1.0, # same-synonym-group candidate bonus
    "noise_max": 0.3,         # upper bound of uniform exploration noise
}


def semantic_pick_tags(
    candidates: list[str],
    count: int,
    parsed: ParsedPrompt,
    rng: random.Random,
    used_globals: set[str],
    selected_tags: list[str] | None = None,
    intent: str | None = None,
) -> list[str]:
    """Интеллектуальный выбор тегов: theme_score + complement + harmony + cooc."""
    if intent is None:
        intent = detect_intent(parsed)

    # Pre-compute categories for selected tags (cache)
    sel_cats_cache: dict[str, set[str]] = {}
    sel_syn_cache: dict[str, set | None] = {}
    selected = selected_tags or []
    for sel in selected:
        sl = sel.lower().strip()
        if sl not in sel_cats_cache:
            sel_cats_cache[sl] = set(get_tag_categories(sel))
        if sl not in sel_syn_cache:
            sel_syn_cache[sl] = _find_synonym_group(sel)

    scored = []
    for tag in candidates:
        tl = tag.lower().strip()
        if tl in used_globals:
            continue

        # Theme score (cached categories)
        theme = 0.0
        tag_cats = sel_cats_cache.get(tl)
        if tag_cats is None:
            tag_cats = set(get_tag_categories(tag))
            sel_cats_cache[tl] = tag_cats
        for cat in tag_cats:
            if cat in _INTENT_ALIGNMENT.get(intent, set()):
                theme = SCORING["theme_aligned"]
                break
        if theme == 0.0:
            for cat in tag_cats:
                if cat in _INTENT_OPPOSITION.get(intent, set()):
                    theme = SCORING["theme_opposed"]
                    break

        # Complement penalty
        complement = 0.0
        tag_syn = _find_synonym_group(tag)
        for s in selected:
            sl = s.lower().strip()
            overlap = tag_cats & sel_cats_cache.get(sl, set())
            if overlap:
                complement += 1.0
            if tag_syn and sl in tag_syn:
                complement += 3.0
        complement = min(complement, SCORING["complement_cap"])

        # Cross-category harmony
        harmony = 0.0
        for s in selected:
            sl = s.lower().strip()
            tag_group = tag_syn or {tl}
            sel_group = sel_syn_cache.get(sl) or {sl}
            for t in tag_group:
                for s2 in sel_group:
                    if (t, s2) in _HARMFUL_PAIRS_SYM:
                        harmony = SCORING["harmony_conflict"]
                        break
                if harmony < 0.0:
                    break
            if harmony < 0.0:
                break

        cooc = _cooccurrence_bonus_fast(tag, selected)
        noise = rng.uniform(0, SCORING["noise_max"])
        final = theme + cooc - complement + harmony + noise
        scored.append((tag, final))

    scored.sort(key=lambda x: (-x[1], x[0]))
    picked = []
    for tag, _ in scored[:count]:
        picked.append(tag)
        selected.append(tag)
    rng.shuffle(picked)
    return picked


# Theme groups for adaptive budget allocation
_THEME_GROUPS: dict[str, set[str]] = {
    "appearance": {"expression", "hair", "body", "eye", "face", "makeup"},
    "environment": {"background", "architecture", "weather", "season"},
    "lighting_atmosphere": {"lighting", "colors", "atmosphere", "bloom"},
    "effects": {"effects", "special_fx", "composition", "framing"},
    "props": {"accessory", "clothing", "weapon", "vehicle", "magic"},
    "thematic": {"demon", "angelic", "horror", "fantasy", "cyberpunk", "gothic",
                 "steampunk", "noir", "retro", "kawaii", "watercolor", "space",
                 "underwater", "warrior", "magical_girl"},
    "misc": {"food", "animal", "furry", "pose"},
}

# Intent → theme priority (higher = more tags from that theme)
_INTENT_THEME_PRIORITY: dict[str, dict[str, int]] = {
    "general": {"appearance": 2, "environment": 2, "lighting_atmosphere": 1,
                "effects": 1, "props": 1, "thematic": 1, "misc": 1},
    "portrait": {"appearance": 3, "lighting_atmosphere": 2, "props": 1,
                 "effects": 1, "environment": 1, "thematic": 1, "misc": 1},
    "action": {"props": 3, "effects": 2, "appearance": 2, "environment": 1,
               "lighting_atmosphere": 1, "thematic": 1, "misc": 1},
    "environment": {"environment": 3, "lighting_atmosphere": 2, "thematic": 1,
                    "appearance": 1, "props": 1, "effects": 1, "misc": 1},
    "horror": {"thematic": 3, "lighting_atmosphere": 2, "effects": 2,
               "appearance": 1, "environment": 1, "props": 1, "misc": 1},
    "romantic": {"lighting_atmosphere": 3, "appearance": 2, "thematic": 2,
                 "props": 1, "effects": 1, "environment": 1, "misc": 1},
    "fantasy": {"thematic": 3, "props": 2, "effects": 2, "appearance": 1,
                "environment": 1, "lighting_atmosphere": 1, "misc": 1},
}


def compute_theme_budget(
    intent: str,
    resolved_categories: list[str],
    base_budget: tuple[int, int],
) -> dict[str, int]:
    """Compute max tags per theme group for a variation.
    Returns dict mapping theme_group -> max_tags (0 means no limit)."""
    priority = _INTENT_THEME_PRIORITY.get(intent, _INTENT_THEME_PRIORITY["general"])
    cat_to_theme: dict[str, str] = {}
    for theme, cats in _THEME_GROUPS.items():
        for cat in cats:
            cat_to_theme[cat] = theme

    # Count how many resolved categories fall into each theme
    theme_counts: dict[str, int] = {}
    for cat in resolved_categories:
        theme = cat_to_theme.get(cat, "misc")
        theme_counts[theme] = theme_counts.get(theme, 0) + 1

    if not theme_counts:
        return {}

    min_t, max_t = base_budget

    # Total tag slots this variation aims to fill. We scale priorities against
    # the SUM of all priorities so a high-priority theme actually gets a larger
    # slice (the old code divided by priority_total and collapsed to 1 for
    # typical medium creativity, rendering priorities a no-op).
    desired_total = max(4, max_t * max(1, len(theme_counts)))
    priority_total = sum(priority.get(t, 1) for t in theme_counts) or 1
    result: dict[str, int] = {}
    for theme, count in theme_counts.items():
        pri = priority.get(theme, 1)
        share = desired_total * pri / priority_total
        result[theme] = max(min_t, int(round(share)))

    return result


def track_variation_theme_usage(
    variation_tags: list[str],
) -> dict[str, int]:
    """Count how many tags from each theme group are in a variation's output."""
    cat_to_theme: dict[str, str] = {}
    for theme, cats in _THEME_GROUPS.items():
        for cat in cats:
            cat_to_theme[cat] = theme

    usage: dict[str, int] = {}
    for tag in variation_tags:
        cats = get_tag_categories(tag)
        for cat in cats:
            theme = cat_to_theme.get(cat, "misc")
            usage[theme] = usage.get(theme, 0) + 1
    return usage


def _cooccurrence_bonus_fast(tag: str, selected: list[str]) -> float:
    """Co-occurrence bonus between tag and already-selected tags."""
    if not selected:
        return 0.0
    related = get_cooccurrence_tags(tag, limit=20)
    if not related:
        return 0.0
    bonus = 0.0
    sel_lower = {s.lower().strip() for s in selected}
    for r in related:
        rt = (r.get("tag") or "").lower().strip()
        if rt in sel_lower:
            bonus += float(r.get("weight", 1.0))
    return min(bonus, SCORING["cooc_cap"])