lib/extractor/scripts/combine_results.py

#!/usr/bin/env python
"""
Скрипт для объединения результатов всех экспериментов в одну Excel-таблицу с форматированием.
Анализирует результаты экспериментов и создает сводную таблицу с метриками в различных разрезах.
Также строит графики через seaborn и сохраняет их в отдельную директорию.
"""

import argparse
import glob
import os

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
from openpyxl import Workbook
from openpyxl.styles import Alignment, Border, Font, PatternFill, Side
from openpyxl.utils import get_column_letter
from openpyxl.utils.dataframe import dataframe_to_rows


def setup_plot_directory(plots_dir: str) -> None:
    """
    Создает директорию для сохранения графиков, если она не существует.
    
    Args:
        plots_dir: Путь к директории для графиков
    """
    if not os.path.exists(plots_dir):
        os.makedirs(plots_dir)
        print(f"Создана директория для графиков: {plots_dir}")
    else:
        print(f"Директория для графиков: {plots_dir}")


def parse_args():
    """Парсит аргументы командной строки."""
    parser = argparse.ArgumentParser(description="Объединение результатов экспериментов в одну Excel-таблицу")
    
    parser.add_argument("--results-dir", type=str, default="data",
                        help="Директория с результатами экспериментов (по умолчанию: data)")
    parser.add_argument("--output-file", type=str, default="combined_results.xlsx",
                        help="Путь к выходному Excel-файлу (по умолчанию: combined_results.xlsx)")
    parser.add_argument("--plots-dir", type=str, default="plots",
                        help="Директория для сохранения графиков (по умолчанию: plots)")
    
    return parser.parse_args()


def parse_file_name(file_name: str) -> dict:
    """
    Парсит имя файла и извлекает параметры эксперимента.
    
    Args:
        file_name: Имя файла для парсинга
        
    Returns:
        Словарь с параметрами (words_per_chunk, overlap_words, model) или None при ошибке
    """
    try:
        # Извлекаем параметры из имени файла
        parts = file_name.split('_')
        if len(parts) < 4:
            return None

        # Ищем части с w (words) и o (overlap)
        words_part = None
        overlap_part = None
        
        for part in parts:
            if part.startswith('w') and part[1:].isdigit():
                words_part = part[1:]
            elif part.startswith('o') and part[1:].isdigit():
                # Убираем потенциальную часть .csv или .xlsx из overlap_part
                overlap_part = part[1:].split('.')[0]
        
        if words_part is None or overlap_part is None:
            return None
        
        # Пытаемся извлечь имя модели из оставшейся части имени файла
        model_part = file_name.split(f"_w{words_part}_o{overlap_part}_", 1)
        if len(model_part) < 2:
            return None
            
        # Получаем имя модели и удаляем возможное расширение файла
        model_name_parts = model_part[1].split('.')
        if len(model_name_parts) > 1:
            model_name_parts = model_name_parts[:-1]
        
        model_name_parts = '_'.join(model_name_parts).split('_')
        model_name = '/'.join(model_name_parts)
        
        return {
            'words_per_chunk': int(words_part),
            'overlap_words': int(overlap_part),
            'model': model_name,
            'overlap_percentage': round(int(overlap_part) / int(words_part) * 100, 1)
        }
    except Exception as e:
        print(f"Ошибка при парсинге файла {file_name}: {e}")
        return None


def load_data_files(results_dir: str, pattern: str, file_type: str, load_function) -> pd.DataFrame:
    """
    Общая функция для загрузки файлов данных с определенным паттерном имени.
    
    Args:
        results_dir: Директория с результатами
        pattern: Glob-паттерн для поиска файлов
        file_type: Тип файлов для сообщений (напр. "результатов", "метрик")
        load_function: Функция для загрузки конкретного типа файла
        
    Returns:
        DataFrame с объединенными данными или None при ошибке
    """
    print(f"Загрузка {file_type} из {results_dir}...")
    
    # Ищем все файлы с указанным паттерном
    data_files = glob.glob(os.path.join(results_dir, pattern))
    
    if not data_files:
        print(f"В директории {results_dir} не найдены файлы {file_type}")
        return None
    
    print(f"Найдено {len(data_files)} файлов {file_type}")
    
    all_data = []
    
    for file_path in data_files:
        # Извлекаем информацию о стратегии и модели из имени файла
        file_name = os.path.basename(file_path)
        print(f"Обрабатываю файл: {file_name}")
        
        # Парсим параметры из имени файла
        params = parse_file_name(file_name)
        
        if params is None:
            print(f"Пропуск файла {file_name}: не удалось извлечь параметры")
            continue
        
        words_part = params['words_per_chunk']
        overlap_part = params['overlap_words']
        model_name = params['model']
        overlap_percentage = params['overlap_percentage']
        
        print(f"  Параметры: words={words_part}, overlap={overlap_part}, model={model_name}")
        
        try:
            # Загружаем данные, используя переданную функцию
            df = load_function(file_path)
            
            # Добавляем информацию о стратегии и модели
            df['model'] = model_name
            df['words_per_chunk'] = words_part
            df['overlap_words'] = overlap_part
            df['overlap_percentage'] = overlap_percentage
            
            all_data.append(df)
        except Exception as e:
            print(f"Ошибка при обработке файла {file_path}: {e}")
    
    if not all_data:
        print(f"Не удалось загрузить ни один файл {file_type}")
        return None
    
    # Объединяем все данные
    combined_data = pd.concat(all_data, ignore_index=True)
    
    return combined_data


def load_results_files(results_dir: str) -> pd.DataFrame:
    """
    Загружает все файлы результатов из указанной директории.
    
    Args:
        results_dir: Директория с результатами
        
    Returns:
        DataFrame с объединенными результатами
    """
    # Используем общую функцию для загрузки CSV файлов
    data = load_data_files(
        results_dir, 
        "results_*.csv", 
        "результатов", 
        lambda f: pd.read_csv(f)
    )
    
    if data is None:
        raise ValueError("Не удалось загрузить файлы с результатами")
    
    return data


