import csv
import re
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

BASE_DIR = Path(__file__).resolve().parent


def generate_plots():
    csv_path = BASE_DIR.parent / "results.csv"
    if not csv_path.exists():
        print(f"Error: {csv_path} not found. Run experiment.py first.")
        return

    results = []
    with open(csv_path, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        for row in reader:
            results.append(
                {
                    "maze": row["maze"],
                    "strategy": row["strategy"],
                    "time_ms": float(row["time_ms"]),
                    "visited_cells": int(row["visited_cells"]),
                    "path_length": int(row["path_length"]),
                }
            )

    # Sort mazes by requested logical order: no_exit, empty, then by size (NxN)
    unique_mazes = list(dict.fromkeys(r["maze"] for r in results))

    def get_sort_key(m_name):
        name = m_name.lower()
        if "no_exit" in name or "noexit" in name:
            return 0
        if "empty" in name:
            return 1

        match = re.search(r"(\d+)x\d+", name)
        if match:
            return 100 + int(match.group(1))

        return 999

    maze_files_keys = sorted(unique_mazes, key=get_sort_key)

    fig, axes = plt.subplots(
        len(maze_files_keys), 3, figsize=(18, 3 * len(maze_files_keys))
    )

    for idx, maze_name in enumerate(maze_files_keys):
        maze_res = [r for r in results if r["maze"] == maze_name]
        if not maze_res:
            continue

        strats = [r["strategy"] for r in maze_res]
        times = [r["time_ms"] for r in maze_res]
        visited = [r["visited_cells"] for r in maze_res]
        path_lens = [r["path_length"] for r in maze_res]

        x = np.arange(len(strats))

        # Check if axes is 1D or 2D depending on number of mazes
        ax_time = axes[0] if len(maze_files_keys) == 1 else axes[idx, 0]
        ax_visited = axes[1] if len(maze_files_keys) == 1 else axes[idx, 1]
        ax_path = axes[2] if len(maze_files_keys) == 1 else axes[idx, 2]

        ax_time.bar(x, times, color=["red", "green", "blue"])
        ax_time.set_xticks(x)
        ax_time.set_xticklabels(strats)
        ax_time.set_title(f"{maze_name}: Execution Time (ms)")

        ax_visited.bar(x, visited, color=["red", "green", "blue"])
        ax_visited.set_xticks(x)
        ax_visited.set_xticklabels(strats)
        ax_visited.set_title(f"{maze_name}: Visited Cells")

        ax_path.bar(x, path_lens, color=["red", "green", "blue"])
        ax_path.set_xticks(x)
        ax_path.set_xticklabels(strats)
        ax_path.set_title(f"{maze_name}: Path Length")

    plt.tight_layout()
    chart_path = BASE_DIR.parent / "benchmark_charts.png"
    plt.savefig(chart_path)
    print(f"Charts exported to {chart_path}")


if __name__ == "__main__":
    generate_plots()