src/yaml_edit/ml/features.py

"""Bounded input-only features and deterministic copy/normalization actions."""

import hashlib
import math
import re
from dataclasses import dataclass

from yaml_edit import engine
from yaml_edit.models import Candidate, FieldKind, Scalar, YamlEditError
from yaml_edit.parsing import Parsed, parse_source
from yaml_edit.suggestion import MAX_INSTRUCTION_BYTES

FEATURE_DIMENSION = 384
ACTION_FEATURES = 13
MAX_ML_SOURCE_BYTES = 64 * 1024
MAX_ML_CANDIDATES = 64
MAX_PROPOSAL_ATTEMPTS = 256
MAX_RETAINED_ACTIONS = 256
_KIND_ORDER = tuple(FieldKind)
_KIND_INDEX = {kind: index for index, kind in enumerate(_KIND_ORDER)}
_WORD_NUMBERS = {
    "zero": 0,
    "one": 1,
    "two": 2,
    "three": 3,
    "four": 4,
    "five": 5,
    "six": 6,
    "seven": 7,
    "eight": 8,
    "nine": 9,
    "ten": 10,
}
_TOKEN = re.compile(r"\S+")
_QUANTITY = re.compile(r"[0-9]+(?:\.[0-9]+)?(?:m|[KMGTPE]i|[kMGTPE])?", re.ASCII)
_TRIM = "\"'`()[]{}<>,;.!?"


class MlInputError(ValueError):
    """An input exceeds the experiment's documented deterministic boundary."""


@dataclass(frozen=True, slots=True)
class Action:
    """One core-valid candidate and instruction-derived replacement value."""

    candidate: Candidate
    value: Scalar
    raw_value: str
    span: tuple[int, int]
    features: tuple[float, ...]


@dataclass(frozen=True, slots=True)
class PreparedInput:
    """Model input derived exclusively from source, instruction, and inspection."""

    vector: tuple[float, ...]
    actions: tuple[Action, ...]


def _bounded_source(source: str) -> str:
    try:
        raw = source.encode("utf-8")
    except UnicodeEncodeError as error:
        raise MlInputError("ML source is not valid UTF-8 text.") from error
    if len(raw) > MAX_ML_SOURCE_BYTES:
        raise MlInputError(f"ML source exceeds the {MAX_ML_SOURCE_BYTES}-byte experiment limit.")
    return source


def _bounded_instruction(instruction: str) -> str:
    try:
        raw = instruction.encode("utf-8")
    except UnicodeEncodeError as error:
        raise MlInputError("Instruction is not valid UTF-8 text.") from error
    if not raw or len(raw) > MAX_INSTRUCTION_BYTES or instruction.splitlines() != [instruction]:
        raise MlInputError(
            f"Instruction must be one line of at most {MAX_INSTRUCTION_BYTES} UTF-8 bytes."
        )
    return instruction


def _hash_feature(vector: list[float], feature: str) -> None:
    digest = hashlib.sha256(feature.encode("utf-8")).digest()
    index = int.from_bytes(digest[:4], "big") % len(vector)
    sign = 1.0 if digest[4] & 1 else -1.0
    vector[index] += sign


def _instruction_vector(instruction: str) -> tuple[float, ...]:
    text = instruction.casefold()
    vector = [0.0] * FEATURE_DIMENSION
    padded = f"^{text}$"
    for size in (2, 3, 4):
        for start in range(len(padded) - size + 1):
            _hash_feature(vector, f"c{size}:{padded[start : start + size]}")
    for word in re.findall(r"[a-z0-9]+", text, re.ASCII):
        _hash_feature(vector, f"w:{word}")
    norm = math.sqrt(sum(value * value for value in vector))
    if norm:
        vector = [value / norm for value in vector]
    return tuple(vector)


def _trimmed_span(instruction: str, start: int, end: int) -> tuple[int, int] | None:
    while start < end and instruction[start] in _TRIM:
        start += 1
    while end > start and instruction[end - 1] in _TRIM:
        end -= 1
    return None if start == end else (start, end)


def _lexemes(instruction: str) -> tuple[tuple[str, str, int, int], ...]:
    result: list[tuple[str, str, int, int]] = []
    for match in _TOKEN.finditer(instruction):
        span = _trimmed_span(instruction, match.start(), match.end())
        if span is not None:
            start, end = span
            raw = instruction[start:end]
            result.append((raw, raw.casefold(), start, end))
    return tuple(result)


def _replica_values(
    lexemes: tuple[tuple[str, str, int, int], ...],
) -> tuple[tuple[Scalar, str, int, int], ...]:
    values: list[tuple[Scalar, str, int, int]] = []
    for raw, folded, start, end in lexemes:
        if folded.isascii() and folded.isdigit():
            values.append((int(folded), raw, start, end))
        elif folded in _WORD_NUMBERS:
            values.append((_WORD_NUMBERS[folded], raw, start, end))
    return tuple(values)


