#!/usr/bin/env python3
"""Заполнение Excel-шаблона результатами из Mathematica Notebook (.nb).

Пример:
    python fill_table_from_nb.py \
        --template table_template.xlsx \
        --notebooks h_174-66_Pu240_2-3_osc_Qw1.nb \
        --output table_filled.xlsx

Можно передать сразу несколько notebook-файлов. Имя каждого файла должно содержать:
    _2-3_ или _4-5_       — блок строк в таблице;
    _osc_, _pr_ или _gid_ — группу столбцов;
    _Qw1 или _Qw2         — левую или правую половину таблицы.
"""

from __future__ import annotations

import argparse
import glob
import math
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable


# Границы блоков в конкретном шаблоне.
PAIR_START_ROW = {
    "2-3": 5,
    "4-5": 19,
}

# Смещение группы относительно первого столбца данных своей половины.
MODEL_COLUMN_OFFSET = {
    "osc": 0,  # Jосц/J
    "pr": 3,   # Jпр/J
    "gid": 6,  # Jгид/J
}

# Порядок столбцов внутри каждой группы шаблона.
METHOD_COLUMN_OFFSET = {
    "heavy": 0,
    "new": 1,
    "light": 2,
}

# Смещение строк относительно строки Qb.
VALUE_ROW_OFFSET = {
    "Qb": 0,
    "Qw": 1,
    "hb": 2,
    "hw": 3,
    "L": 4,
    "Q": 5,
    "int1": 6,
    "int2": 7,
}

QB_TO_METHOD = {
    "Qb3": "heavy",
    "Qb2": "new",
    "Qb1": "light",
}

GROUP_RE = re.compile(r"Cell\[CellGroupData\[\{(.*?)\},\s*Open\s*\]\]", re.S)
INPUT_LABEL_RE = re.compile(r'CellLabel->\s*"In\[(\d+)\]:="')
ASSIGNMENT_RE = re.compile(r'RowBox\[\{"([A-Za-z][A-Za-z0-9]*)",\s*"="')
SIMPLE_OUTPUT_RE = re.compile(r'Cell\[BoxData\["([^"]+)"\],\s*"Output"', re.S)


@dataclass(frozen=True)
class NbRecord:
    order: int
    input_number: int
    expression: str
    assignments: tuple[str, ...]
    value: float | None


@dataclass(frozen=True)
class Target:
    pair: str
    model: str
    qw_side: int


@dataclass
class MethodValues:
    Qb: float | None = None
    Qw: float | None = None
    hb: float | None = None
    hw: float | None = None
    L: float | None = None
    Q: float | None = None
    int1: float | None = None
    int2: float | None = None

    def as_dict(self) -> dict[str, float | None]:
        return {
            "Qb": self.Qb,
            "Qw": self.Qw,
            "hb": self.hb,
            "hw": self.hw,
            "L": self.L,
            "Q": self.Q,
            "int1": self.int1,
            "int2": self.int2,
        }


def mathematica_number_to_float(raw: str) -> float | None:
    """Преобразует 1.23`*^-4 и подобные числа Mathematica в float."""
    value = raw.strip()
    value = re.sub(r"`(?:\d+(?:\.\d*)?)?", "", value)
    value = value.replace("*^", "e")

    special = {
        "Infinity": math.inf,
        "-Infinity": -math.inf,
        "Indeterminate": None,
    }
    if value in special:
        return special[value]

    try:
        return float(value)
    except ValueError:
        return None