def load_question_metrics_files(results_dir: str) -> pd.DataFrame:
    """
    Загружает все файлы с метриками по вопросам из указанной директории.
    
    Args:
        results_dir: Директория с результатами
        
    Returns:
        DataFrame с объединенными метриками по вопросам или None, если файлов нет
    """
    # Используем общую функцию для загрузки Excel файлов
    return load_data_files(
        results_dir, 
        "question_metrics_*.xlsx", 
        "метрик по вопросам", 
        lambda f: pd.read_excel(f)
    )


def prepare_summary_by_model_top_n(df: pd.DataFrame, macro_metrics: pd.DataFrame = None) -> pd.DataFrame:
    """
    Подготавливает сводную таблицу по моделям и top_n значениям.
    Если доступны macro метрики, они также включаются в сводную таблицу.
    
    Args:
        df: DataFrame с объединенными результатами
        macro_metrics: DataFrame с macro метриками (опционально)
        
    Returns:
        DataFrame со сводной таблицей
    """
    # Определяем группировочные колонки и метрики
    group_by_columns = ['model', 'top_n']
    metrics = ['text_precision', 'text_recall', 'text_f1', 'doc_precision', 'doc_recall', 'doc_f1']
    
    # Используем общую функцию для подготовки сводки
    return prepare_summary(df, group_by_columns, metrics, macro_metrics)


def prepare_summary_by_chunking_params_top_n(df: pd.DataFrame, macro_metrics: pd.DataFrame = None) -> pd.DataFrame:
    """
    Подготавливает сводную таблицу по параметрам чанкинга и top_n значениям.
    Если доступны macro метрики, они также включаются в сводную таблицу.
    
    Args:
        df: DataFrame с объединенными результатами
        macro_metrics: DataFrame с macro метриками (опционально)
        
    Returns:
        DataFrame со сводной таблицей
    """
    # Определяем группировочные колонки и метрики
    group_by_columns = ['words_per_chunk', 'overlap_words', 'top_n']
    metrics = ['text_precision', 'text_recall', 'text_f1', 'doc_precision', 'doc_recall', 'doc_f1']
    
    # Используем общую функцию для подготовки сводки
    return prepare_summary(df, group_by_columns, metrics, macro_metrics)


def prepare_summary(df: pd.DataFrame, group_by_columns: list, metrics: list, macro_metrics: pd.DataFrame = None) -> pd.DataFrame:
    """
    Общая функция для подготовки сводной таблицы по указанным группировочным колонкам.
    Если доступны macro метрики, они также включаются в сводную таблицу.
    
    Args:
        df: DataFrame с объединенными результатами
        group_by_columns: Колонки для группировки
        metrics: Список метрик для расчета среднего
        macro_metrics: DataFrame с macro метриками (опционально)
        
    Returns:
        DataFrame со сводной таблицей
    """
    # Группируем по указанным колонкам, вычисляем средние значения метрик
    summary = df.groupby(group_by_columns).agg({
        metric: 'mean' for metric in metrics
    }).reset_index()
    
    # Если среди группировочных колонок есть 'overlap_words' и 'words_per_chunk',
    # добавляем процент перекрытия
    if 'overlap_words' in group_by_columns and 'words_per_chunk' in group_by_columns:
        summary['overlap_percentage'] = (summary['overlap_words'] / summary['words_per_chunk'] * 100).round(1)
    
    # Если доступны macro метрики, объединяем их с summary
    if macro_metrics is not None:
        # Преобразуем метрики в macro_метрики
        macro_metric_names = [f"macro_{metric}" for metric in metrics]
        
        # Группируем macro метрики по тем же колонкам
        macro_summary = macro_metrics.groupby(group_by_columns).agg({
            metric: 'mean' for metric in macro_metric_names
        }).reset_index()
        
        # Если нужно, добавляем процент перекрытия для согласованности
        if 'overlap_words' in group_by_columns and 'words_per_chunk' in group_by_columns:
            macro_summary['overlap_percentage'] = (macro_summary['overlap_words'] / macro_summary['words_per_chunk'] * 100).round(1)
            merge_on = group_by_columns + ['overlap_percentage']
        else:
            merge_on = group_by_columns
            
        # Объединяем с основной сводкой
        summary = pd.merge(summary, macro_summary, on=merge_on, how='left')
    
    # Сортируем по группировочным колонкам
    summary = summary.sort_values(group_by_columns)
    
    # Округляем метрики до 4 знаков после запятой
    for col in summary.columns:
        if any(col.endswith(suffix) for suffix in ['precision', 'recall', 'f1']):
            summary[col] = summary[col].round(4)
    
    return summary


def prepare_best_configurations(df: pd.DataFrame, macro_metrics: pd.DataFrame = None) -> pd.DataFrame:
    """
    Подготавливает таблицу с лучшими конфигурациями для каждой модели и различных top_n.
    Выбирает конфигурацию только на основе macro_text_recall и text_recall (weighted),
    игнорируя F1 метрики как менее важные.
    
    Args:
        df: DataFrame с объединенными результатами
        macro_metrics: DataFrame с macro метриками (опционально)
        
    Returns:
        DataFrame с лучшими конфигурациями
    """
    # Выбираем ключевые значения top_n
    key_top_n = [10, 20, 50, 100]
    
    # Определяем источник метрик и акцентируем только на recall-метриках
    if macro_metrics is not None:
        print("Выбор лучших конфигураций на основе macro метрик (macro_text_recall)")
        metrics_source = macro_metrics
        text_recall_metric = 'macro_text_recall'
        doc_recall_metric = 'macro_doc_recall'
    else:
        print("Выбор лучших конфигураций на основе weighted метрик (text_recall)")
        metrics_source = df
        text_recall_metric = 'text_recall'
        doc_recall_metric = 'doc_recall'
    
    # Фильтруем только по ключевым значениям top_n
    filtered_df = metrics_source[metrics_source['top_n'].isin(key_top_n)]
    
    # Для каждой модели и top_n находим конфигурацию только с лучшим recall
    best_configs = []
    
    for model in metrics_source['model'].unique():
        for top_n in key_top_n:
            model_top_n_df = filtered_df[(filtered_df['model'] == model) & (filtered_df['top_n'] == top_n)]
            
            if len(model_top_n_df) == 0:
                continue
            
            # Находим конфигурацию с лучшим text_recall
            best_text_recall_idx = model_top_n_df[text_recall_metric].idxmax()
            best_text_recall_config = model_top_n_df.loc[best_text_recall_idx].copy()
            best_text_recall_config['metric_type'] = 'text_recall'
            
            # Находим конфигурацию с лучшим doc_recall
            best_doc_recall_idx = model_top_n_df[doc_recall_metric].idxmax()
            best_doc_recall_config = model_top_n_df.loc[best_doc_recall_idx].copy()
            best_doc_recall_config['metric_type'] = 'doc_recall'
            
            best_configs.append(best_text_recall_config)
            best_configs.append(best_doc_recall_config)
    
    if not best_configs:
        return pd.DataFrame()
    
    best_configs_df = pd.DataFrame(best_configs)
    
    # Выбираем и сортируем нужные столбцы
    cols_to_keep = ['model', 'top_n', 'metric_type', 'words_per_chunk', 'overlap_words', 'overlap_percentage']
    
    # Добавляем столбцы метрик в зависимости от того, какие доступны
    if macro_metrics is not None:
        # Для macro метрик сначала выбираем recall-метрики
        recall_cols = [col for col in best_configs_df.columns if col.endswith('recall')]
        # Затем добавляем остальные метрики
        other_cols = [col for col in best_configs_df.columns if any(col.endswith(m) for m in 
                      ['precision', 'f1']) and col.startswith('macro_')]
        metric_cols = recall_cols + other_cols
    else:
        # Для weighted метрик сначала выбираем recall-метрики
        recall_cols = [col for col in best_configs_df.columns if col.endswith('recall')]
        # Затем добавляем остальные метрики
        other_cols = [col for col in best_configs_df.columns if any(col.endswith(m) for m in 
                      ['precision', 'f1']) and not col.startswith('macro_')]
        metric_cols = recall_cols + other_cols
    
    result = best_configs_df[cols_to_keep + metric_cols].sort_values(['model', 'top_n', 'metric_type'])
    
    return result