def _image_values(
    lexemes: tuple[tuple[str, str, int, int], ...],
) -> tuple[tuple[Scalar, str, int, int], ...]:
    values: list[tuple[Scalar, str, int, int]] = []
    for index, (raw, _, start, end) in enumerate(lexemes):
        previous = "" if index == 0 else lexemes[index - 1][1]
        marked = any(marker in raw for marker in ("/", ":", "@"))
        if marked or previous in {"to", "using", "become"}:
            values.append((raw, raw, start, end))
    return tuple(values)


def _resource_values(
    instruction: str,
    lexemes: tuple[tuple[str, str, int, int], ...],
) -> tuple[tuple[Scalar, str, int, int], ...]:
    values: list[tuple[Scalar, str, int, int]] = []
    for index, (raw, folded, start, end) in enumerate(lexemes):
        if _QUANTITY.fullmatch(raw) is not None:
            value = raw
            if (
                raw.isascii()
                and raw.isdigit()
                and index + 1 < len(lexemes)
                and lexemes[index + 1][1].rstrip("s") == "millicore"
            ):
                value = f"{raw}m"
            values.append((value, raw, start, end))
        if (
            folded == "five"
            and index + 2 < len(lexemes)
            and lexemes[index + 1][1] == "hundred"
            and lexemes[index + 2][1].rstrip("s") == "millicore"
        ):
            phrase_end = lexemes[index + 2][3]
            values.append(("500m", instruction[start:phrase_end], start, phrase_end))
    return tuple(values)


def _values(
    instruction: str,
    candidate: Candidate,
    lexemes: tuple[tuple[str, str, int, int], ...],
) -> tuple[tuple[Scalar, str, int, int], ...]:
    if candidate.kind == FieldKind.REPLICAS:
        return _replica_values(lexemes)
    if candidate.kind == FieldKind.IMAGE:
        return _image_values(lexemes)
    return _resource_values(instruction, lexemes)


def _container_mentioned(instruction: str, candidate: Candidate) -> float:
    if candidate.container is None:
        return 0.0
    pattern = rf"(?<![A-Za-z0-9_-]){re.escape(candidate.container)}(?:'s)?(?![A-Za-z0-9_-])"
    return float(re.search(pattern, instruction, re.IGNORECASE | re.ASCII) is not None)


def _action_features(
    instruction: str, candidate: Candidate, value: Scalar, start: int
) -> tuple[float, ...]:
    features = [0.0] * ACTION_FEATURES
    features[_KIND_INDEX[candidate.kind]] = 1.0
    features[6] = _container_mentioned(instruction, candidate)
    features[7] = float(candidate.container is None)
    features[8] = float(type(value) is int)
    features[9] = float(type(value) is str)
    relative = start / max(1, len(instruction))
    position = min(2, int(relative * 3))
    features[10 + position] = 1.0
    return tuple(features)


type _Attempt = tuple[Candidate, Scalar, str, int, int]


def _attempts(instruction: str, candidates: tuple[Candidate, ...]) -> tuple[_Attempt, ...]:
    lexemes = _lexemes(instruction)
    result: list[_Attempt] = []
    seen: set[tuple[str, type[object], object, int, int]] = set()
    for candidate in candidates:
        for value, raw, start, end in _values(instruction, candidate, lexemes):
            key = (candidate.id, type(value), value, start, end)
            if key in seen:
                continue
            seen.add(key)
            result.append((candidate, value, raw, start, end))
            if len(result) > MAX_PROPOSAL_ATTEMPTS:
                raise MlInputError(
                    f"ML input exceeds the {MAX_PROPOSAL_ATTEMPTS}-attempt proposal budget."
                )
    return tuple(result)


def _actions(
    parsed: Parsed, source: str, instruction: str, attempts: tuple[_Attempt, ...]
) -> tuple[Action, ...]:
    result: list[Action] = []
    for candidate, value, raw, start, end in attempts:
        try:
            # The parsed value was created above for this exact source and never escapes.
            engine._propose_parsed(parsed, source, candidate.id, value)  # noqa: SLF001
        except YamlEditError:
            continue
        result.append(
            Action(
                candidate,
                value,
                raw,
                (start, end),
                _action_features(instruction, candidate, value, start),
            )
        )
        if len(result) > MAX_RETAINED_ACTIONS:
            raise MlInputError(
                f"ML input exceeds the {MAX_RETAINED_ACTIONS}-action retention budget."
            )
    return tuple(result)


def prepare_input(source: str, instruction: str) -> PreparedInput:
    """Build bounded fixed instruction features and core-valid variable actions."""
    bounded_source = _bounded_source(source)
    bounded_instruction = _bounded_instruction(instruction)
    parsed = parse_source(bounded_source)
    if len(parsed.candidates) > MAX_ML_CANDIDATES:
        raise MlInputError(f"ML input exceeds the {MAX_ML_CANDIDATES}-candidate experiment limit.")
    attempts = _attempts(bounded_instruction, parsed.candidates)
    return PreparedInput(
        _instruction_vector(bounded_instruction),
        _actions(parsed, bounded_source, bounded_instruction, attempts),
    )