pylib_ai_arm_statements/arm_statements/tests/test_segments.py

from decimal import Decimal as D

from arm_statements.header import parse_header
from arm_statements.loader import Rules
from arm_statements.model import RowRole
from arm_statements.segments import Layout, find_layout
from arm_statements.tests.builders import excel, hdr, ocr, place

DATES3 = ["На 30 сентября 2025 г.", "На 31 декабря 2024 г.", "На 31 декабря 2023 г."]


def test_npz_sheet_with_own_header_maps_by_date(rules: Rules) -> None:
    first = [
        place({1: "Наименование показателя", 12: "Код", 16: DATES3[0], 21: DATES3[1], 24: DATES3[2]}),
        place({1: "Нематериальные активы", 12: "1110", 16: D("10"), 21: D("9"), 24: D("8")}),
        place({1: "Итого по разделу I", 12: "1100", 16: D("10"), 21: D("9"), 24: D("8")}),
    ]
    sheet4 = [
        place({0: "Наименование показателя", 6: "Код", 8: DATES3[0], 9: DATES3[1], 10: DATES3[2]}),
        place({0: "Заемные средства", 6: "1510", 8: D("3"), 9: D("2"), 10: D("1")}),
    ]
    cand = excel([("Бухгалтерский баланс", first), ("Бухгалтерский баланс. Лист 4", sheet4)])
    layout = find_layout(cand, parse_header(hdr(report_date="30.09.2025"), rules), rules)
    assert [c.column_index for c in layout.columns] == [16, 21, 24]
    assert layout.segments[1].value_cols == ((8, 16), (9, 21), (10, 24))
    assert not [w for w in layout.warnings if w.code.startswith("SEGMENT_")]


def test_repeat_header_mid_table(rules: Rules) -> None:
    rows = [
        ["Наименование показателя", "Код", "На 30.09.2025", "На 31.12.2024"],
        ["Запасы", "1210", "5", "4"],
        ["Наименование показателя", "Код", "На 30.09.2025", "На 31.12.2024"],
        ["Итого", "1200", "5", "4"],
    ]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    assert [s.first_row for s in layout.segments] == [0, 2]
    assert layout.roles[2] is RowRole.HEADER


def test_repeat_header_without_dates_maps_by_order(rules: Rules) -> None:
    rows = [
        ["Наименование", "Код", "На 30.09.2025", "На 31.12.2024"],
        ["Запасы", "1210", "5", "4"],
        ["Наименование", "Код", "отчетн", "прошл"],
        ["Итого", "1200", "5", "4"],
    ]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    assert layout.segments[1].value_cols == ((2, 2), (3, 3))
    assert layout.warnings == ()


def test_wider_segment_is_mismatch(rules: Rules) -> None:
    rows = [
        ["Наименование", "Код", "На 30.09.2025", "На 31.12.2024", "На 31.12.2023"],
        ["Запасы", "1210", "5", "4", "3"],
        ["Наименование", "Код", "графа 3", "графа 4", "графа 5", "графа 6"],
        ["Итого", "1200", "5", "4", "3", "2"],
    ]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    (warning,) = [w for w in layout.warnings if w.code == "SEGMENT_COLUMNS_MISMATCH"]
    assert warning.severity.value == "error" and "5" in warning.message
    assert layout.segments[1].value_cols == ((2, 2), (3, 3), (4, 4)) and layout.segments[1].extra_cols == (5,)


def test_second_sheet_single_code_in_other_column(rules: Rules) -> None:
    first = [
        place({0: "Наименование", 3: "Код", 4: "За 9 месяцев 2025 г.", 5: "За 9 месяцев 2024 г."}),
        place({0: "Выручка", 3: "2110", 4: D("100"), 5: D("90")}),
        place({0: "Себестоимость", 3: "2120", 4: D("-60"), 5: D("-50")}),
        place({0: "Валовая прибыль", 3: "2100", 4: D("40"), 5: D("40")}),
    ]
    second = [place({0: "Чистая прибыль", 1: "2400", 2: D("7"), 3: D("6")})]
    cand = excel([("ОФР", first), ("ОФР_2", second)])
    layout = find_layout(cand, parse_header(hdr(), rules), rules)
    assert layout.segments[1].code_col == 1
    assert layout.segments[1].value_cols == ((2, 4), (3, 5))
    assert not [w for w in layout.warnings if w.code.startswith(("SEGMENT_", "HEADER_"))]