def get_grouping_columns(sheet) -> dict:
    """
    Определяет подходящие колонки для группировки данных на листе.
    
    Args:
        sheet: Лист Excel
        
    Returns:
        Словарь с данными о группировке или None
    """
    # Возможные варианты группировки
    grouping_possibilities = [
        {'columns': ['model', 'words_per_chunk', 'overlap_words']},
        {'columns': ['model']},
        {'columns': ['words_per_chunk', 'overlap_words']},
        {'columns': ['top_n']},
        {'columns': ['model', 'top_n', 'metric_type']}
    ]
    
    # Для каждого варианта группировки проверяем наличие всех колонок
    for grouping in grouping_possibilities:
        column_indices = {}
        all_columns_present = True
        
        for column_name in grouping['columns']:
            column_idx = None
            for col_idx, cell in enumerate(sheet[1], start=1):
                if cell.value == column_name:
                    column_idx = col_idx
                    break
            
            if column_idx is None:
                all_columns_present = False
                break
            else:
                column_indices[column_name] = column_idx
        
        if all_columns_present:
            return {
                'columns': grouping['columns'],
                'indices': column_indices
            }
    
    return None


def apply_header_formatting(sheet):
    """
    Применяет форматирование к заголовкам.
    
    Args:
        sheet: Лист Excel
    """
    # Форматирование заголовков
    for cell in sheet[1]:
        cell.font = Font(bold=True)
        cell.fill = PatternFill(start_color="D9D9D9", end_color="D9D9D9", fill_type="solid")
        cell.alignment = Alignment(horizontal='center', vertical='center', wrap_text=True)


def adjust_column_width(sheet):
    """
    Настраивает ширину столбцов на основе содержимого.
    
    Args:
        sheet: Лист Excel
    """
    # Авторазмер столбцов
    for column in sheet.columns:
        max_length = 0
        column_letter = get_column_letter(column[0].column)
        
        for cell in column:
            if cell.value:
                try:
                    if len(str(cell.value)) > max_length:
                        max_length = len(str(cell.value))
                except:
                    pass
        
        adjusted_width = (max_length + 2) * 1.1
        sheet.column_dimensions[column_letter].width = adjusted_width


def apply_cell_formatting(sheet):
    """
    Применяет форматирование к ячейкам (границы, выравнивание и т.д.).
    
    Args:
        sheet: Лист Excel
    """
    # Тонкие границы для всех ячеек
    thin_border = Border(
        left=Side(style='thin'),
        right=Side(style='thin'),
        top=Side(style='thin'),
        bottom=Side(style='thin')
    )
    
    for row in sheet.iter_rows(min_row=1, max_row=sheet.max_row, min_col=1, max_col=sheet.max_column):
        for cell in row:
            cell.border = thin_border
    
    # Форматирование числовых значений
    numeric_columns = [
        'text_precision', 'text_recall', 'text_f1',
        'doc_precision', 'doc_recall', 'doc_f1',
        'macro_text_precision', 'macro_text_recall', 'macro_text_f1',
        'macro_doc_precision', 'macro_doc_recall', 'macro_doc_f1'
    ]
    
    for col_idx, header in enumerate(sheet[1], start=1):
        if header.value in numeric_columns or (header.value and str(header.value).endswith(('precision', 'recall', 'f1'))):
            for row_idx in range(2, sheet.max_row + 1):
                cell = sheet.cell(row=row_idx, column=col_idx)
                if isinstance(cell.value, (int, float)):
                    cell.number_format = '0.0000'
    
    # Выравнивание для всех ячеек
    for row in sheet.iter_rows(min_row=2, max_row=sheet.max_row, min_col=1, max_col=sheet.max_column):
        for cell in row:
            cell.alignment = Alignment(horizontal='center', vertical='center')


