import csv
import os
import statistics
import matplotlib.pyplot as plt

from maze_builder import TextFileMazeBuilder
from solver import MazeSolver
from strategies import BFSStrategy, DFSStrategy, AStarStrategy

# --- НАСТРОЙКИ ---
MAZES_DIR = "mazes"
OUTPUT_CSV = "results.csv"
RUNS = 10  # Количество запусков для усреднения
PLOTS_DIR = "plots" # Новая папка для графиков

STRATEGIES = {
    "BFS": BFSStrategy,
    "DFS": DFSStrategy,
    "A*": AStarStrategy,
}

# Словарь для хранения всех данных для графиков
all_data = {}

# Создаем папку для графиков, если её нет
os.makedirs(PLOTS_DIR, exist_ok=True)

builder = TextFileMazeBuilder()
maze_files = sorted(f for f in os.listdir(MAZES_DIR) if f.endswith(".txt"))
rows = []

print("=== СТАРТ ЭКСПЕРИМЕНТА ===\n")

# --- ОСНОВНОЙ ЦИКЛ ЭКСПЕРИМЕНТА ---
for maze_file in maze_files:
    maze_name = maze_file.replace(".txt", "")
    filepath = os.path.join(MAZES_DIR, maze_file)

    try:
        maze = builder.build_from_file(filepath)
    except ValueError as e:
        print(f"  [!] Пропуск {maze_file}: {e}")
        continue # Переходим к следующему файлу, если этот не загрузился

    # Эта строка теперь выполняется для каждого успешного лабиринта
    print(f"\n{'='*50}")
    print(f"Лабиринт: {maze_name}  ({maze.width}×{maze.height})")
    
    all_data[maze_name] = {}

    for strat_name, StratClass in STRATEGIES.items():
        times, visited_counts, path_lengths = [], [], []
        run_stats = []

        for run_num in range(1, RUNS + 1):
            solver = MazeSolver(maze, StratClass())
            stats = solver.solve()
            
            times.append(stats.time_ms)
            visited_counts.append(stats.visited_cells)
            path_lengths.append(stats.path_length)
            
            # Сохраняем данные каждой попытки
            run_stats.append({
                'попытка': run_num,
                'время_мс': stats.time_ms,
                'посещено_клеток': stats.visited_cells,
                'длина_пути': stats.path_length
            })
            
            print(f"  {strat_name} | Попытка {run_num}/{RUNS} | Время: {stats.time_ms:.2f} мс")

        # Вычисляем средние значения
        avg_time = statistics.mean(times)
        avg_visited = statistics.mean(visited_counts)
        
        valid_path_lengths = [p for p in path_lengths if p is not None]
        avg_path = statistics.mean(valid_path_lengths) if valid_path_lengths else None

        print(f"  {strat_name:10s} | СРЕДНЕЕ: время {avg_time:.2f} мс | "
              f"посещено {avg_visited:.1f} | путь {avg_path if avg_path is not None else '—'}")

        rows.append({
            "лабиринт": maze_name,
            "стратегия": strat_name,
            "время_мс": round(avg_time, 6),
            "посещено_клеток": round(avg_visited, 1),
            "длина_пути": round(avg_path, 1) if avg_path is not None else None,
        })
        
        all_data[maze_name][strat_name] = run_stats

# --- СОХРАНЕНИЕ CSV ---
with open(OUTPUT_CSV, "w", newline="", encoding="utf-8") as csvfile:
    fieldnames = ["лабиринт", "стратегия", "время_мс", "посещено_клеток", "длина_пути"]
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()
    writer.writerows(rows)
print(f"\n✓ Результаты сохранены в {OUTPUT_CSV}")

# --- ПОСТРОЕНИЕ ГРАФИКОВ ---
print("\n=== ПОСТРОЕНИЕ ГРАФИКОВ ===")
for maze_name, strategies_data in all_data.items():
    print(f"\nСтроим графики для лабиринта: {maze_name}")
    
    # Создаем ОДИН график для времени выполнения
    fig, ax = plt.subplots(figsize=(10, 6))
    fig.suptitle(f"Сравнение времени выполнения алгоритмов\nЛабиринт '{maze_name}'", fontsize=14)
    
    ax.set_title("Время выполнения (мс)")
    ax.set_xlabel("Номер попытки")
    ax.set_ylabel("Время (мс)")
    
    for strat_name in STRATEGIES.keys():
        # Извлекаем только данные о времени выполнения
        data_points = [
            run['время_мс'] for run in strategies_data.get(strat_name, [])
        ]
        
        if data_points:
            x_values = range(1, len(data_points) + 1)
            ax.plot(x_values, data_points, 
                    marker='o', 
                    label=strat_name, 
                    linewidth=2)
    
    ax.legend(title="Алгоритм")
    ax.grid(True, which='both', linestyle='--', linewidth=0.5)
    
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])
    
    # Сохраняем график в папку 'plots'
    plot_filename = os.path.join(PLOTS_DIR, f"plot_{maze_name}.png")
    plt.savefig(plot_filename)
    plt.close()

print(f"\n✓ Графики времени выполнения сохранены в папку '{PLOTS_DIR}'")
print("=== ЭКСПЕРИМЕНТ ЗАВЕРШЕН ===")