def parse_records(nb_path: Path) -> list[NbRecord]:
    text = nb_path.read_text(encoding="utf-8", errors="replace")
    records: list[NbRecord] = []

    for order, group_match in enumerate(GROUP_RE.finditer(text)):
        group = group_match.group(1)
        input_label = INPUT_LABEL_RE.search(group)
        if input_label is None:
            continue

        input_end = group.find(']], "Input"')
        if input_end == -1:
            input_end = group.find('], "Input"')
        expression = group[:input_end] if input_end != -1 else group

        outputs = SIMPLE_OUTPUT_RE.findall(group)
        value = mathematica_number_to_float(outputs[-1]) if outputs else None
        assignments = tuple(ASSIGNMENT_RE.findall(expression))

        records.append(
            NbRecord(
                order=order,
                input_number=int(input_label.group(1)),
                expression=expression,
                assignments=assignments,
                value=value,
            )
        )

    if not records:
        raise ValueError(f"В файле {nb_path.name} не найдены группы Input/Output.")
    return records


def is_q_result(record: NbRecord) -> bool:
    expr = record.expression
    return (
        'FractionBox["q"' in expr
        and "SqrtBox" in expr
        and 'RowBox[{"S1", "*", "S2"}]' in expr
    )


def is_integral(record: NbRecord) -> bool:
    return 'SubsuperscriptBox["\\[Integral]"' in record.expression


def first_value_for_assignment(
    records: Iterable[NbRecord], assignment_name: str
) -> float | None:
    for record in records:
        if assignment_name in record.assignments and record.value is not None:
            return record.value
    return None


def extract_methods(records: list[NbRecord]) -> dict[str, MethodValues]:
    """Выделяет heavy/new/light без привязки к номерам In[]/Out[]."""
    q_indexes = [i for i, record in enumerate(records) if is_q_result(record)]
    if len(q_indexes) < 3:
        raise ValueError(
            "Найдено меньше трёх выражений q/Sqrt[S1*S2]. "
            "Проверьте, что notebook полностью вычислен и сохранён с Output-ячейками."
        )

    extracted: dict[str, MethodValues] = {}
    previous_q_index = -1

    for q_index in q_indexes:
        search_window = records[previous_q_index + 1 : q_index + 1]
        qb_candidates: list[tuple[int, str]] = []
        for local_index, record in enumerate(search_window):
            for qb_name in QB_TO_METHOD:
                if qb_name in record.assignments and record.value is not None:
                    qb_candidates.append((local_index, qb_name))

        if not qb_candidates:
            previous_q_index = q_index
            continue

        qb_local_index, qb_name = qb_candidates[-1]
        method = QB_TO_METHOD[qb_name]
        section = search_window[qb_local_index:]
        values = MethodValues()
        values.Qb = first_value_for_assignment(section, qb_name)
        values.hb = first_value_for_assignment(section, "hb")
        values.Q = records[q_index].value
        if values.Q is None:
            q_value = first_value_for_assignment(section, "q")
            s1_value = first_value_for_assignment(section, "S1")
            s2_value = first_value_for_assignment(section, "S2")
            if (
                q_value is not None
                and s1_value is not None
                and s2_value is not None
                and s1_value * s2_value > 0
            ):
                values.Q = q_value / math.sqrt(s1_value * s2_value)

        if method == "heavy":
            values.Qw = first_value_for_assignment(section, "Qw3")
            values.hw = first_value_for_assignment(section, "hw")
            values.L = first_value_for_assignment(section, "L")

        # Первый числовой интеграл после o1 — int1, после o2 — int2.
        pending_integral: str | None = None
        oscillator_values: dict[str, float | None] = {"int1": None, "int2": None}
        for record in section:
            if "o1" in record.assignments:
                pending_integral = "int1"
                oscillator_values["int1"] = record.value
                continue
            if "o2" in record.assignments:
                pending_integral = "int2"
                oscillator_values["int2"] = record.value
                continue
            if pending_integral and is_integral(record) and record.value is not None:
                setattr(values, pending_integral, record.value)
                pending_integral = None

        # Аналитический запасной вариант:
        # integral(0..inf) 2*x^2*o*exp(-o*x^2) dx = sqrt(pi)/(2*sqrt(o)).
        for integral_name, oscillator_value in oscillator_values.items():
            if (
                getattr(values, integral_name) is None
                and oscillator_value is not None
                and oscillator_value > 0
            ):
                setattr(
                    values,
                    integral_name,
                    math.sqrt(math.pi) / (2.0 * math.sqrt(oscillator_value)),
                )

        extracted[method] = values
        previous_q_index = q_index

    required_methods = {"heavy", "new", "light"}
    missing_methods = required_methods - extracted.keys()
    if missing_methods:
        raise ValueError(
            "Не удалось выделить методы: " + ", ".join(sorted(missing_methods))
        )

    required_by_method = {
        "heavy": ("Qb", "Qw", "hb", "hw", "L", "Q", "int1", "int2"),
        "new": ("Qb", "hb", "Q", "int1", "int2"),
        "light": ("Qb", "hb", "Q", "int1", "int2"),
    }
    for method, required_names in required_by_method.items():
        missing_values = [
            name
            for name in required_names
            if getattr(extracted[method], name) is None
        ]
        if missing_values:
            raise ValueError(
                f"Для метода {method} не найдены значения: "
                + ", ".join(missing_values)
            )

    return extracted


