#!/usr/bin/env python
"""
Скрипт для построения специализированных графиков на основе макрометрик из Excel-файла.
Строит несколько типов графиков:
1. Зависимость macro_text_recall от top_N для разных моделей при фиксированных параметрах чанкинга
2. Зависимость macro_text_recall от top_N для разных подходов к чанкингу при фиксированных моделях
3. Зависимость macro_text_recall от подхода к чанкингу для разных моделей при фиксированных top_N
"""

import os

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

# Константы
EXCEL_FILE_PATH = "../../Белагропромбанк/test_vectors/combined_results.xlsx"
PLOTS_DIR = "../../Белагропромбанк/test_vectors/plots"

# Настройки для графиков
plt.rcParams['font.family'] = 'DejaVu Sans'
sns.set_style("whitegrid")
FIGSIZE = (14, 10)
DPI = 300


def setup_plots_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 load_macro_metrics(excel_path: str) -> pd.DataFrame:
    """
    Загружает макрометрики из Excel-файла.
    
    Args:
        excel_path: Путь к Excel-файлу с данными
        
    Returns:
        DataFrame с макрометриками
    """
    try:
        df = pd.read_excel(excel_path, sheet_name="Macro метрики")
        print(f"Загружены данные из {excel_path}, лист 'Macro метрики'")
        print(f"Количество строк: {len(df)}")
        return df
    except Exception as e:
        print(f"Ошибка при загрузке данных: {e}")
        raise


def plot_top_n_vs_recall_by_model(df: pd.DataFrame, plots_dir: str) -> None:
    """
    Строит графики зависимости macro_text_recall от top_N для разных моделей
    при фиксированных параметрах чанкинга (50/25 и 200/75).
    
    Args:
        df: DataFrame с данными
        plots_dir: Директория для сохранения графиков
    """
    # Фиксированные параметры чанкинга
    chunking_params = [
        {"words": 50, "overlap": 25, "title": "Чанкинг 50/25"},
        {"words": 200, "overlap": 75, "title": "Чанкинг 200/75"}
    ]
    
    # Создаем субплоты: 1 строка, 2 столбца
    fig, axes = plt.subplots(1, 2, figsize=FIGSIZE, sharey=True)
    
    for i, params in enumerate(chunking_params):
        # Фильтруем данные для текущих параметров чанкинга
        filtered_df = df[
            (df['words_per_chunk'] == params['words']) & 
            (df['overlap_words'] == params['overlap'])
        ]
        
        if len(filtered_df) == 0:
            print(f"Предупреждение: нет данных для чанкинга {params['words']}/{params['overlap']}")
            axes[i].text(0.5, 0.5, f"Нет данных для чанкинга {params['words']}/{params['overlap']}", 
                        ha='center', va='center', fontsize=12)
            axes[i].set_title(params['title'])
            continue
        
        # Находим уникальные модели
        models = filtered_df['model'].unique()
        
        # Создаем палитру цветов
        palette = sns.color_palette("viridis", len(models))
        
        # Строим график для каждой модели
        for j, model in enumerate(models):
            model_df = filtered_df[filtered_df['model'] == model].sort_values('top_n')
            
            if len(model_df) <= 1:
                print(f"Предупреждение: недостаточно данных для модели {model} при чанкинге {params['words']}/{params['overlap']}")
                continue
                
            # Строим ломаную линию
            axes[i].plot(model_df['top_n'], model_df['macro_text_recall'], 
                         marker='o', linestyle='-', linewidth=2,
                         label=model, color=palette[j])
        
        # Настраиваем оси и заголовок
        axes[i].set_title(params['title'], fontsize=14)
        axes[i].set_xlabel('top_N', fontsize=12)
        if i == 0:
            axes[i].set_ylabel('macro_text_recall', fontsize=12)
        
        # Добавляем сетку
        axes[i].grid(True, linestyle='--', alpha=0.7)
        
        # Добавляем легенду
        axes[i].legend(title="Модель", fontsize=10, loc='best')
    
    # Общий заголовок
    plt.suptitle('Зависимость macro_text_recall от top_N для разных моделей', fontsize=16)
    
    # Настраиваем макет
    plt.tight_layout(rect=[0, 0, 1, 0.96])
    
    # Сохраняем график
    file_path = os.path.join(plots_dir, "top_n_vs_recall_by_model.png")
    plt.savefig(file_path, dpi=DPI)
    plt.close()
    
    print(f"Создан график: {file_path}")


