import matplotlib.pyplot as plt
import numpy as np
import csv
from typing import Dict, List
import os

plt.style.use('seaborn-v0_8-darkgrid')

def load_results(filename: str = 'maze_experiment_results.csv') -> List[Dict]:
    if not os.path.exists(filename):
        print(f"Файл {filename} не найден! Сначала запустите main.py")
        return []
    
    with open(filename, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        results = []
        for row in reader:
            row['time_ms'] = float(row['time_ms'])
            row['visited'] = int(float(row['visited']))
            row['path_length'] = int(float(row['path_length']))
            results.append(row)
    return results

def organize_data(results: List[Dict]) -> Dict:
    data = {}
    for row in results:
        maze = row['maze']
        strategy = row['strategy'].split('(')[0].strip()
        
        if maze not in data:
            data[maze] = {}
        if strategy not in data[maze]:
            data[maze][strategy] = {'time_ms': [], 'visited': [], 'path_length': []}
        
        data[maze][strategy]['time_ms'].append(row['time_ms'])
        data[maze][strategy]['visited'].append(row['visited'])
        data[maze][strategy]['path_length'].append(row['path_length'])
    return data

def create_main_chart(data: Dict, save_dir: str = 'charts'):
    
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    
    maze_names = list(data.keys())
    strategies = ['BFS', 'DFS', 'A*']
    colors = {'BFS': '#3498db', 'DFS': '#e74c3c', 'A*': '#2ecc71'}
    
    fig, axes = plt.subplots(1, 3, figsize=(16, 6))
    fig.suptitle('Сравнение алгоритмов поиска пути', fontsize=16, fontweight='bold')
    
    x = np.arange(len(maze_names))
    width = 0.25
    
    # График 1: Время
    for i, strat in enumerate(strategies):
        values = []
        for maze in maze_names:
            if strat in data[maze] and data[maze][strat]['time_ms']:
                values.append(np.mean(data[maze][strat]['time_ms']))
            else:
                values.append(0)
        
        bars = axes[0].bar(x + i * width, values, width, label=strat, 
                          color=colors[strat], edgecolor='black', linewidth=0.5)
        
        for bar, val in zip(bars, values):
            if val > 0:
                axes[0].text(bar.get_x() + bar.get_width()/2., bar.get_height(),
                           f'{val:.1f}', ha='center', va='bottom', fontsize=7)
    
    axes[0].set_title('Время выполнения (мс)', fontweight='bold')
    axes[0].set_xticks(x + width)
    axes[0].set_xticklabels([n.split('(')[0].strip() for n in maze_names], rotation=15, fontsize=8)
    axes[0].legend(fontsize=8)
    
    # График 2: Посещённые клетки
    for i, strat in enumerate(strategies):
        values = []
        for maze in maze_names:
            if strat in data[maze] and data[maze][strat]['visited']:
                values.append(np.mean(data[maze][strat]['visited']))
            else:
                values.append(0)
        
        bars = axes[1].bar(x + i * width, values, width, label=strat,
                          color=colors[strat], edgecolor='black', linewidth=0.5)
        
        for bar, val in zip(bars, values):
            if val > 0:
                axes[1].text(bar.get_x() + bar.get_width()/2., bar.get_height(),
                           f'{int(val)}', ha='center', va='bottom', fontsize=7)
    
    axes[1].set_title('Посещённые клетки', fontweight='bold')
    axes[1].set_xticks(x + width)
    axes[1].set_xticklabels([n.split('(')[0].strip() for n in maze_names], rotation=15, fontsize=8)
    
    # График 3: Длина пути
    for i, strat in enumerate(strategies):
        values = []
        for maze in maze_names:
            if strat in data[maze] and data[maze][strat]['path_length']:
                values.append(np.mean(data[maze][strat]['path_length']))
            else:
                values.append(0)
        
        bars = axes[2].bar(x + i * width, values, width, label=strat,
                          color=colors[strat], edgecolor='black', linewidth=0.5)
        
        for bar, val in zip(bars, values):
            if val > 0:
                axes[2].text(bar.get_x() + bar.get_width()/2., bar.get_height(),
                           f'{int(val)}', ha='center', va='bottom', fontsize=7)
    
    axes[2].set_title('Длина пути (шагов)', fontweight='bold')
    axes[2].set_xticks(x + width)
    axes[2].set_xticklabels([n.split('(')[0].strip() for n in maze_names], rotation=15, fontsize=8)
    
    plt.tight_layout()
    filepath = os.path.join(save_dir, 'comparison_all_metrics.png')
    plt.savefig(filepath, dpi=150, bbox_inches='tight')
    plt.close()
    print(f" График сохранён: {filepath}")

def print_table(data: Dict):
    print("\n" + "=" * 80)
    print("РЕЗУЛЬТАТЫ ТЕСТИРОВАНИЯ")
    print("=" * 80)
    
    for maze_name, strategies in data.items():
        print(f"\n{maze_name}")
        print("-" * 60)
        print(f"{'Алгоритм':<10} {'Время(мс)':<12} {'Посещено':<12} {'Длина пути':<12}")
        print("-" * 46)
        
        for strat in ['BFS', 'DFS', 'A*']:
            if strat in strategies and strategies[strat]['time_ms']:
                t = np.mean(strategies[strat]['time_ms'])
                v = np.mean(strategies[strat]['visited'])
                p = np.mean(strategies[strat]['path_length'])
                print(f"{strat:<10} {t:<12.2f} {int(v):<12} {int(p):<12}")

def main():
    print("=" * 60)
    print("ВИЗУАЛИЗАЦИЯ РЕЗУЛЬТАТОВ")
    print("=" * 60)
    
    results = load_results('maze_experiment_results.csv')
    
    if not results:
        print("Нет данных. Запустите сначала main.py")
        return
    
    data = organize_data(results)
    
    print_table(data)
    
    # Один график
    print("\nСоздание графика...")
    create_main_chart(data)

if __name__ == "__main__":
    main()