def apply_group_formatting(sheet, grouping):
    """
    Применяет форматирование к группам строк.
    
    Args:
        sheet: Лист Excel
        grouping: Словарь с данными о группировке
    """
    if not grouping or sheet.max_row <= 1:
        return
    
    # Для каждой строки проверяем изменение значений группировочных колонок
    last_values = {column: None for column in grouping['columns']}
    
    # Применяем жирную верхнюю границу к первой строке данных
    for col_idx in range(1, sheet.max_column + 1):
        cell = sheet.cell(row=2, column=col_idx)
        cell.border = Border(
            left=cell.border.left,
            right=cell.border.right,
            top=Side(style='thick'),
            bottom=cell.border.bottom
        )
    
    for row_idx in range(2, sheet.max_row + 1):
        current_values = {}
        for column in grouping['columns']:
            col_idx = grouping['indices'][column]
            current_values[column] = sheet.cell(row=row_idx, column=col_idx).value
        
        # Если значения изменились, добавляем жирные границы
        values_changed = False
        for column in grouping['columns']:
            if current_values[column] != last_values[column]:
                values_changed = True
                break
        
        if values_changed and row_idx > 2:
            # Жирная верхняя граница для текущей строки
            for col_idx in range(1, sheet.max_column + 1):
                cell = sheet.cell(row=row_idx, column=col_idx)
                cell.border = Border(
                    left=cell.border.left,
                    right=cell.border.right,
                    top=Side(style='thick'),
                    bottom=cell.border.bottom
                )
            
            # Жирная нижняя граница для предыдущей строки
            for col_idx in range(1, sheet.max_column + 1):
                cell = sheet.cell(row=row_idx-1, column=col_idx)
                cell.border = Border(
                    left=cell.border.left,
                    right=cell.border.right,
                    top=cell.border.top,
                    bottom=Side(style='thick')
                )
        
        # Запоминаем текущие значения для следующей итерации
        for column in grouping['columns']:
            last_values[column] = current_values[column]
    
    # Добавляем жирную нижнюю границу для последней строки
    for col_idx in range(1, sheet.max_column + 1):
        cell = sheet.cell(row=sheet.max_row, column=col_idx)
        cell.border = Border(
            left=cell.border.left,
            right=cell.border.right,
            top=cell.border.top,
            bottom=Side(style='thick')
        )


def apply_formatting(workbook: Workbook) -> None:
    """
    Применяет форматирование к Excel-файлу.
    Добавляет автофильтры для всех столбцов и улучшает визуальное представление.
    
    Args:
        workbook: Workbook-объект openpyxl
    """
    for sheet_name in workbook.sheetnames:
        sheet = workbook[sheet_name]
        
        # Добавляем автофильтры для всех столбцов
        if sheet.max_row > 1:  # Проверяем, что в листе есть данные
            sheet.auto_filter.ref = sheet.dimensions
        
        # Применяем форматирование
        apply_header_formatting(sheet)
        adjust_column_width(sheet)
        apply_cell_formatting(sheet)
        
        # Определяем группирующие колонки и применяем форматирование к группам
        grouping = get_grouping_columns(sheet)
        if grouping:
            apply_group_formatting(sheet, grouping)


def create_model_comparison_plot(df: pd.DataFrame, metrics: list | str, top_n: int, plots_dir: str) -> None:
    """
    Создает график сравнения моделей по указанным метрикам для заданного top_n.
    
    Args:
        df: DataFrame с данными
        metrics: Список метрик или одна метрика для сравнения
        top_n: Значение top_n для фильтрации
        plots_dir: Директория для сохранения графиков
    """
    if isinstance(metrics, str):
        metrics = [metrics]
        
    # Фильтруем данные
    filtered_df = df[df['top_n'] == top_n]
    
    if len(filtered_df) == 0:
        print(f"Нет данных для top_n={top_n}")
        return
    
    # Определяем тип метрик (macro или weighted)
    metrics_type = "macro" if metrics[0].startswith("macro_") else "weighted"
    
    # Создаем фигуру с несколькими подграфиками
    fig, axes = plt.subplots(1, len(metrics), figsize=(6 * len(metrics), 8))
    
    # Если только одна метрика, преобразуем axes в список для единообразного обращения
    if len(metrics) == 1:
        axes = [axes]
    
    # Для каждой метрики создаем subplot
    for i, metric in enumerate(metrics):
        # Группируем данные по модели
        columns_to_agg = {metric: 'mean'}
        model_data = filtered_df.groupby('model').agg(columns_to_agg).reset_index()
        
        # Сортируем по значению метрики (по убыванию)
        model_data = model_data.sort_values(metric, ascending=False)
        
        # Определяем цветовую схему
        palette = sns.color_palette("viridis", len(model_data))
        
        # Строим столбчатую диаграмму на соответствующем subplot
        ax = sns.barplot(x='model', y=metric, data=model_data, palette=palette, ax=axes[i])
        
        # Добавляем значения над столбцами
        for j, v in enumerate(model_data[metric]):
            ax.text(j, v + 0.01, f"{v:.4f}", ha='center', fontsize=8)
        
        # Устанавливаем заголовок и метки осей
        ax.set_title(f"{metric} (top_n={top_n})", fontsize=12)
        ax.set_xlabel("Модель", fontsize=10)
        ax.set_ylabel(f"{metric}", fontsize=10)
        
        # Поворачиваем подписи по оси X для лучшей читаемости
        ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha='right', fontsize=8)
    
    # Настраиваем макет
    plt.tight_layout()
    
    # Сохраняем график
    metric_names = '_'.join([m.replace('macro_', '') for m in metrics])
    file_name = f"model_comparison_{metrics_type}_{metric_names}_top{top_n}.png"
    plt.savefig(os.path.join(plots_dir, file_name), dpi=300)
    plt.close()
    
    print(f"Создан график сравнения моделей: {file_name}")