def plot_top_n_vs_recall_by_chunking(df: pd.DataFrame, plots_dir: str) -> None:
    """
    Строит графики зависимости macro_text_recall от top_N для разных параметров чанкинга
    при фиксированных моделях (bge и frida).
    
    Args:
        df: DataFrame с данными
        plots_dir: Директория для сохранения графиков
    """
    # Фиксированные модели
    models = ["BAAI/bge", "frida"]
    
    # Создаем субплоты: 1 строка, 2 столбца
    fig, axes = plt.subplots(1, 2, figsize=FIGSIZE, sharey=True)
    
    for i, model_name in enumerate(models):
        # Находим все строки с моделями, содержащими указанное название
        model_df = df[df['model'].str.contains(model_name, case=False)]
        
        if len(model_df) == 0:
            print(f"Предупреждение: нет данных для модели {model_name}")
            axes[i].text(0.5, 0.5, f"Нет данных для модели {model_name}", 
                        ha='center', va='center', fontsize=12)
            axes[i].set_title(f"Модель: {model_name}")
            continue
        
        # Находим уникальные комбинации параметров чанкинга
        chunking_combinations = model_df.drop_duplicates(['words_per_chunk', 'overlap_words'])[['words_per_chunk', 'overlap_words']]
        
        # Ограничиваем количество комбинаций до 7 для читаемости
        if len(chunking_combinations) > 7:
            print(f"Предупреждение: слишком много комбинаций чанкинга для модели {model_name}, ограничиваем до 7")
            chunking_combinations = chunking_combinations.head(7)
        
        # Создаем палитру цветов
        palette = sns.color_palette("viridis", len(chunking_combinations))
        
        # Строим график для каждой комбинации параметров чанкинга
        for j, (_, row) in enumerate(chunking_combinations.iterrows()):
            words = row['words_per_chunk']
            overlap = row['overlap_words']
            
            # Фильтруем данные для текущей модели и параметров чанкинга
            chunking_df = model_df[
                (model_df['words_per_chunk'] == words) & 
                (model_df['overlap_words'] == overlap)
            ].sort_values('top_n')
            
            if len(chunking_df) <= 1:
                print(f"Предупреждение: недостаточно данных для модели {model_name} с чанкингом {words}/{overlap}")
                continue
                
            # Строим ломаную линию
            axes[i].plot(chunking_df['top_n'], chunking_df['macro_text_recall'], 
                         marker='o', linestyle='-', linewidth=2,
                         label=f"w={words}, o={overlap}", color=palette[j])
        
        # Настраиваем оси и заголовок
        axes[i].set_title(f"Модель: {model_name}", fontsize=14)
        axes[i].set_xlabel('top_N', fontsize=12)
        if i == 0:
            axes[i].set_ylabel('macro_text_recall', fontsize=12)
        
        # Добавляем сетку
        axes[i].grid(True, linestyle='--', alpha=0.7)
        
        # Добавляем легенду
        axes[i].legend(title="Чанкинг", fontsize=10, loc='best')
    
    # Общий заголовок
    plt.suptitle('Зависимость macro_text_recall от top_N для разных параметров чанкинга', fontsize=16)
    
    # Настраиваем макет
    plt.tight_layout(rect=[0, 0, 1, 0.96])
    
    # Сохраняем график
    file_path = os.path.join(plots_dir, "top_n_vs_recall_by_chunking.png")
    plt.savefig(file_path, dpi=DPI)
    plt.close()
    
    print(f"Создан график: {file_path}")