def extract_metadata(records: list[NbRecord]) -> dict[str, float | None]:
    return {
        name: first_value_for_assignment(records, name)
        for name in ("Z1", "Z2", "d", "b1", "b2")
    }


def detect_target(nb_path: Path) -> Target:
    name = nb_path.stem
    lower = name.lower()

    pair_candidates = re.findall(r"(?:^|_)(\d+-\d+)(?=_|$)", name)
    pair = next((item for item in pair_candidates if item in PAIR_START_ROW), None)
    if pair is None:
        raise ValueError(
            f"В имени {nb_path.name} не найден поддерживаемый блок 2-3 или 4-5."
        )

    model: str | None = None
    for candidate in MODEL_COLUMN_OFFSET:
        if re.search(rf"(?:^|[_-]){candidate}(?:[_-]|$)", lower):
            model = candidate
            break
    if model is None:
        raise ValueError(
            f"В имени {nb_path.name} не найден тип расчёта osc, pr или gid."
        )

    qw_match = re.search(r"qw[_-]?([12])(?:[_-]|$)", lower)
    if qw_match is None:
        raise ValueError(f"В имени {nb_path.name} не найдено Qw1 или Qw2.")

    return Target(pair=pair, model=model, qw_side=int(qw_match.group(1)))


def format_metadata(pair: str, metadata: dict[str, float | None]) -> str:
    def integer_or_blank(value: float | None) -> str:
        if value is None:
            return ""
        if float(value).is_integer():
            return str(int(value))
        return f"{value:.3g}"

    def compact_decimal(value: float | None) -> str:
        if value is None:
            return ""
        text = f"{value:.2f}"
        if text.startswith("0."):
            text = text[1:]
        elif text.startswith("-0."):
            text = "-" + text[2:]
        return text

    return (
        f"{pair}\n"
        f"Z1={integer_or_blank(metadata['Z1'])}\n"
        f"Z2={integer_or_blank(metadata['Z2'])}\n"
        f"d={integer_or_blank(metadata['d'])}\n"
        f"b1={compact_decimal(metadata['b1'])}\n"
        f"b2={compact_decimal(metadata['b2'])}"
    )