def create_top_n_plot(df: pd.DataFrame, models: list | str, metric: str, plots_dir: str) -> None:
    """
    Создает график зависимости метрики от top_n для заданных моделей.
    
    Args:
        df: DataFrame с данными
        models: Список моделей или одна модель для сравнения
        metric: Название метрики
        plots_dir: Директория для сохранения графиков
    """
    if isinstance(models, str):
        models = [models]
        
    # Создаем фигуру
    plt.figure(figsize=(12, 8))
    
    # Определяем цветовую схему
    palette = sns.color_palette("viridis", len(models))
    
    # Ограничиваем количество моделей для читаемости
    if len(models) > 5:
        models = models[:5]
        print("Слишком много моделей для графика, ограничиваем до 5")
    
    # Для каждой модели строим линию
    for i, model in enumerate(models):
        # Находим наиболее часто используемые параметры чанкинга для этой модели
        model_df = df[df['model'] == model]
        
        if len(model_df) == 0:
            print(f"Нет данных для модели {model}")
            continue
        
        # Группируем по параметрам чанкинга и подсчитываем частоту
        common_configs = model_df.groupby(['words_per_chunk', 'overlap_words']).size().reset_index(name='count')
        
        if len(common_configs) == 0:
            continue
            
        # Берем наиболее частую конфигурацию
        common_config = common_configs.sort_values('count', ascending=False).iloc[0]
        
        # Фильтруем для этой конфигурации
        config_df = model_df[
            (model_df['words_per_chunk'] == common_config['words_per_chunk']) &
            (model_df['overlap_words'] == common_config['overlap_words'])
        ].sort_values('top_n')
        
        if len(config_df) <= 1:
            continue
        
        # Строим линию
        plt.plot(config_df['top_n'], config_df[metric], marker='o', linewidth=2, 
                 label=f"{model} (w={common_config['words_per_chunk']}, o={common_config['overlap_words']})",
                 color=palette[i])
    
    # Добавляем легенду, заголовок и метки осей
    plt.legend(title="Модель (параметры)", fontsize=10, loc='best')
    plt.title(f"Зависимость {metric} от top_n для разных моделей", fontsize=16)
    plt.xlabel("top_n", fontsize=14)
    plt.ylabel(metric, fontsize=14)
    
    # Включаем сетку
    plt.grid(True, linestyle='--', alpha=0.7)
    
    # Настраиваем макет
    plt.tight_layout()
    
    # Сохраняем график
    is_macro = "macro" if "macro" in metric else "weighted"
    file_name = f"top_n_comparison_{is_macro}_{metric.replace('macro_', '')}.png"
    plt.savefig(os.path.join(plots_dir, file_name), dpi=300)
    plt.close()
    
    print(f"Создан график зависимости от top_n: {file_name}")


def create_chunk_size_plot(df: pd.DataFrame, model: str, metrics: list | str, top_n: int, plots_dir: str) -> None:
    """
    Создает график зависимости метрик от размера чанка для заданной модели и top_n.
    
    Args:
        df: DataFrame с данными
        model: Название модели
        metrics: Список метрик или одна метрика
        top_n: Значение top_n
        plots_dir: Директория для сохранения графиков
    """
    if isinstance(metrics, str):
        metrics = [metrics]
        
    # Фильтруем данные
    filtered_df = df[(df['model'] == model) & (df['top_n'] == top_n)]
    
    if len(filtered_df) <= 1:
        print(f"Недостаточно данных для модели {model} и top_n={top_n}")
        return
    
    # Создаем фигуру
    plt.figure(figsize=(14, 8))
    
    # Определяем цветовую схему для метрик
    palette = sns.color_palette("viridis", len(metrics))
    
    # Группируем по размеру чанка и проценту перекрытия
    # Вычисляем среднее только для указанных метрик, а не для всех столбцов
    columns_to_agg = {metric: 'mean' for metric in metrics}
    chunk_data = filtered_df.groupby(['words_per_chunk', 'overlap_percentage']).agg(columns_to_agg).reset_index()
    
    # Получаем уникальные значения процента перекрытия
    overlap_percentages = sorted(chunk_data['overlap_percentage'].unique())
    
    # Настраиваем маркеры и линии для разных перекрытий
    markers = ['o', 's', '^', 'D', 'x', '*']
    
    # Для каждого перекрытия строим линии с разными метриками
    for i, overlap in enumerate(overlap_percentages):
        subset = chunk_data[chunk_data['overlap_percentage'] == overlap].sort_values('words_per_chunk')
        
        for j, metric in enumerate(metrics):
            plt.plot(subset['words_per_chunk'], subset[metric], 
                     marker=markers[i % len(markers)], linewidth=2, 
                     label=f"{metric}, overlap={overlap}%",
                     color=palette[j])
    
    # Добавляем легенду и заголовок
    plt.legend(title="Метрика и перекрытие", fontsize=10, loc='best')
    plt.title(f"Зависимость метрик от размера чанка для {model} (top_n={top_n})", fontsize=16)
    plt.xlabel("Размер чанка (слов)", fontsize=14)
    plt.ylabel("Значение метрики", fontsize=14)
    
    # Включаем сетку
    plt.grid(True, linestyle='--', alpha=0.7)
    
    # Настраиваем макет
    plt.tight_layout()
    
    # Сохраняем график
    metrics_type = "macro" if metrics[0].startswith("macro_") else "weighted"
    model_name = model.replace('/', '_')
    metric_names = '_'.join([m.replace('macro_', '') for m in metrics])
    file_name = f"chunk_size_{metrics_type}_{metric_names}_{model_name}_top{top_n}.png"
    plt.savefig(os.path.join(plots_dir, file_name), dpi=300)
    plt.close()
    
    print(f"Создан график зависимости от размера чанка: {file_name}")


