tower_game/behavior.py

from __future__ import annotations

from collections import Counter
from copy import deepcopy
from typing import Any

from .schemas import BehaviorSummary


SCALING_LABELS = {
    "str_scaling_dial": "strength",
    "agi_scaling_dial": "agility",
    "int_scaling_dial": "intelligence",
    "def_scaling_dial": "endurance",
}


def summarize_behavior(
    state: dict[str, Any],
    *,
    start_index: int | None = None,
) -> BehaviorSummary:
    checkpoint = int(state.get("behavior_checkpoint_index", 0))
    start = checkpoint if start_index is None else max(0, int(start_index))
    history = [str(entry) for entry in state.get("action_history", [])[start:]]
    actions = Counter()
    skill_use = Counter()
    element_use = Counter()
    scaling_use = Counter()
    effect_use = Counter()
    equipment = Counter()
    intent_responses = Counter()

    for entry in history:
        if entry == "used_strike":
            actions["strike"] += 1
        elif entry == "used_defend":
            actions["defend"] += 1
        elif entry == "used_item":
            actions["item"] += 1
        elif entry.startswith("used_skill:"):
            skill_use[entry.split(":", 1)[1]] += 1
            actions["skill"] += 1
        elif entry.startswith("used_signature:"):
            continue
        elif entry.startswith("equipped_"):
            equipment[entry.removeprefix("equipped_")] += 1

    skill_by_id = {
        skill.get("id"): skill
        for skill in (
            list(state.get("run_skill_registry", {}).values())
            + list(state.get("active_skills", []))
        )
    }
    for skill_id, count in skill_use.items():
        skill = skill_by_id.get(skill_id)
        if skill:
            element_use[str(skill.get("element", "Physical"))] += count
            scaling_use[str(skill.get("scaling_stat", "str_scaling_dial"))] += count
            effect_use[str(skill.get("effect_id", "damage"))] += count
            rider = str(skill.get("rider_effect", "none"))
            if rider != "none":
                effect_use[rider] += count

    events = state.get("combat_events", [])
    defensive_reactions = sum(
        1
        for event in events
        if event.get("actor") == "Player"
        and any(word in str(event.get("text", "")).lower() for word in ("brace", "guard", "fortif"))
    )
    for index, event in enumerate(events):
        if event.get("actor") != "Enemy":
            continue
        text = str(event.get("text", "")).lower()
        intent = "attack" if "damage" in text else "defend" if "defense up" in text else "tactic"
        if index + 1 < len(events) and events[index + 1].get("actor") == "Player":
            intent_responses[intent] += 1

    total = sum(actions.values())
    ratios = {
        key: round(actions[key] / total, 3) if total else 0.0
        for key in ("strike", "skill", "defend", "item")
    }
    patterns = [
        f"frequent_{name}"
        for name, count in actions.most_common(2)
        if total and count / total >= 0.35
    ]
    if state.get("heals_refused", 0):
        patterns.append("refuses_healing")

    affinities = Counter()
    affinities["strength"] += actions["strike"] ** 1.25
    affinities["Physical"] += actions["strike"] ** 1.25
    affinities["damage"] += actions["strike"] ** 1.25
    affinities["endurance"] += actions["defend"] ** 1.25
    affinities["guard"] += actions["defend"] ** 1.25
    affinities["armor"] += actions["defend"] ** 1.25
    for scaling, count in scaling_use.items():
        affinities[SCALING_LABELS.get(scaling, scaling)] += count ** 1.25
    for element, count in element_use.items():
        affinities[element] += count ** 1.25
    for effect, count in effect_use.items():
        affinities[effect] += count ** 1.25

    dominant_action = actions.most_common(1)[0][0] if actions else ""
    dominant_scaling = scaling_use.most_common(1)[0][0] if scaling_use else (
        "str_scaling_dial" if dominant_action == "strike"
        else "def_scaling_dial" if dominant_action == "defend"
        else ""
    )
    dominant_element = element_use.most_common(1)[0][0] if element_use else (
        "Physical" if dominant_action == "strike" else ""
    )

    return BehaviorSummary(
        total_actions=total,
        action_ratios=ratios,
        skill_use=dict(skill_use),
        element_use=dict(element_use),
        defensive_reactions=defensive_reactions,
        intent_responses=dict(intent_responses),
        healing_accepted=state.get("healing_decisions", []).count("accepted"),
        healing_refused=int(state.get("heals_refused", 0)),
        equipment_preferences=dict(equipment),
        low_health_turns=sum(
            1
            for event in events
            if "hp" in str(event.get("text", "")).lower()
            and any(token in str(event.get("text", "")) for token in ("1/", "2/", "3/"))
        ),
        repeated_patterns=patterns[:6],
        affinity_scores={key: round(value, 3) for key, value in affinities.items()},
        dominant_action=dominant_action,
        dominant_scaling_stat=dominant_scaling,
        dominant_element=dominant_element,
    )


def blend_behavior_summaries(
    earlier: BehaviorSummary,
    recent: BehaviorSummary,
    *,
    recent_weight: float = 0.7,
) -> BehaviorSummary:
    earlier_weight = 1.0 - recent_weight

    def blend_dict(left: dict[str, float], right: dict[str, float]) -> dict[str, float]:
        keys = set(left) | set(right)
        return {
            key: round(float(left.get(key, 0)) * earlier_weight + float(right.get(key, 0)) * recent_weight, 3)
            for key in keys
        }

    result = deepcopy(recent.model_dump())
    result["total_actions"] = round(
        earlier.total_actions * earlier_weight + recent.total_actions * recent_weight
    )
    result["action_ratios"] = blend_dict(earlier.action_ratios, recent.action_ratios)
    result["skill_use"] = {
        key: round(value)
        for key, value in blend_dict(earlier.skill_use, recent.skill_use).items()
    }
    result["element_use"] = {
        key: round(value)
        for key, value in blend_dict(earlier.element_use, recent.element_use).items()
    }
    result["affinity_scores"] = blend_dict(
        earlier.affinity_scores, recent.affinity_scores
    )
    dominant = max(result["affinity_scores"], key=result["affinity_scores"].get, default="")
    scaling_by_affinity = {
        "strength": "str_scaling_dial",
        "agility": "agi_scaling_dial",
        "intelligence": "int_scaling_dial",
        "endurance": "def_scaling_dial",
    }
    result["dominant_scaling_stat"] = scaling_by_affinity.get(
        dominant, recent.dominant_scaling_stat or earlier.dominant_scaling_stat
    )
    result["dominant_element"] = recent.dominant_element or earlier.dominant_element
    result["dominant_action"] = recent.dominant_action or earlier.dominant_action
    result["repeated_patterns"] = list(
        dict.fromkeys([*recent.repeated_patterns, *earlier.repeated_patterns])
    )[:6]
    return BehaviorSummary.model_validate(result)


def evolution_behavior_summary(
    state: dict[str, Any],
    stage: int,
) -> BehaviorSummary:
    recent = summarize_behavior(state)
    stored = state.get("behavior_summaries", [])
    if stage >= 2 and stored:
        earlier = BehaviorSummary.model_validate(stored[0])
        return blend_behavior_summaries(earlier, recent, recent_weight=0.7)
    return recent