arm-core/api/src/domain/vgo_distribution.py

"""Чистый выбор агрегатов и строк ВГО после поиска встречных сумм."""

import uuid
from collections import Counter, defaultdict
from collections.abc import Mapping, Sequence, Set
from dataclasses import dataclass
from decimal import Decimal
from typing import Literal

from src.domain.consolidated_osv import ConsolidatedRow

RowKey = tuple[uuid.UUID, ...]
ZERO = Decimal("0")
Side = Literal["D", "K"]
# Сторона → счёт → правила, в которые счёт входит именно этой стороной (с исключениями стороны, §4).
RuleOwners = Mapping[Side, Mapping[uuid.UUID, Set[uuid.UUID]]]


@dataclass(frozen=True)
class Distribution:
    """Результат выбора: расхождения, строки без пары и остаток для ручного разбора."""

    selected_diff: tuple[int, ...]
    auto_rules: Mapping[RowKey, uuid.UUID]
    unassigned: frozenset[RowKey]
    used_ids: frozenset[uuid.UUID]


@dataclass(frozen=True)
class DistributionInput:
    rows: Sequence[ConsolidatedRow]
    exact_ids: Set[uuid.UUID]
    diff_ids: Sequence[Set[uuid.UUID]]
    manual_keys: Set[RowKey]
    owners: RuleOwners
    members: Set[uuid.UUID]


def account_owners(
    selections: Sequence[tuple[uuid.UUID, Set[uuid.UUID], Set[uuid.UUID]]],
) -> dict[Side, dict[uuid.UUID, frozenset[uuid.UUID]]]:
    """Для каждого раскрытого счёта — правила, в которые он входит, отдельно по дебету и кредиту (§5 п. 2a).

    Вход — (правило, счета дебетовой стороны, счета кредитовой стороны) уже с исключениями своей стороны. Стороны не
    сливаются: исключение одной стороны (например, 76.07.1 из кредита правила 6) не должно возвращаться через маску
    другой (дебетовую «76»), иначе кредитовая строка ложно попадает в несколько правил (AURORA-2034).
    """
    owners: dict[Side, defaultdict[uuid.UUID, set[uuid.UUID]]] = {"D": defaultdict(set), "K": defaultdict(set)}
    for rule_id, debit_ids, credit_ids in selections:
        for account_id in debit_ids:
            owners["D"][account_id].add(rule_id)
        for account_id in credit_ids:
            owners["K"][account_id].add(rule_id)
    return {
        side: {account_id: frozenset(rule_ids) for account_id, rule_ids in by_account.items()}
        for side, by_account in owners.items()
    }


def row_side(closing_debit: Decimal | None, closing_credit: Decimal | None) -> Side:
    """Сторона строки свода: дебет, если есть дебетовое сальдо или нет кредитового; иначе кредит (§5 п. 2a)."""
    return "D" if (closing_debit or ZERO) != ZERO or (closing_credit or ZERO) == ZERO else "K"


def rules_for(owners: RuleOwners, account_id: uuid.UUID | None, side: Side) -> frozenset[uuid.UUID]:
    """Правила, в которые счёт входит заданной стороной; без счёта — ни одного."""
    if account_id is None:
        return frozenset()
    return frozenset(owners.get(side, {}).get(account_id, frozenset()))


def select_exact_aggregates(
    used_ids: Set[uuid.UUID], candidate_ids: Sequence[Set[uuid.UUID]]
) -> tuple[tuple[int, ...], frozenset[uuid.UUID]]:
    """Закрепить точные агрегаты по порядку правил без повторного использования строк."""
    used = set(used_ids)
    selected = []
    for index, ids in enumerate(candidate_ids):
        if ids and used.isdisjoint(ids):
            selected.append(index)
            used.update(ids)
    return tuple(selected), frozenset(used)


def distribute_vgo_rows(data: DistributionInput) -> Distribution:
    """Выбрать непересекающиеся расхождения, затем однозначные строки без пары и остаток.

    Точные агрегаты уже закреплены по порядку правил. Пересекающиеся расхождения все остаются
    нераспределёнными; ручные решения не заменяются автоматическими.
    """
    candidates = [(index, ids) for index, ids in enumerate(data.diff_ids) if ids and data.exact_ids.isdisjoint(ids)]
    counts = Counter(entry_id for _, ids in candidates for entry_id in ids)
    selected = tuple(index for index, ids in candidates if all(counts[entry_id] == 1 for entry_id in ids))
    used = set(data.exact_ids)
    for index in selected:
        used.update(data.diff_ids[index])

    auto_rules: dict[RowKey, uuid.UUID] = {}
    unassigned: set[RowKey] = set()
    for row in data.rows:
        if not used.isdisjoint(row.entry_ids):
            continue
        manual = row.entry_ids in data.manual_keys
        # §5 п. 2a: разносятся только строки между разными компаниями периметра; строка с контрагентом вне периметра
        # или «сам на себя» в таблицу не попадает, если по ней нет ручного решения аналитика (AURORA-2034)
        in_perimeter = (
            row.legal_entity_id in data.members
            and row.counterparty_id in data.members
            and row.legal_entity_id != row.counterparty_id
        )
        if not manual and not in_perimeter:
            continue
        rule_ids = rules_for(data.owners, row.account_id, row_side(row.closing_debit, row.closing_credit))
        if (
            not manual
            and len(rule_ids) == 1
            and (row.closing_debit not in (None, ZERO) or row.closing_credit not in (None, ZERO))
        ):
            auto_rules[row.entry_ids] = next(iter(rule_ids))
        elif rule_ids or manual:
            unassigned.add(row.entry_ids)
    return Distribution(selected, auto_rules, frozenset(unassigned), frozenset(used))