def test_header_not_found_anywhere(rules: Rules) -> None:
    rows = [["Запасы", "1210", "5", "4"], ["Дебиторы", "1230", "3", "2"], ["Итого", "1200", "8", "6"]]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    assert [w.code for w in layout.warnings] == ["HEADER_ROW_NOT_FOUND"]
    assert [(c.column_index, c.label) for c in layout.columns] == [(2, ""), (3, "")]


def test_segment_without_codes(rules: Rules) -> None:
    first = [["Наименование", "Код", "На 30.09.2025"], ["Запасы", "1210", "5"], ["Итого", "1200", "5"]]
    tail = [["Примечание к форме", "", "7"]]
    layout = find_layout(excel([("Баланс", first), ("Хвост", tail)]), parse_header(hdr(), rules), rules)
    assert [w.code for w in layout.warnings] == ["SEGMENT_WITHOUT_CODES"]
    assert layout.roles[3] is RowRole.DATA and layout.code_found


def test_first_segment_without_header_columns_from_first_headed(rules: Rules) -> None:
    """P1-18: колонки документа задаёт первый сегмент с заголовком, у первого значения — по порядку."""
    rows = [
        ["Запасы", "1210", "5", "4"],
        ["Дебиторы", "1230", "3", "2"],
        ["Итого", "1200", "8", "6"],
        ["Наименование", "Код", "На 30.09.2025", "На 31.12.2024"],
        ["Прочее", "1250", "1", "1"],
    ]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    assert [s.first_row for s in layout.segments] == [0, 3]
    assert [c.column_index for c in layout.columns] == [2, 3]
    assert layout.segments[0].value_cols == ((2, 2), (3, 3))
    assert layout.warnings == ()


def test_no_codes_anywhere_has_no_segment_warnings(rules: Rules) -> None:
    layout = find_layout(ocr([["Просто", "текст"], ["ещё", "строка"]]), parse_header(hdr(), rules), rules)
    assert not layout.code_found and layout.columns == () and layout.warnings == ()
    assert layout.roles == (RowRole.DATA, RowRole.DATA)


def _no_date_noise(layout: Layout) -> None:
    assert not [w for w in layout.warnings if w.code == "COLUMN_WITHOUT_DATE"]


def test_two_row_header_with_caption_on_top_ocr(rules: Rules) -> None:
    """I-1: годы во второй строке шапки — не строка данных, колонки с абсолютными датами."""
    rows = [
        ["Наименование", "Код", "На 30 сентября", "На 31 декабря", "На 31 декабря"],
        ["", "", "2025 г.", "2024 г.", "2023 г."],
        ["Запасы", "1210", "5", "4", "3"],
    ]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    assert layout.roles[:2] == (RowRole.HEADER, RowRole.HEADER) and layout.roles[2] is RowRole.DATA
    assert [c.period_kind.value != "none" for c in layout.columns] == [True, True, True]
    _no_date_noise(layout)


def test_two_row_header_with_caption_on_top_excel(rules: Rules) -> None:
    sheet = [
        place({0: "Наименование", 1: "Код", 2: "На 30 сентября", 3: "На 31 декабря", 4: "На 31 декабря"}),
        place({2: "2025 г.", 3: "2024 г.", 4: "2023 г."}),
        place({0: "Запасы", 1: "1210", 2: D("5"), 3: D("4"), 4: D("3")}),
    ]
    layout = find_layout(excel([("Баланс", sheet)]), parse_header(hdr(), rules), rules)
    assert layout.roles[1] is RowRole.HEADER and layout.roles[2] is RowRole.DATA
    assert [c.period_kind.value != "none" for c in layout.columns] == [True, True, True]
    _no_date_noise(layout)