def create_heatmap(df: pd.DataFrame, models: list | str, metric: str, top_n: int, plots_dir: str) -> None:
    """
    Создает тепловые карты зависимости метрики от размера чанка и процента перекрытия
    для заданных моделей.
    
    Args:
        df: DataFrame с данными
        models: Список моделей или одна модель
        metric: Название метрики
        top_n: Значение top_n
        plots_dir: Директория для сохранения графиков
    """
    if isinstance(models, str):
        models = [models]
        
    # Ограничиваем количество моделей для наглядности
    if len(models) > 4:
        models = models[:4]
    
    # Создаем фигуру с подграфиками
    fig, axes = plt.subplots(1, len(models), figsize=(6 * len(models), 6), squeeze=False)
    
    # Для каждой модели создаем тепловую карту
    for i, model in enumerate(models):
        # Фильтруем данные для указанной модели и top_n
        filtered_df = df[(df['model'] == model) & (df['top_n'] == top_n)]
        
        # Проверяем, достаточно ли данных для построения тепловой карты
        chunk_sizes = filtered_df['words_per_chunk'].unique()
        overlap_percentages = filtered_df['overlap_percentage'].unique()
        
        if len(chunk_sizes) <= 1 or len(overlap_percentages) <= 1:
            print(f"Недостаточно данных для построения тепловой карты для модели {model} и top_n={top_n}")
            # Пропускаем этот subplot
            axes[0, i].text(0.5, 0.5, f"Недостаточно данных для {model}", 
                           horizontalalignment='center', verticalalignment='center')
            axes[0, i].set_title(model)
            axes[0, i].axis('off')
            continue
        
        # Создаем сводную таблицу для тепловой карты, используя только нужную метрику
        # Сначала выберем только колонки для pivot_table
        pivot_columns = ['words_per_chunk', 'overlap_percentage', metric]
        pivot_df = filtered_df[pivot_columns].copy()
        
        # Теперь создаем сводную таблицу
        pivot_data = pivot_df.pivot_table(
            index='words_per_chunk',
            columns='overlap_percentage',
            values=metric,
            aggfunc='mean'
        )
        
        # Строим тепловую карту
        sns.heatmap(pivot_data, annot=True, fmt=".4f", cmap="viridis",
                    linewidths=.5, annot_kws={"size": 8}, ax=axes[0, i])
        
        # Устанавливаем заголовок и метки осей
        axes[0, i].set_title(model, fontsize=12)
        axes[0, i].set_xlabel("Процент перекрытия (%)", fontsize=10)
        axes[0, i].set_ylabel("Размер чанка (слов)", fontsize=10)
    
    # Добавляем общий заголовок
    plt.suptitle(f"Тепловые карты {metric} для разных моделей (top_n={top_n})", fontsize=16)
    
    # Настраиваем макет
    plt.tight_layout(rect=[0, 0, 1, 0.96])  # Оставляем место для общего заголовка
    
    # Сохраняем график
    is_macro = "macro" if "macro" in metric else "weighted"
    file_name = f"heatmap_{is_macro}_{metric.replace('macro_', '')}_top{top_n}.png"
    plt.savefig(os.path.join(plots_dir, file_name), dpi=300)
    plt.close()
    
    print(f"Созданы тепловые карты: {file_name}")


def find_best_combinations(df: pd.DataFrame, metrics: list | str = None) -> pd.DataFrame:
    """
    Находит наилучшие комбинации параметров на основе агрегированных recall-метрик.
    
    Args:
        df: DataFrame с данными
        metrics: Список метрик для анализа или None (тогда используются все recall-метрики)
        
    Returns:
        DataFrame с лучшими комбинациями параметров
    """
    if metrics is None:
        # По умолчанию выбираем все метрики с "recall" в названии
        metrics = [col for col in df.columns if "recall" in col]
    elif isinstance(metrics, str):
        metrics = [metrics]
    
    print(f"Поиск лучших комбинаций на основе метрик: {metrics}")
    
    # Создаем новую метрику - сумму всех указанных recall-метрик
    df_copy = df.copy()
    df_copy['combined_recall'] = df_copy[metrics].sum(axis=1)
    
    # Находим лучшие комбинации для различных значений top_n
    best_combinations = []
    
    for top_n in df_copy['top_n'].unique():
        top_n_df = df_copy[df_copy['top_n'] == top_n]
        
        if len(top_n_df) == 0:
            continue
        
        # Находим строку с максимальным combined_recall
        best_idx = top_n_df['combined_recall'].idxmax()
        best_row = top_n_df.loc[best_idx].copy()
        best_row['best_for_top_n'] = top_n
        
        best_combinations.append(best_row)
    
    # Находим лучшие комбинации для разных моделей
    for model in df_copy['model'].unique():
        model_df = df_copy[df_copy['model'] == model]
        
        if len(model_df) == 0:
            continue
        
        # Находим строку с максимальным combined_recall
        best_idx = model_df['combined_recall'].idxmax()
        best_row = model_df.loc[best_idx].copy()
        best_row['best_for_model'] = model
        
        best_combinations.append(best_row)
    
    # Находим лучшие комбинации для разных размеров чанков
    for chunk_size in df_copy['words_per_chunk'].unique():
        chunk_df = df_copy[df_copy['words_per_chunk'] == chunk_size]
        
        if len(chunk_df) == 0:
            continue
        
        # Находим строку с максимальным combined_recall
        best_idx = chunk_df['combined_recall'].idxmax()
        best_row = chunk_df.loc[best_idx].copy()
        best_row['best_for_chunk_size'] = chunk_size
        
        best_combinations.append(best_row)
    
    # Находим абсолютно лучшую комбинацию
    if len(df_copy) > 0:
        best_idx = df_copy['combined_recall'].idxmax()
        best_row = df_copy.loc[best_idx].copy()
        best_row['absolute_best'] = True
        
        best_combinations.append(best_row)
    
    if not best_combinations:
        return pd.DataFrame()
    
    result = pd.DataFrame(best_combinations)
    
    # Сортируем по combined_recall (по убыванию)
    result = result.sort_values('combined_recall', ascending=False)
    
    print(f"Найдено {len(result)} лучших комбинаций")
    
    return result


def create_best_combinations_plot(best_df: pd.DataFrame, metrics: list | str, plots_dir: str) -> None:
    """
    Создает график сравнения лучших комбинаций параметров.
    
    Args:
        best_df: DataFrame с лучшими комбинациями
        metrics: Список метрик для визуализации
        plots_dir: Директория для сохранения графиков
    """
    if isinstance(metrics, str):
        metrics = [metrics]
    
    if len(best_df) == 0:
        print("Нет данных для построения графика лучших комбинаций")
        return
    
    # Создаем новый признак для идентификации комбинаций
    best_df['combo_label'] = best_df.apply(
        lambda row: f"{row['model']} (w={row['words_per_chunk']}, o={row['overlap_words']}, top_n={row['top_n']})",
        axis=1
    )
    
    # Берем только лучшие N комбинаций для читаемости
    max_combos = 10
    if len(best_df) > max_combos:
        plot_df = best_df.head(max_combos).copy()
        print(f"Ограничиваем график до {max_combos} лучших комбинаций")
    else:
        plot_df = best_df.copy()
    
    # Создаем длинный формат данных для seaborn
    plot_data = plot_df.melt(
        id_vars=['combo_label', 'combined_recall'],
        value_vars=metrics,
        var_name='metric',
        value_name='value'
    )
    
    # Сортируем по суммарному recall (комбинации) и метрике (для группировки)
    plot_data = plot_data.sort_values(['combined_recall', 'metric'], ascending=[False, True])
    
    # Создаем фигуру для графика
    plt.figure(figsize=(14, 10))
    
    # Создаем bar plot
    sns.barplot(
        x='combo_label',
        y='value',
        hue='metric',
        data=plot_data,
        palette='viridis'
    )
    
    # Настраиваем оси и заголовок
    plt.title('Лучшие комбинации параметров по recall-метрикам', fontsize=16)
    plt.xlabel('Комбинация параметров', fontsize=14)
    plt.ylabel('Значение метрики', fontsize=14)
    
    # Поворачиваем подписи по оси X для лучшей читаемости
    plt.xticks(rotation=45, ha='right', fontsize=10)
    
    # Настраиваем легенду
    plt.legend(title='Метрика', fontsize=12)
    
    # Добавляем сетку
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    
    # Настраиваем макет
    plt.tight_layout()
    
    # Сохраняем график
    file_name = f"best_combinations_comparison.png"
    plt.savefig(os.path.join(plots_dir, file_name), dpi=300)
    plt.close()
    
    print(f"Создан график сравнения лучших комбинаций: {file_name}")