def plot_chunking_vs_recall_by_model(df: pd.DataFrame, plots_dir: str) -> None:
    """
    Строит графики зависимости macro_text_recall от подхода к чанкингу 
    для разных моделей при фиксированных top_N (5, 20, 100).
    
    Args:
        df: DataFrame с данными
        plots_dir: Директория для сохранения графиков
    """
    # Фиксированные значения top_N
    top_n_values = [5, 20, 100]
    
    # Создаем субплоты: 1 строка, 3 столбца
    fig, axes = plt.subplots(1, 3, figsize=FIGSIZE, sharey=True)
    
    # Создаем порядок чанкинга - сортируем по возрастанию размера и оверлапа
    chunking_order = df.drop_duplicates(['words_per_chunk', 'overlap_words'])[['words_per_chunk', 'overlap_words']]
    chunking_order = chunking_order.sort_values(['words_per_chunk', 'overlap_words'])
    
    # Создаем словарь для маппинга комбинаций чанкинга на индексы
    chunking_labels = [f"{row['words_per_chunk']}/{row['overlap_words']}" for _, row in chunking_order.iterrows()]
    chunking_map = {f"{row['words_per_chunk']}/{row['overlap_words']}": i for i, (_, row) in enumerate(chunking_order.iterrows())}
    
    for i, top_n in enumerate(top_n_values):
        # Фильтруем данные для текущего top_N
        top_n_df = df[df['top_n'] == top_n]
        
        if len(top_n_df) == 0:
            print(f"Предупреждение: нет данных для top_N={top_n}")
            axes[i].text(0.5, 0.5, f"Нет данных для top_N={top_n}", 
                        ha='center', va='center', fontsize=12)
            axes[i].set_title(f"top_N={top_n}")
            continue
        
        # Находим уникальные модели
        models = top_n_df['model'].unique()
        
        # Ограничиваем количество моделей до 5 для читаемости
        if len(models) > 5:
            print(f"Предупреждение: слишком много моделей для top_N={top_n}, ограничиваем до 5")
            models = models[:5]
        
        # Создаем палитру цветов
        palette = sns.color_palette("viridis", len(models))
        
        # Строим график для каждой модели
        for j, model in enumerate(models):
            model_df = top_n_df[top_n_df['model'] == model].copy()
            
            if len(model_df) <= 1:
                print(f"Предупреждение: недостаточно данных для модели {model} при top_N={top_n}")
                continue
            
            # Создаем новую колонку с индексом чанкинга для сортировки
            model_df['chunking_index'] = model_df.apply(
                lambda row: chunking_map.get(f"{row['words_per_chunk']}/{row['overlap_words']}", -1), 
                axis=1
            )
            
            # Отбрасываем строки с неизвестными комбинациями чанкинга
            model_df = model_df[model_df['chunking_index'] >= 0]
            
            if len(model_df) <= 1:
                continue
                
            # Сортируем по индексу чанкинга
            model_df = model_df.sort_values('chunking_index')
            
            # Создаем список индексов и значений для графика
            x_indices = model_df['chunking_index'].tolist()
            y_values = model_df['macro_text_recall'].tolist()
            
            # Строим ломаную линию
            axes[i].plot(x_indices, y_values, marker='o', linestyle='-', linewidth=2,
                        label=model, color=palette[j])
        
        # Настраиваем оси и заголовок
        axes[i].set_title(f"top_N={top_n}", fontsize=14)
        axes[i].set_xlabel('Подход к чанкингу', fontsize=12)
        if i == 0:
            axes[i].set_ylabel('macro_text_recall', fontsize=12)
        
        # Устанавливаем метки на оси X (подходы к чанкингу)
        axes[i].set_xticks(range(len(chunking_labels)))
        axes[i].set_xticklabels(chunking_labels, rotation=45, ha='right', fontsize=10)
        
        # Добавляем сетку
        axes[i].grid(True, linestyle='--', alpha=0.7)
        
        # Добавляем легенду
        axes[i].legend(title="Модель", fontsize=10, loc='best')
    
    # Общий заголовок
    plt.suptitle('Зависимость macro_text_recall от подхода к чанкингу для разных моделей', fontsize=16)
    
    # Настраиваем макет
    plt.tight_layout(rect=[0, 0, 1, 0.96])
    
    # Сохраняем график
    file_path = os.path.join(plots_dir, "chunking_vs_recall_by_model.png")
    plt.savefig(file_path, dpi=DPI)
    plt.close()
    
    print(f"Создан график: {file_path}")


def main():
    """Основная функция скрипта."""
    # Создаем директорию для графиков
    setup_plots_directory(PLOTS_DIR)
    
    # Загружаем данные
    try:
        macro_metrics = load_macro_metrics(EXCEL_FILE_PATH)
    except Exception as e:
        print(f"Критическая ошибка: {e}")
        return
    
    # Строим графики
    plot_top_n_vs_recall_by_model(macro_metrics, PLOTS_DIR)
    plot_top_n_vs_recall_by_chunking(macro_metrics, PLOTS_DIR)
    plot_chunking_vs_recall_by_model(macro_metrics, PLOTS_DIR)
    
    print("Готово! Все графики созданы.")


if __name__ == "__main__":
    main()