def test_repeat_header_in_two_rows(rules: Rules) -> None:
    """T10.5: повторная шапка в два ряда использует ту же склейку."""
    rows = [
        ["Наименование", "Код", "На 30 сентября", "На 31 декабря"],
        ["", "", "2025 г.", "2024 г."],
        ["Запасы", "1210", "5", "4"],
        ["Наименование", "Код", "На 30 сентября", "На 31 декабря"],
        ["", "", "2025 г.", "2024 г."],
        ["Итого", "1200", "5", "4"],
    ]
    layout = find_layout(ocr(rows), parse_header(hdr(), rules), rules)
    assert [s.first_row for s in layout.segments] == [0, 3]
    assert layout.roles[4] is RowRole.HEADER and layout.roles[5] is RowRole.DATA
    _no_date_noise(layout)


def _npz_sheets(second_dates: list[str]) -> list[tuple[str, list[list[str | D | None]]]]:
    first = [
        place({1: "Наименование показателя", 12: "Код", 16: DATES3[0], 21: DATES3[1], 24: DATES3[2]}),
        place({1: "Нематериальные активы", 12: "1110", 16: D("10"), 21: D("9"), 24: D("8")}),
    ]
    second = [
        place({0: "Наименование показателя", 6: "Код", 8: second_dates[0], 9: second_dates[1], 10: second_dates[2]}),
        place({0: "Заемные средства", 6: "1510", 8: D("3"), 9: D("2"), 10: D("1")}),
    ]
    return [("Бухгалтерский баланс", first), ("Бухгалтерский баланс. Лист 4", second)]


def test_continuation_sheet_with_permuted_dates_maps_by_date_not_order(rules: Rules) -> None:
    """M-3: даты листа-продолжения переставлены — колонки ищутся по дате, а не по порядку."""
    permuted = [DATES3[2], DATES3[0], DATES3[1]]  # 31.12.2023, 30.09.2025, 31.12.2024
    cand = excel(_npz_sheets(permuted))
    layout = find_layout(cand, parse_header(hdr(report_date="30.09.2025"), rules), rules)
    assert layout.segments[1].value_cols == ((8, 24), (9, 16), (10, 21))
    assert not [w for w in layout.warnings if w.code.startswith("SEGMENT_")]


def test_continuation_sheet_with_duplicate_date_maps_by_order(rules: Rules) -> None:
    """M-3: две колонки на одну дату — сопоставление по порядку."""
    duplicate = [DATES3[0], DATES3[0], DATES3[2]]
    cand = excel(_npz_sheets(duplicate))
    layout = find_layout(cand, parse_header(hdr(report_date="30.09.2025"), rules), rules)
    assert layout.segments[1].value_cols == ((8, 16), (9, 21), (10, 24))


def test_numbers_right_of_table_in_all_segments_give_one_warning(rules: Rules) -> None:
    head: list[str | D | None] = ["Наименование", "Код", "На 30.09.2025", "На 31.12.2024", None]
    rows: list[list[str | D | None]] = [
        head,
        ["Запасы", "1210", D("5"), D("4"), D("9")],
        head,
        ["Итого", "1200", D("5"), D("4"), D("8")],
    ]
    layout = find_layout(excel([("Лист1", rows)]), parse_header(hdr(), rules), rules)
    assert [w.code for w in layout.warnings] == ["NUMBERS_OUTSIDE_TABLE"]
    assert layout.warnings[0].locator == "Лист1!E2\nЛист1!E4"
    assert layout.segments[1].value_cols == ((2, 2), (3, 3))


def test_headerless_continuation_sheet_with_outside_numbers(rules: Rules) -> None:
    first: list[list[str | D | None]] = [
        ["Наименование", "", "Код", "На 30.09.2025", "На 31.12.2024", None],
        ["Запасы", "", "1210", D("5"), D("4"), D("9")],
    ]
    second: list[list[str | D | None]] = [["Итого", "", "1200", D("5"), D("4"), D("8")]]
    cand = excel([("Лист1", first), ("Лист2", second)])
    layout = find_layout(cand, parse_header(hdr(), rules), rules)
    assert [w.code for w in layout.warnings] == ["NUMBERS_OUTSIDE_TABLE"]
    assert layout.warnings[0].locator == "Лист1!F2\nЛист2!F1"
    assert layout.segments[1].value_cols == ((3, 3), (4, 4))
    assert layout.segments[1].extra_cols == ()