"""
features.py -- the single source of truth for what the Modular Mind sees and the
action set it chooses from. Imported by the duel simulator (training env), the
torch model, and the numpy inference brain so all three agree exactly.
"""
from __future__ import annotations

import numpy as np

ACTIONS = ["IDLE", "APPROACH", "RETREAT", "CLEAVE", "BLOCK"]

FEATURES = [
    "abs_dist",          # 0  normalised |player - boss|, 0 close .. 1 far
    "in_cleave_range",   # 1  player within the boss's cleave reach
    "in_mid_range",      # 2  player in the "step in and swing" band
    "boss_hp",           # 3  boss hp fraction
    "player_hp",         # 4  player hp fraction
    "boss_ready",        # 5  1 = cleave off cooldown, 0 = recovering
    "player_pressuring", # 6  player attacking or moving in
    "player_open",       # 7  player in attack/roll recovery = punishable
    "player_blocking",   # 8  player is holding block
    "player_threat",     # 9  player is swinging AT the boss in melee range (block cue)
    "bias",              # 10
]
NF = len(FEATURES)


def extract_features(s: dict) -> np.ndarray:
    arena = float(s.get("arenaW", 900))
    dist = abs(float(s["playerX"]) - float(s["bossX"]))
    abs_dist = min(dist / (arena * 0.6), 1.0)
    cleave_reach = float(s.get("cleaveReach", 160))
    f = np.zeros(NF, dtype=np.float32)
    f[0] = abs_dist
    f[1] = 1.0 if dist <= cleave_reach else 0.0
    f[2] = 1.0 if cleave_reach < dist <= cleave_reach * 2.2 else 0.0
    f[3] = float(s.get("bossHP", 1.0))
    f[4] = float(s.get("playerHP", 1.0))
    f[5] = 1.0 if float(s.get("bossCooldown", 0.0)) <= 0.0 else 0.0
    f[6] = 1.0 if (s.get("playerAttacking") or s.get("playerApproaching")) else 0.0
    f[7] = 1.0 if s.get("playerRecovering") else 0.0
    f[8] = 1.0 if s.get("playerBlocking") else 0.0
    f[9] = 1.0 if s.get("playerThreat") else 0.0
    f[10] = 1.0
    return f


def legal_mask(s: dict) -> np.ndarray:
    """CLEAVE is illegal while on cooldown; everything else always allowed."""
    m = np.ones(len(ACTIONS), dtype=np.float32)
    if float(s.get("bossCooldown", 0.0)) > 0.0:
        m[ACTIONS.index("CLEAVE")] = 0.0
    return m