def write_notebook_to_sheet(
    worksheet,
    target: Target,
    methods: dict[str, MethodValues],
    metadata: dict[str, float | None],
    written_cells: dict[str, str],
    source_name: str,
) -> list[str]:
    start_row = PAIR_START_ROW[target.pair]

    # Qw1: подпись в A, данные с C. Qw2: подпись в M, данные с N.
    if target.qw_side == 1:
        metadata_column = 1
        first_data_column = 3
    else:
        metadata_column = 13
        first_data_column = 14

    worksheet.cell(row=start_row, column=metadata_column).value = format_metadata(
        target.pair, metadata
    )

    report: list[str] = []
    model_start_column = first_data_column + MODEL_COLUMN_OFFSET[target.model]

    for method_name in ("heavy", "new", "light"):
        column = model_start_column + METHOD_COLUMN_OFFSET[method_name]
        values = methods[method_name].as_dict()

        for value_name, row_offset in VALUE_ROW_OFFSET.items():
            value = values[value_name]
            if value is None:
                continue

            row = start_row + row_offset
            coordinate = worksheet.cell(row=row, column=column).coordinate
            previous_source = written_cells.get(coordinate)
            if previous_source is not None and previous_source != source_name:
                print(
                    f"Предупреждение: {coordinate} уже заполнялась из "
                    f"{previous_source}; новое значение берётся из {source_name}.",
                    file=sys.stderr,
                )

            cell = worksheet.cell(row=row, column=column)
            cell.value = value
            cell.number_format = "0.000"
            written_cells[coordinate] = source_name
            report.append(f"{coordinate}={value:.12g}")

    return report


def expand_notebook_arguments(patterns: list[str]) -> list[Path]:
    paths: list[Path] = []
    for pattern in patterns:
        matches = [Path(p) for p in glob.glob(pattern)]
        if matches:
            paths.extend(matches)
        else:
            path = Path(pattern)
            if path.exists():
                paths.append(path)
            else:
                raise FileNotFoundError(f"Не найден notebook: {pattern}")

    unique_paths: list[Path] = []
    seen: set[Path] = set()
    for path in paths:
        resolved = path.resolve()
        if resolved not in seen:
            seen.add(resolved)
            unique_paths.append(path)
    return unique_paths


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--template", required=True, type=Path, help="Excel-шаблон")
    parser.add_argument(
        "--notebooks",
        required=True,
        nargs="+",
        help="Один или несколько .nb файлов; разрешены маски *.nb",
    )
    parser.add_argument("--output", required=True, type=Path, help="Итоговый .xlsx")
    parser.add_argument("--sheet", default="Лист1", help="Имя листа шаблона")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Только разобрать notebook и показать найденные значения",
    )
    args = parser.parse_args()

    notebook_paths = expand_notebook_arguments(args.notebooks)
    parsed_items = []

    for nb_path in notebook_paths:
        records = parse_records(nb_path)
        target = detect_target(nb_path)
        methods = extract_methods(records)
        metadata = extract_metadata(records)
        parsed_items.append((nb_path, target, methods, metadata))

        print(
            f"{nb_path.name}: pair={target.pair}, model={target.model}, "
            f"Qw{target.qw_side}"
        )
        for method_name in ("heavy", "new", "light"):
            values = methods[method_name].as_dict()
            visible = ", ".join(
                f"{key}={value:.12g}"
                for key, value in values.items()
                if value is not None
            )
            print(f"  {method_name}: {visible}")

    if args.dry_run:
        return 0

    try:
        from openpyxl import load_workbook
    except ImportError as exc:
        raise SystemExit(
            "Не установлен openpyxl. Выполните: python -m pip install openpyxl"
        ) from exc

    if not args.template.exists():
        raise FileNotFoundError(f"Не найден шаблон: {args.template}")

    # Открытие через бинарный поток позволяет использовать шаблон без расширения .xlsx.
    with args.template.open("rb") as template_file:
        workbook = load_workbook(template_file)

    if args.sheet not in workbook.sheetnames:
        raise KeyError(
            f"В шаблоне нет листа {args.sheet!r}. Доступны: "
            + ", ".join(workbook.sheetnames)
        )
    worksheet = workbook[args.sheet]

    written_cells: dict[str, str] = {}
    for nb_path, target, methods, metadata in parsed_items:
        report = write_notebook_to_sheet(
            worksheet=worksheet,
            target=target,
            methods=methods,
            metadata=metadata,
            written_cells=written_cells,
            source_name=nb_path.name,
        )
        print(f"  Записано {len(report)} ячеек: " + ", ".join(report))

    args.output.parent.mkdir(parents=True, exist_ok=True)
    workbook.save(args.output)
    print(f"Готово: {args.output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())