import matplotlib.pyplot as plt
import csv
import numpy as np
from collections import defaultdict

# Читаем результаты из CSV и усредняем по 5 повторам
data = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))

with open('docs/data/results.csv', 'r', encoding='utf-8') as f:
    reader = csv.reader(f)
    header = next(reader)  # пропускаем заголовок
    
    # Проверяем, есть ли колонка "Повтор"
    has_repeat = 'Повтор' in header
    
    for row in reader:
        struct = row[0]
        mode = row[1]
        op = row[2]
        time_val = float(row[3])
        
        # Сохраняем все замеры
        data[op][struct][mode].append(time_val)

# Усредняем
avg_data = {}
for op in data:
    avg_data[op] = {}
    for struct in data[op]:
        avg_data[op][struct] = {}
        for mode in data[op][struct]:
            times = data[op][struct][mode]
            avg_data[op][struct][mode] = sum(times) / len(times)

# Создаём графики
operations = ['вставка', 'поиск', 'удаление']
fig, axes = plt.subplots(1, 3, figsize=(15, 5))

colors = {'случайный': '#1f77b4', 'отсортированный': '#d62728'}

for idx, op in enumerate(operations):
    ax = axes[idx]
    
    structures = ['Связный список', 'Хеш-таблица', 'Двоичное дерево']
    data_keys = ['LinkedList', 'HashTable', 'BST']
    
    x = np.arange(len(structures))
    width = 0.35
    
    random_times = [avg_data[op][key].get('случайный', 0) for key in data_keys]
    sorted_times = [avg_data[op][key].get('отсортированный', 0) for key in data_keys]
    
    bars1 = ax.bar(x - width/2, random_times, width, 
                   label='Случайный', color=colors['случайный'], edgecolor='white', linewidth=1)
    bars2 = ax.bar(x + width/2, sorted_times, width, 
                   label='Отсортированный', color=colors['отсортированный'], edgecolor='white', linewidth=1)
    
    ax.set_ylabel('Время (секунды)', fontsize=11)
    ax.set_title(f'{op.upper()}', fontsize=13, fontweight='bold')
    ax.set_xticks(x)
    ax.set_xticklabels(structures, fontsize=10)
    
    ax.grid(True, axis='y', alpha=0.3, linestyle='--')
    ax.set_axisbelow(True)
    
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    
    for bar in bars1:
        height = bar.get_height()
        if height > 0:
            ax.annotate(f'{height:.4f}',
                       xy=(bar.get_x() + bar.get_width()/2, height),
                       xytext=(0, 3), textcoords="offset points",
                       ha='center', va='bottom', fontsize=8)
    
    for bar in bars2:
        height = bar.get_height()
        if height > 0:
            ax.annotate(f'{height:.4f}',
                       xy=(bar.get_x() + bar.get_width()/2, height),
                       xytext=(0, 3), textcoords="offset points",
                       ha='center', va='bottom', fontsize=8)

fig.legend(labels=['Случайный', 'Отсортированный'], 
           loc='lower center', bbox_to_anchor=(0.5, -0.05), 
           ncol=2, fontsize=11, frameon=True, fancybox=True, shadow=True)

plt.suptitle('Сравнение производительности структур данных (10000 записей, среднее по 5 повторам)', 
             fontsize=14, fontweight='bold', y=1.02)
plt.tight_layout()
plt.subplots_adjust(bottom=0.12)
plt.savefig('docs/data/graph.png', dpi=150, bbox_inches='tight', facecolor='white')
plt.show()

print("График сохранён в docs/data/graph.png (использованы средние значения по 5 повторам)")