pylib_ai_arm_statements/arm_statements/structure.py

"""Шаг 8 — подтверждение неуверенного кода структурой бланка и выбор редакции (спека «Шаг 8»)."""

from collections.abc import Sequence
from dataclasses import dataclass
from datetime import date

from arm_statements.checks import form_variant, subline_parents
from arm_statements.evaluation import EvalLine, is_check_line
from arm_statements.loader import Rules
from arm_statements.model import DocumentType, Edition, FormVariant

FSBU4_FROM = date(2025, 1, 1)  # эффективная дата после 2024-12-31 → fsbu4

_CODE_LEN = 4
_SUBLINE_LEN = 5
_CONFIDENT = 100


def edition_by_date(effective: date | None) -> Edition | None:
    if effective is None:
        return None
    return Edition.FSBU4 if effective >= FSBU4_FROM else Edition.RAS_66N


@dataclass(frozen=True, slots=True)
class StructureResult:
    edition: Edition | None
    unconfirmed: frozenset[str]  # ключи строк с неуверенным кодом, не подтверждённым структурой


# Элемент последовательности: (код, уверенный).
_Item = tuple[str, bool]


def _optimum(items: Sequence[_Item], order: Sequence[str], scale: int, skip: int | None = None) -> int:
    """Взвешенная LCS `items` × `order` по равенству кода; вес позиции `scale * уверенный + 1`.

    `skip` — индекс позиции `items`, которой запрещено входить. Сложность O(|items|·|order|)."""
    if not items or not order:
        return 0
    previous = [0] * (len(order) + 1)
    for index, (code, confident) in enumerate(items):
        weight = None if index == skip else scale * int(confident) + 1
        current = [0] * (len(order) + 1)
        for j, order_code in enumerate(order, start=1):
            best = max(previous[j], current[j - 1])
            if weight is not None and code == order_code:
                best = max(best, previous[j - 1] + weight)
            current[j] = best
        previous = current
    return previous[-1]


def _editions(rules: Rules) -> tuple[Edition, ...]:
    return rules.code_orders.editions or tuple(Edition)


def _pick(optima: dict[Edition, int], effective: date | None) -> tuple[Edition | None, tuple[Edition, ...]]:
    """Редакция с большим оптимумом; равенство — по дате. Второе значение — редакции, чьи оптимумы лучшие."""
    top = max(optima.values())
    tied = tuple(edition for edition, value in optima.items() if value == top)
    if len(tied) == 1:
        return tied[0], tied
    by_date = edition_by_date(effective)
    return (by_date if by_date in tied else None), tied


def _filtered(items: Sequence[_Item], order: Sequence[str]) -> tuple[list[_Item], list[int]]:
    allowed = set(order)
    kept = [(item, index) for index, item in enumerate(items) if item[0] in allowed]
    return [item for item, _ in kept], [index for _, index in kept]


def _select(
    document_type: DocumentType,
    items: Sequence[_Item],
    variant: FormVariant,
    effective: date | None,
    rules: Rules,
) -> tuple[Edition | None, tuple[Edition, ...], dict[Edition, int]]:
    scale = len(items) + 1
    optima = {
        edition: _optimum(
            _filtered(items, order := rules.code_orders.sequence(edition, document_type, variant))[0],
            order,
            scale,
        )
        for edition in _editions(rules)
    }
    edition, tied = _pick(optima, effective)
    return edition, tied, optima


def choose_edition(
    document_type: DocumentType,
    codes: Sequence[tuple[str, bool]],
    effective: date | None,
    rules: Rules,
    *,
    variant: FormVariant | None = None,
) -> Edition | None:
    """Шаги 2–3 и 6 без подтверждения кодов; `codes` — пары `(код, уверенный)` строк `ok` по порядку."""
    if document_type is DocumentType.DISCLOSURE:
        return edition_by_date(effective)
    if variant is None:
        variant = form_variant(document_type, [code for code, _ in codes]) or FormVariant.FULL
    items = [(code, confident) for code, confident in codes if len(code) == _CODE_LEN]
    return _select(document_type, items, variant, effective, rules)[0]


def _is_uncertain(line: EvalLine) -> bool:
    return (line.code_confidence is None or line.code_confidence < _CONFIDENT) and not line.code_fixed


def _confirm_sublines(checked: Sequence[EvalLine], codes: Sequence[str], confirmed_keys: set[str]) -> set[str]:
    """Пятизначные неуверенные коды: подтверждены родителем, стоящим раньше (§11.9); возвращает неподтверждённые."""
    unconfirmed: set[str] = set()
    parents = subline_parents(codes)
    for position, line in enumerate(checked):
        code = line.line_code or ""
        if len(code) == _CODE_LEN or not _is_uncertain(line):
            continue
        parent = parents[position]
        found = None
        if len(code) == _SUBLINE_LEN and parent is not None:
            found = next((e for e in reversed(checked[:position]) if e.line_code == parent), None)
        if found is not None and (not _is_uncertain(found) or found.key in confirmed_keys):
            confirmed_keys.add(line.key)
        else:
            unconfirmed.add(line.key)
    return unconfirmed


def confirm_codes(
    lines: Sequence[EvalLine],
    document_type: DocumentType,
    variant: FormVariant | None,
    effective: date | None,
    rules: Rules,
) -> StructureResult:
    if document_type is DocumentType.DISCLOSURE:
        return StructureResult(edition_by_date(effective), frozenset())
    checked = sorted(
        (line for line in lines if is_check_line(line) and line.order is not None),
        key=lambda line: line.order or 0,
    )
    codes = [line.line_code or "" for line in checked]
    if variant is None:
        variant = form_variant(document_type, codes) or FormVariant.FULL
    four = [line for line in checked if len(line.line_code or "") == _CODE_LEN]
    items: list[_Item] = [(line.line_code or "", not _is_uncertain(line)) for line in four]
    edition, tied, _ = _select(document_type, items, variant, effective, rules)

    # Четырёхзначные неуверенные коды: входят во все лучшие подпоследовательности выбранной редакции
    # (при равенстве без даты — во все лучшие обеих редакций, P1-4).
    unconfirmed: set[str] = set()
    confirmed_keys: set[str] = set()
    scale = len(items) + 1
    for index, line in enumerate(four):
        if not _is_uncertain(line):
            confirmed_keys.add(line.key)
            continue
        confirmed = True
        for candidate in tied if edition is None else (edition,):
            order = rules.code_orders.sequence(candidate, document_type, variant)
            if items[index][0] not in order:
                confirmed = False
                break
            kept, indices = _filtered(items, order)
            total = _optimum(kept, order, scale)
            if _optimum(kept, order, scale, skip=indices.index(index)) >= total:
                confirmed = False
                break
        (confirmed_keys if confirmed else unconfirmed).add(line.key)

    unconfirmed |= _confirm_sublines(checked, codes, confirmed_keys)
    return StructureResult(edition, frozenset(unconfirmed))