scripts/testing/plot_results.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Скрипт для визуализации агрегированных результатов тестирования RAG.

Читает данные из Excel-файла, сгенерированного aggregate_results.py,
и строит различные графики для анализа влияния параметров на метрики.
"""

import argparse
import json
import os

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

# --- Настройки --- 
DEFAULT_RESULTS_FILE = "data/output/aggregated_results.xlsx" # Файл с агрегированными данными
DEFAULT_PLOTS_DIR = "data/output/plots" # Куда сохранять графики

# Настройки графиков
plt.rcParams['font.family'] = 'DejaVu Sans' # Шрифт с поддержкой кириллицы
sns.set_style("whitegrid")
FIGSIZE = (16, 10) # Увеличенный размер для сложных графиков
DPI = 300
PALETTE = "viridis" # Цветовая палитра

# --- Маппинг названий столбцов (копия из aggregate_results.py) ---
COLUMN_NAME_MAPPING = {
    # Параметры запуска из pipeline.py
    'run_id': 'ID Запуска',
    'model_name': 'Модель',
    'chunking_strategy': 'Стратегия Чанкинга',
    'strategy_params': 'Параметры Стратегии',
    'process_tables': 'Обраб. Таблиц',
    'top_n': 'Top N',
    'use_injection': 'Сборка Контекста',
    'use_qe': 'Query Expansion',
    'neighbors_included': 'Вкл. Соседей',
    'similarity_threshold': 'Порог Схожести',

    # Идентификаторы из датасета (для детальных результатов)
    'question_id': 'ID Вопроса',
    'question_text': 'Текст Вопроса',

    # Детальные метрики из pipeline.py
    'chunk_text_precision': 'Точность (Чанк-Текст)',
    'chunk_text_recall': 'Полнота (Чанк-Текст)',
    'chunk_text_f1': 'F1 (Чанк-Текст)',
    'found_puncts': 'Найдено Пунктов',
    'total_puncts': 'Всего Пунктов',
    'relevant_chunks': 'Релевантных Чанков',
    'total_chunks_in_top_n': 'Всего Чанков в Топ-N',
    'assembly_punct_recall': 'Полнота (Сборка-Пункт)',
    'assembled_context_preview': 'Предпросмотр Сборки',
    # 'top_chunk_ids': 'Индексы Топ-Чанков', # Списки, могут плохо отображаться
    # 'top_chunk_similarities': 'Схожести Топ-Чанков', # Списки

    # Агрегированные метрики (добавляются в calculate_aggregated_metrics)
    'weighted_chunk_text_precision': 'Weighted Точность (Чанк-Текст)',
    'weighted_chunk_text_recall': 'Weighted Полнота (Чанк-Текст)',
    'weighted_chunk_text_f1': 'Weighted F1 (Чанк-Текст)',
    'weighted_assembly_punct_recall': 'Weighted Полнота (Сборка-Пункт)',

    'macro_chunk_text_precision': 'Macro Точность (Чанк-Текст)',
    'macro_chunk_text_recall': 'Macro Полнота (Чанк-Текст)',
    'macro_chunk_text_f1': 'Macro F1 (Чанк-Текст)',
    'macro_assembly_punct_recall': 'Macro Полнота (Сборка-Пункт)',

    'micro_text_precision': 'Micro Точность (Текст)',
    'micro_text_recall': 'Micro Полнота (Текст)',
    'micro_text_f1': 'Micro F1 (Текст)',
}
# --- Конец маппинга ---

def parse_args():
    """Парсит аргументы командной строки."""
    parser = argparse.ArgumentParser(description="Визуализация результатов тестирования RAG")

    parser.add_argument("--results-file", type=str, default=DEFAULT_RESULTS_FILE,
                        help=f"Путь к Excel-файлу с агрегированными результатами (по умолчанию: {DEFAULT_RESULTS_FILE})")
    parser.add_argument("--plots-dir", type=str, default=DEFAULT_PLOTS_DIR,
                        help=f"Директория для сохранения графиков (по умолчанию: {DEFAULT_PLOTS_DIR})")
    parser.add_argument("--sheet-name", type=str, default="Агрегированные метрики",
                        help="Название листа в Excel-файле для чтения данных")

    return parser.parse_args()

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

def load_aggregated_data(file_path: str, sheet_name: str) -> pd.DataFrame:
    """Загружает данные из указанного листа Excel-файла."""
    print(f"Загрузка данных из файла: {file_path}, лист: {sheet_name}")
    try:
        df = pd.read_excel(file_path, sheet_name=sheet_name)
        print(f"Загружено {len(df)} строк.")
        print(f"Колонки: {df.columns.tolist()}")
        # Добавим проверку на необходимые колонки (РУССКИЕ НАЗВАНИЯ)
        required_cols_rus = [
            COLUMN_NAME_MAPPING['model_name'], COLUMN_NAME_MAPPING['chunking_strategy'],
            COLUMN_NAME_MAPPING['strategy_params'], COLUMN_NAME_MAPPING['process_tables'],
            COLUMN_NAME_MAPPING['top_n'], COLUMN_NAME_MAPPING['use_injection'],
            COLUMN_NAME_MAPPING['use_qe'], COLUMN_NAME_MAPPING['neighbors_included'],
            COLUMN_NAME_MAPPING['similarity_threshold']
            ]
        # Проверяем только те, что есть в маппинге
        missing_required = [col for col in required_cols_rus if col not in df.columns]
        if missing_required:
            print(f"Предупреждение: Не все ожидаемые колонки параметров найдены в данных: {missing_required}")

        # --- Добавим парсинг strategy_params из JSON строки в словарь ---
        params_col = COLUMN_NAME_MAPPING['strategy_params']
        if params_col in df.columns:
            def safe_json_loads(x):
                try:
                    # Обработка NaN и пустых строк
                    if pd.isna(x) or not isinstance(x, str) or not x.strip():
                        return {}
                    return json.loads(x)
                except (json.JSONDecodeError, TypeError):
                    return {} # Возвращаем пустой словарь при ошибке

            df[params_col] = df[params_col].apply(safe_json_loads)
            # Создаем строковое представление для группировки и лейблов
            df[f"{params_col}_str"] = df[params_col].apply(
                lambda d: json.dumps(d, sort_keys=True, ensure_ascii=False)
            )
            print(f"Колонка '{params_col}' преобразована из JSON строк.")
        # --------------------------------------------------------------

        return df
    except FileNotFoundError:
        print(f"Ошибка: Файл не найден: {file_path}")
        return pd.DataFrame()
    except ValueError as e:
        print(f"Ошибка: Лист '{sheet_name}' не найден в файле {file_path}. Доступные листы: {pd.ExcelFile(file_path).sheet_names}")
        return pd.DataFrame()
    except Exception as e:
        print(f"Ошибка при чтении Excel файла: {e}")
        return pd.DataFrame()

# --- Функции построения графиков --- #

def plot_metric_vs_top_n(
    df: pd.DataFrame,
    metric_name_rus: str, # Ожидаем русское имя метрики
    fixed_strategy: str | None,
    fixed_strategy_params: str | None, # Ожидаем строку JSON или None
    plots_dir: str
) -> None:
    """
    Строит график зависимости метрики от top_n для разных моделей
    (при фиксированных параметрах чанкинга).
    Разделяет линии по значению use_injection.
    Использует русские названия колонок.
    """
    # Используем русские названия колонок из маппинга
    metric_col_rus = metric_name_rus # Передаем уже готовое русское имя
    top_n_col_rus = COLUMN_NAME_MAPPING['top_n']
    model_col_rus = COLUMN_NAME_MAPPING['model_name']
    injection_col_rus = COLUMN_NAME_MAPPING['use_injection']
    strategy_col_rus = COLUMN_NAME_MAPPING['chunking_strategy']
    params_str_col_rus = f"{COLUMN_NAME_MAPPING['strategy_params']}_str" # Используем строковое представление

    if metric_col_rus not in df.columns:
        print(f"График пропущен: Колонка '{metric_col_rus}' не найдена.")
        return

    plot_df = df.copy()

    # Фильтруем по параметрам чанкинга, если задано
    chunk_suffix = "all_strategies_all_params"
    if fixed_strategy and strategy_col_rus in plot_df.columns:
        plot_df = plot_df[plot_df[strategy_col_rus] == fixed_strategy]
        chunk_suffix = f"strategy_{fixed_strategy}"
        # Фильтруем по строковому представлению параметров
        if fixed_strategy_params and params_str_col_rus in plot_df.columns:
            plot_df = plot_df[plot_df[params_str_col_rus] == fixed_strategy_params]
            # Генерируем короткий хэш для параметров в названии файла
            params_hash = hash(fixed_strategy_params) # Хэш от строки
            chunk_suffix += f"_params-{params_hash:x}" # Hex hash

        if plot_df.empty:
            print(f"График Metric vs Top-N пропущен: Нет данных для strategy={fixed_strategy}, params={fixed_strategy_params}")
            return

    plt.figure(figsize=FIGSIZE)
    sns.lineplot(
        data=plot_df,
        x=top_n_col_rus,
        y=metric_col_rus,
        hue=model_col_rus,
        style=injection_col_rus, # Разные стили линий для True/False
        markers=True,
        markersize=8,
        linewidth=2,
        palette=PALETTE
    )

    plt.title(f"Зависимость {metric_col_rus} от top_n ({chunk_suffix})")
    plt.xlabel("Top N")
    plt.ylabel(metric_col_rus.replace("_", " ").title())
    plt.legend(title="Модель / Сборка", bbox_to_anchor=(1.05, 1), loc='upper left')
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout(rect=[0, 0, 0.85, 1]) # Оставляем место для легенды

    filename = f"plot_{metric_col_rus.replace(' ', '_').replace('(', '').replace(')', '')}_vs_top_n_{chunk_suffix}.png"
    filepath = os.path.join(plots_dir, filename)
    plt.savefig(filepath, dpi=DPI)
    plt.close()
    print(f"Создан график: {filepath}")

def plot_injection_comparison(
    df: pd.DataFrame,
    metric_name_rus: str, # Ожидаем русское имя метрики
    plots_dir: str
) -> None:
    """
    Сравнивает метрики с использованием и без использования сборки контекста
    в виде парных столбчатых диаграмм для разных моделей и параметров чанкинга.
    Использует русские названия колонок.
    """
    # Русские названия колонок
    metric_col_rus = metric_name_rus
    injection_col_rus = COLUMN_NAME_MAPPING['use_injection']
    model_col_rus = COLUMN_NAME_MAPPING['model_name']
    strategy_col_rus = COLUMN_NAME_MAPPING['chunking_strategy']
    params_str_col_rus = f"{COLUMN_NAME_MAPPING['strategy_params']}_str"
    tables_col_rus = COLUMN_NAME_MAPPING['process_tables']
    qe_col_rus = COLUMN_NAME_MAPPING['use_qe']
    neighbors_col_rus = COLUMN_NAME_MAPPING['neighbors_included']
    top_n_col_rus = COLUMN_NAME_MAPPING['top_n']
    threshold_col_rus = COLUMN_NAME_MAPPING['similarity_threshold']

    if metric_col_rus not in df.columns or injection_col_rus not in df.columns:
        print(f"График сравнения сборки пропущен: Колонки '{metric_col_rus}' или '{injection_col_rus}' не найдены.")
        return

    plot_df = df.copy()
    # Используем русские названия при создании лейбла
    plot_df['config_label'] = plot_df.apply(
        lambda r: (
            f"{r.get(model_col_rus, 'N/A')}\n"
            f"Стратегия: {r.get(strategy_col_rus, 'N/A')}\n"
            # Используем строковое представление параметров
            f"Параметры: {r.get(params_str_col_rus, '{}')[:30]}...\n"
            f"Табл: {r.get(tables_col_rus, 'N/A')}, QE: {r.get(qe_col_rus, 'N/A')}, Соседи: {r.get(neighbors_col_rus, 'N/A')}\n"
            f"TopN: {int(r.get(top_n_col_rus, 0))}, Порог: {r.get(threshold_col_rus, 0):.2f}"
        ),
        axis=1
    )

    # Оставляем только строки, где есть и True, и False для данного флага
    # Группируем по config_label, считаем уникальные значения флага use_injection
    counts = plot_df.groupby('config_label')[injection_col_rus].nunique()
    configs_with_both = counts[counts >= 2].index # Используем >= 2 на случай дубликатов
    plot_df = plot_df[plot_df['config_label'].isin(configs_with_both)]

    if plot_df.empty:
        print(f"График сравнения сборки пропущен: Нет конфигураций с обоими вариантами {injection_col_rus}.")
        return

    # Ограничим количество конфигураций для читаемости (по средней метрике)
    top_configs = plot_df.groupby('config_label')[metric_col_rus].mean().nlargest(10).index # Уменьшил до 10
    plot_df = plot_df[plot_df['config_label'].isin(top_configs)]

    if plot_df.empty:
         print(f"График сравнения сборки пропущен: Не осталось данных после фильтрации топ-конфигураций.")
         return

    plt.figure(figsize=(FIGSIZE[0]*0.9, FIGSIZE[1]*0.7)) # Уменьшил размер
    sns.barplot(
        data=plot_df,
        x='config_label',
        y=metric_col_rus,
        hue=injection_col_rus,
        palette=PALETTE
    )

    plt.title(f"Сравнение {metric_col_rus} с/без {injection_col_rus}")
    plt.xlabel("Конфигурация")
    plt.ylabel(metric_col_rus)
    plt.xticks(rotation=60, ha='right', fontsize=8) # Уменьшил шрифт, увеличил поворот
    plt.legend(title=injection_col_rus)
    plt.grid(True, axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()

    filename = f"plot_{metric_col_rus.replace(' ', '_').replace('(', '').replace(')', '')}_injection_comparison.png"
    filepath = os.path.join(plots_dir, filename)
    plt.savefig(filepath, dpi=DPI)
    plt.close()
    print(f"Создан график: {filepath}")

# --- Новая функция для сравнения булевых флагов --- 
def plot_boolean_flag_comparison(
    df: pd.DataFrame,
    metric_name_rus: str, # Ожидаем русское имя метрики
    flag_column_eng: str, # Ожидаем английское имя флага для поиска в маппинге
    plots_dir: str
) -> None:
    """
    Сравнивает метрики при True/False значениях указанного булева флага
    в виде парных столбчатых диаграмм для разных конфигураций.
    Использует русские названия колонок.
    """
    # Русские названия колонок
    metric_col_rus = metric_name_rus
    try:
        flag_col_rus = COLUMN_NAME_MAPPING[flag_column_eng]
    except KeyError:
        print(f"Ошибка: Английское имя флага '{flag_column_eng}' не найдено в COLUMN_NAME_MAPPING.")
        return

    model_col_rus = COLUMN_NAME_MAPPING['model_name']
    strategy_col_rus = COLUMN_NAME_MAPPING['chunking_strategy']
    params_str_col_rus = f"{COLUMN_NAME_MAPPING['strategy_params']}_str"
    injection_col_rus = COLUMN_NAME_MAPPING['use_injection']
    top_n_col_rus = COLUMN_NAME_MAPPING['top_n']
    # Другие флаги
    tables_col_rus = COLUMN_NAME_MAPPING['process_tables']
    qe_col_rus = COLUMN_NAME_MAPPING['use_qe']
    neighbors_col_rus = COLUMN_NAME_MAPPING['neighbors_included']


    if metric_col_rus not in df.columns or flag_col_rus not in df.columns:
        print(f"График сравнения флага '{flag_col_rus}' пропущен: Колонки '{metric_col_rus}' или '{flag_col_rus}' не найдены.")
        return

    plot_df = df.copy()
    # Создаем обобщенный лейбл конфигурации, исключая сам флаг
    plot_df['config_label'] = plot_df.apply(
        lambda r: (
            f"{r.get(model_col_rus, 'N/A')}\n"
            f"Стратегия: {r.get(strategy_col_rus, 'N/A')} Параметры: {r.get(params_str_col_rus, '{}')[:20]}...\n"
            f"Сборка: {r.get(injection_col_rus, 'N/A')}, TopN: {int(r.get(top_n_col_rus, 0))}"
            # Динамически добавляем другие флаги, кроме сравниваемого
            + (f", Табл: {r.get(tables_col_rus, 'N/A')}" if flag_col_rus != tables_col_rus else "")
            + (f", QE: {r.get(qe_col_rus, 'N/A')}" if flag_col_rus != qe_col_rus else "")
            + (f", Соседи: {r.get(neighbors_col_rus, 'N/A')}" if flag_col_rus != neighbors_col_rus else "")
        ),
        axis=1
    )

    # Оставляем только строки, где есть и True, и False для данного флага
    counts = plot_df.groupby('config_label')[flag_col_rus].nunique()
    configs_with_both = counts[counts >= 2].index # Используем >= 2
    plot_df = plot_df[plot_df['config_label'].isin(configs_with_both)]

    if plot_df.empty:
        print(f"График сравнения флага '{flag_col_rus}' пропущен: Нет конфигураций с обоими вариантами {flag_col_rus}.")
        return

    # Ограничим количество конфигураций для читаемости (по средней метрике)
    top_configs = plot_df.groupby('config_label')[metric_col_rus].mean().nlargest(10).index # Уменьшил до 10
    plot_df = plot_df[plot_df['config_label'].isin(top_configs)]

    if plot_df.empty:
         print(f"График сравнения флага '{flag_col_rus}' пропущен: Не осталось данных после фильтрации топ-конфигураций.")
         return

    plt.figure(figsize=(FIGSIZE[0]*0.9, FIGSIZE[1]*0.7)) # Уменьшил размер
    sns.barplot(
        data=plot_df,
        x='config_label',
        y=metric_col_rus,
        hue=flag_col_rus,
        palette=PALETTE
    )

    plt.title(f"Сравнение {metric_col_rus} в зависимости от '{flag_col_rus}'")
    plt.xlabel("Конфигурация")
    plt.ylabel(metric_col_rus)
    plt.xticks(rotation=60, ha='right', fontsize=8) # Уменьшил шрифт, увеличил поворот
    plt.legend(title=f"{flag_col_rus}")
    plt.grid(True, axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()

    filename = f"plot_{metric_col_rus.replace(' ', '_').replace('(', '').replace(')', '')}_{flag_column_eng}_comparison.png"
    filepath = os.path.join(plots_dir, filename)
    plt.savefig(filepath, dpi=DPI)
    plt.close()
    print(f"Создан график: {filepath}")

# --- Основная функция --- 
def main():
    """Основная функция скрипта."""
    args = parse_args()

    setup_plots_directory(args.plots_dir)
    df = load_aggregated_data(args.results_file, args.sheet_name)

    if df.empty:
        print("Нет данных для построения графиков. Завершение.")
        return

    # Определяем метрики для построения графиков (используем английские ключи для поиска русских имен)
    metric_keys = [
        'weighted_chunk_text_recall', 'weighted_chunk_text_f1', 'weighted_assembly_punct_recall',
        'macro_chunk_text_recall', 'macro_chunk_text_f1', 'macro_assembly_punct_recall',
        'micro_text_recall', 'micro_text_f1'
    ]

    # Получаем существующие русские имена метрик в DataFrame
    existing_metrics_rus = [COLUMN_NAME_MAPPING.get(key) for key in metric_keys if COLUMN_NAME_MAPPING.get(key) in df.columns]

    # Определяем фиксированные параметры для некоторых графиков
    strategy_col_rus = COLUMN_NAME_MAPPING.get('chunking_strategy')
    params_str_col_rus = f"{COLUMN_NAME_MAPPING.get('strategy_params')}_str"
    model_col_rus = COLUMN_NAME_MAPPING.get('model_name')

    fixed_strategy_example = df[strategy_col_rus].unique()[0] if strategy_col_rus in df.columns and len(df[strategy_col_rus].unique()) > 0 else None
    fixed_strategy_params_example = None
    if fixed_strategy_example and params_str_col_rus in df.columns:
        params_list = df[df[strategy_col_rus] == fixed_strategy_example][params_str_col_rus].unique()
        if len(params_list) > 0:
            fixed_strategy_params_example = params_list[0]

    fixed_model_example = df[model_col_rus].unique()[0] if model_col_rus in df.columns and len(df[model_col_rus].unique()) > 0 else None
    fixed_top_n_example = 20

    print("--- Построение графиков ---")

    # 1. Графики Metric vs Top-N
    print("\n1. Зависимость метрик от Top-N:")
    for metric_name_rus in existing_metrics_rus:
            # Проверяем, что метрика не micro (у micro нет зависимости от top_n)
            if 'Micro' in metric_name_rus:
                 continue
            plot_metric_vs_top_n(
                df, metric_name_rus,
                fixed_strategy_example, fixed_strategy_params_example,
                args.plots_dir
            )

    # 2. Графики Metric vs Chunking
    print("\n2. Зависимость метрик от параметров чанкинга: [Пропущено - требует переосмысления]")
    # plot_metric_vs_chunking(...) # Закомментировано

    # 3. Графики сравнения Use Injection
    print("\n3. Сравнение метрик с/без сборки контекста:")
    for metric_name_rus in existing_metrics_rus:
             plot_injection_comparison(df, metric_name_rus, args.plots_dir)

    # 4. Графики сравнения других булевых флагов
    boolean_flags_eng = ['process_tables', 'use_qe', 'neighbors_included']
    print("\n4. Сравнение метрик в зависимости от булевых флагов:")
    for flag_eng in boolean_flags_eng:
        flag_rus = COLUMN_NAME_MAPPING.get(flag_eng)
        if not flag_rus or flag_rus not in df.columns:
            print(f"  Пропуск сравнения для флага: '{flag_eng}' (колонка '{flag_rus}' не найдена)")
            continue
        print(f"  Сравнение для флага: '{flag_rus}'")
        for metric_name_rus in existing_metrics_rus:
                 plot_boolean_flag_comparison(df, metric_name_rus, flag_eng, args.plots_dir)

    print("\n--- Построение графиков завершено ---")

if __name__ == "__main__":
    main()