def generate_plots(combined_results: pd.DataFrame, macro_metrics: pd.DataFrame, plots_dir: str) -> None:
    """
    Генерирует набор графиков с помощью seaborn и сохраняет их в указанную директорию.
    Фокусируется в первую очередь на recall-метриках как наиболее важных.
    
    Args:
        combined_results: DataFrame с объединенными результатами (weighted метрики)
        macro_metrics: DataFrame с macro метриками
        plots_dir: Директория для сохранения графиков
    """
    # Создаем директорию для графиков, если она не существует
    setup_plot_directory(plots_dir)
    
    # Настраиваем стиль для графиков
    sns.set_style("whitegrid")
    plt.rcParams['font.family'] = 'DejaVu Sans'
    
    # Получаем список моделей для построения графиков
    models = combined_results['model'].unique()
    top_n_values = [10, 20, 50, 100]
    
    print(f"Генерация графиков для {len(models)} моделей...")
    
    # 0. Добавляем анализ наилучших комбинаций параметров
    # Определяем метрики для анализа - фокусируемся на recall
    weighted_recall_metrics = ['text_recall', 'doc_recall']
    
    # Находим лучшие комбинации параметров
    best_combinations = find_best_combinations(combined_results, weighted_recall_metrics)
    
    # Создаем график сравнения лучших комбинаций
    if not best_combinations.empty:
        create_best_combinations_plot(best_combinations, weighted_recall_metrics, plots_dir)
    
    # Если доступны macro метрики, делаем то же самое для них
    if macro_metrics is not None:
        macro_recall_metrics = ['macro_text_recall', 'macro_doc_recall']
        macro_best_combinations = find_best_combinations(macro_metrics, macro_recall_metrics)
        
        if not macro_best_combinations.empty:
            create_best_combinations_plot(macro_best_combinations, macro_recall_metrics, plots_dir)
    
    # 1. Создаем графики сравнения моделей для weighted метрик
    # Фокусируемся на recall-метриках
    weighted_metrics = {
        'text': ['text_recall'],  # Только text_recall
        'doc': ['doc_recall']     # Только doc_recall
    }
    
    for top_n in top_n_values:
        for metrics_group, metrics in weighted_metrics.items():
            create_model_comparison_plot(combined_results, metrics, top_n, plots_dir)
    
    # 2. Если доступны macro метрики, создаем графики на их основе
    if macro_metrics is not None:
        print("Создание графиков на основе macro метрик...")
        macro_metrics_groups = {
            'text': ['macro_text_recall'],  # Только macro_text_recall
            'doc': ['macro_doc_recall']     # Только macro_doc_recall
        }
        
        for top_n in top_n_values:
            for metrics_group, metrics in macro_metrics_groups.items():
                create_model_comparison_plot(macro_metrics, metrics, top_n, plots_dir)
    
    # 3. Создаем графики зависимости от top_n
    for metrics_type, df in [("weighted", combined_results), ("macro", macro_metrics)]:
        if df is None:
            continue
            
        metrics_to_plot = []
        if metrics_type == "weighted":
            metrics_to_plot = ['text_recall', 'doc_recall']  # Только recall-метрики
        else:
            metrics_to_plot = ['macro_text_recall', 'macro_doc_recall']  # Только macro recall-метрики
        
        for metric in metrics_to_plot:
            create_top_n_plot(df, models, metric, plots_dir)
    
    # 4. Для каждой модели создаем графики по размеру чанка
    for model in models:
        # Выбираем 2 значения top_n для анализа
        for top_n in [20, 50]:
            # Создаем графики с recall-метриками
            weighted_metrics_to_combine = ['text_recall']
            create_chunk_size_plot(combined_results, model, weighted_metrics_to_combine, top_n, plots_dir)
            
            doc_metrics_to_combine = ['doc_recall']
            create_chunk_size_plot(combined_results, model, doc_metrics_to_combine, top_n, plots_dir)
            
            # Если есть macro метрики, создаем соответствующие графики
            if macro_metrics is not None:
                macro_metrics_to_combine = ['macro_text_recall']
                create_chunk_size_plot(macro_metrics, model, macro_metrics_to_combine, top_n, plots_dir)
                
                macro_doc_metrics_to_combine = ['macro_doc_recall']
                create_chunk_size_plot(macro_metrics, model, macro_doc_metrics_to_combine, top_n, plots_dir)
    
    # 5. Создаем тепловые карты для моделей
    for top_n in [20, 50]:
        for metric_prefix in ["", "macro_"]:
            for metric_type in ["text_recall", "doc_recall"]:
                metric = f"{metric_prefix}{metric_type}"
                # Используем соответствующий DataFrame
                if metric_prefix and macro_metrics is None:
                    continue
                df_to_use = macro_metrics if metric_prefix else combined_results
                create_heatmap(df_to_use, models, metric, top_n, plots_dir)
    
    print(f"Создание графиков завершено в директории {plots_dir}")


