pylib_ai_arm_statements/arm_statements/columns.py

"""Шаг 2: колонки значений, подписи и даты колонок."""

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

from arm_statements.amounts import parse_amount
from arm_statements.codes import is_code_cell
from arm_statements.dates import expected_dates, parse_period_label, resolve_relative
from arm_statements.issues import WarningCode, make_warning
from arm_statements.layout import Grid, cell_at
from arm_statements.loader import Rules
from arm_statements.model import Cell, PeriodKind, Source, StatementColumn, StatementHeader, StatementWarning

_FULL_CONFIDENCE = 100.0


@dataclass(frozen=True, slots=True)
class ColumnDraft:
    """Колонка сегмента до сопоставления; `date` — только у абсолютной подписи."""

    grid_col: int
    label: str
    label_confidence: float | None
    kind: PeriodKind
    date: date | None


@dataclass(frozen=True, slots=True)
class DocumentDates:
    effective: date | None
    confident_effective: date | None
    latest_absolute: date | None
    confident_latest_absolute: date | None


def _text(cell: Cell | None) -> str:
    return "" if cell is None else cell.text.strip()


def _code_rows(rows: Grid, data_rows: Sequence[int], code_col: int, rules: Rules) -> list[int]:
    return [r for r in data_rows if is_code_cell(_text(cell_at(rows, r, code_col)), rules)]


def _is_number(cell: Cell | None, source: Source, rules: Rules) -> bool:
    if cell is None:
        return False
    if isinstance(cell.value, Decimal):
        return True
    return parse_amount(cell.text, source, rules).amount is not None


def has_numbers(  # noqa: PLR0913, PLR0917 - сигнатура задана планом (Task 9)
    rows: Grid, data_rows: Sequence[int], col: int, code_col: int, source: Source, rules: Rules
) -> bool:
    """В строках с кодом есть ячейка столбца с числом (значение-число или разобранная сумма)."""
    return any(_is_number(cell_at(rows, r, col), source, rules) for r in _code_rows(rows, data_rows, code_col, rules))


def _width(rows: Grid, row_indexes: Sequence[int]) -> int:
    return max((len(rows[r]) for r in row_indexes if 0 <= r < len(rows)), default=0)


def amount_columns(rows: Grid, code_col: int, rules: Rules, *, source: Source = Source.EXCEL) -> tuple[int, ...]:
    """Столбцы правее кодов, где в строках с кодом есть числа, по возрастанию."""
    every = range(len(rows))
    return tuple(
        col
        for col in range(code_col + 1, _width(rows, every))
        if has_numbers(rows, every, col, code_col, source, rules)
    )


def _label(rows: Grid, header_rows: tuple[int, ...], col: int) -> Cell | None:
    """Подпись колонки: строки заголовка через пробел, уверенность — минимум; одна ячейка берётся как есть."""
    cells = [c for r in header_rows if (c := cell_at(rows, r, col)) is not None and c.text.strip()]
    if not cells:
        return None
    if len(cells) == 1:
        return cells[0]
    text = " ".join(c.text.strip() for c in cells)
    return Cell(text, None, min(c.confidence for c in cells), cells[0].ref)


def segment_columns(  # noqa: PLR0913, PLR0917 - сигнатура задана планом (Task 9)
    rows: Grid,
    header_rows: tuple[int, ...],
    data_rows: Sequence[int],
    code_col: int,
    source: Source,
    rules: Rules,
) -> tuple[list[ColumnDraft], list[StatementWarning]]:
    """Колонки значений сегмента и предупреждения о пропущенных/недатированных столбцах."""
    drafts: list[ColumnDraft] = []
    warnings: list[StatementWarning] = []
    for col in range(code_col + 1, _width(rows, [*header_rows, *data_rows])):
        numbers = has_numbers(rows, data_rows, col, code_col, source, rules)
        if not header_rows:
            if numbers:
                drafts.append(ColumnDraft(col, "", None, PeriodKind.NONE, None))
            continue
        cell = _label(rows, header_rows, col)
        parsed = parse_period_label(cell, rules)
        if cell is not None and parsed.excluded:
            message = f"столбец {col} («{cell.text}») пропущен: не период"
            warnings.append(make_warning(WarningCode.COLUMN_SKIPPED, message, rules, column_index=col))
        elif cell is not None and parsed.kind is not PeriodKind.NONE:
            drafts.append(ColumnDraft(col, cell.text, cell.confidence, parsed.kind, parsed.date))
        elif numbers:
            label, conf = (cell.text, cell.confidence) if cell is not None else ("", None)
            message = f"у столбца {col} с числами нет разобранной подписи периода"
            warnings.append(make_warning(WarningCode.COLUMN_WITHOUT_DATE, message, rules, column_index=col))
            drafts.append(ColumnDraft(col, label, conf, PeriodKind.NONE, None))
    return drafts, warnings


def _latest(drafts: Sequence[ColumnDraft], *, confident: bool) -> date | None:
    days = [
        d.date
        for d in drafts
        if d.kind is PeriodKind.ABSOLUTE
        and d.date is not None
        and (not confident or d.label_confidence == _FULL_CONFIDENCE)
    ]
    return max(days, default=None)


def document_dates(drafts: Sequence[ColumnDraft], header: StatementHeader) -> DocumentDates:
    """Эффективная дата и её «уверенный» двойник: неуверенное само себя не подтверждает."""
    latest = _latest(drafts, confident=False)
    confident_latest = _latest(drafts, confident=True)
    report = header.report_date
    sure_report = report.value if report.confidence == _FULL_CONFIDENCE else None
    return DocumentDates(
        effective=report.value or latest,
        confident_effective=sure_report or confident_latest,
        latest_absolute=latest,
        confident_latest_absolute=confident_latest,
    )


def finalize_columns(
    drafts: Sequence[ColumnDraft], dates: DocumentDates, rules: Rules
) -> tuple[tuple[StatementColumn, ...], list[StatementWarning]]:
    """Колонки документа: даты относительных подписей и проверка неуверенных абсолютных."""
    expected = None if dates.confident_effective is None else expected_dates(dates.confident_effective)
    columns: list[StatementColumn] = []
    warnings: list[StatementWarning] = []
    for draft in drafts:
        if draft.kind is PeriodKind.ABSOLUTE:
            period = draft.date
            unsure = draft.label_confidence is not None and draft.label_confidence < _FULL_CONFIDENCE
            if unsure and (expected is None or period not in expected):
                message = f"дата подписи столбца {draft.grid_col} («{draft.label}») не подтверждена"
                warnings.append(
                    make_warning(WarningCode.COLUMN_DATE_UNCERTAIN, message, rules, column_index=draft.grid_col)
                )
        else:
            period = resolve_relative(draft.kind, dates.effective)
        columns.append(StatementColumn(draft.grid_col, draft.label, draft.label_confidence, period, draft.kind))
    return tuple(columns), warnings