def print_best_combinations(best_df: pd.DataFrame) -> None:
    """
    Выводит информацию о лучших комбинациях параметров.
    
    Args:
        best_df: DataFrame с лучшими комбинациями
    """
    if best_df.empty:
        print("Не найдено лучших комбинаций")
        return
    
    print("\n=== ЛУЧШИЕ КОМБИНАЦИИ ПАРАМЕТРОВ ===")
    
    # Выводим абсолютно лучшую комбинацию, если она есть
    absolute_best = best_df[best_df.get('absolute_best', False) == True]
    if not absolute_best.empty:
        row = absolute_best.iloc[0]
        print(f"\nАБСОЛЮТНО ЛУЧШАЯ КОМБИНАЦИЯ:")
        print(f"  Модель: {row['model']}")
        print(f"  Размер чанка: {row['words_per_chunk']} слов")
        print(f"  Перекрытие: {row['overlap_words']} слов ({row['overlap_percentage']}%)")
        print(f"  top_n: {row['top_n']}")
        
        # Выводим значения метрик
        recall_metrics = [col for col in best_df.columns if 'recall' in col and col != 'combined_recall']
        for metric in recall_metrics:
            print(f"  {metric}: {row[metric]:.4f}")
    
    print("\n=== ТОП-5 ЛУЧШИХ КОМБИНАЦИЙ ===")
    for i, row in best_df.head(5).iterrows():
        print(f"\n#{i+1}: {row['model']}, w={row['words_per_chunk']}, o={row['overlap_words']}, top_n={row['top_n']}")
        
        # Выводим значения метрик
        recall_metrics = [col for col in best_df.columns if 'recall' in col and col != 'combined_recall']
        for metric in recall_metrics:
            print(f"  {metric}: {row[metric]:.4f}")
    
    print("\n=======================================")


def create_combined_excel(combined_results: pd.DataFrame, question_metrics: pd.DataFrame, 
                       macro_metrics: pd.DataFrame = None, output_file: str = "combined_results.xlsx") -> None:
    """
    Создает Excel-файл с несколькими листами, содержащими различные срезы данных.
    Добавляет автофильтры и применяет форматирование.
    
    Args:
        combined_results: DataFrame с объединенными результатами
        question_metrics: DataFrame с метриками по вопросам
        macro_metrics: DataFrame с macro метриками (опционально)
        output_file: Путь к выходному Excel-файлу
    """
    print(f"Создание Excel-файла {output_file}...")
    
    # Создаем новый Excel-файл
    workbook = Workbook()
    
    # Удаляем стандартный лист
    default_sheet = workbook.active
    workbook.remove(default_sheet)
    
    # Подготавливаем данные для различных листов
    sheets_data = {
        "Исходные данные": combined_results,
        "Сводка по моделям": prepare_summary_by_model_top_n(combined_results, macro_metrics),
        "Сводка по чанкингу": prepare_summary_by_chunking_params_top_n(combined_results, macro_metrics),
        "Лучшие конфигурации": prepare_best_configurations(combined_results, macro_metrics)
    }
    
    # Если есть метрики по вопросам, добавляем лист с ними
    if question_metrics is not None:
        sheets_data["Метрики по вопросам"] = question_metrics
    
    # Если есть macro метрики, добавляем лист с ними
    if macro_metrics is not None:
        sheets_data["Macro метрики"] = macro_metrics
    
    # Создаем листы и добавляем данные
    for sheet_name, data in sheets_data.items():
        if data is not None and not data.empty:
            sheet = workbook.create_sheet(title=sheet_name)
            for r in dataframe_to_rows(data, index=False, header=True):
                sheet.append(r)
    
    # Применяем форматирование
    apply_formatting(workbook)
    
    # Сохраняем файл
    workbook.save(output_file)
    print(f"Excel-файл создан: {output_file}")


def calculate_macro_metrics(question_metrics: pd.DataFrame) -> pd.DataFrame:
    """
    Вычисляет macro метрики на основе результатов по вопросам.
    
    Args:
        question_metrics: DataFrame с метриками по вопросам
        
    Returns:
        DataFrame с macro метриками
    """
    if question_metrics is None:
        return None
    
    print("Вычисление macro метрик на основе метрик по вопросам...")
    
    # Группируем по конфигурации (модель, параметры чанкинга, top_n)
    grouped_metrics = question_metrics.groupby(['model', 'words_per_chunk', 'overlap_words', 'top_n'])
    
    # Для каждой группы вычисляем среднее значение метрик (macro)
    macro_metrics = grouped_metrics.agg({
        'text_precision': 'mean',  # Macro precision = среднее precision по всем вопросам
        'text_recall': 'mean',     # Macro recall = среднее recall по всем вопросам
        'text_f1': 'mean',         # Macro F1 = среднее F1 по всем вопросам
        'doc_precision': 'mean',
        'doc_recall': 'mean',
        'doc_f1': 'mean'
    }).reset_index()
    
    # Добавляем префикс "macro_" к названиям метрик для ясности
    for col in ['text_precision', 'text_recall', 'text_f1', 'doc_precision', 'doc_recall', 'doc_f1']:
        macro_metrics.rename(columns={col: f'macro_{col}'}, inplace=True)
    
    # Добавляем процент перекрытия
    macro_metrics['overlap_percentage'] = (macro_metrics['overlap_words'] / macro_metrics['words_per_chunk'] * 100).round(1)
    
    print(f"Вычислено {len(macro_metrics)} наборов macro метрик")
    
    return macro_metrics


def main():
    """Основная функция скрипта."""
    args = parse_args()
    
    # Загружаем результаты из CSV-файлов
    combined_results = load_results_files(args.results_dir)
    
    # Загружаем метрики по вопросам (если есть)
    question_metrics = load_question_metrics_files(args.results_dir)
    
    # Вычисляем macro метрики на основе метрик по вопросам
    macro_metrics = calculate_macro_metrics(question_metrics)
    
    # Находим лучшие комбинации параметров
    best_combinations_weighted = find_best_combinations(combined_results, ['text_recall', 'doc_recall'])
    print_best_combinations(best_combinations_weighted)
    
    if macro_metrics is not None:
        best_combinations_macro = find_best_combinations(macro_metrics, ['macro_text_recall', 'macro_doc_recall'])
        print_best_combinations(best_combinations_macro)
    
    # Создаем объединенный Excel-файл с данными
    create_combined_excel(combined_results, question_metrics, macro_metrics, args.output_file)
    
    # Генерируем графики с помощью seaborn
    print(f"Генерация графиков и сохранение их в директорию: {args.plots_dir}")
    generate_plots(combined_results, macro_metrics, args.plots_dir)
    
    print("Готово! Результаты сохранены в Excel и графики созданы.")


if __name__ == "__main__":
    main()