from pathlib import Path
from statistics import mean
import csv
import random

import matplotlib.pyplot as plt

from core.cell import Cell
from core.maze import Maze
from solver.maze_solver import MazeSolver
from strategies.astar_strategy import AStarStrategy
from strategies.bfs_strategy import BFSStrategy
from strategies.dfs_strategy import DFSStrategy
from strategies.dijkstra_strategy import DijkstraStrategy


BASE_DIR = Path(__file__).resolve().parent
OUT_DIR = BASE_DIR / "experiment_results"


def build_maze_from_symbols(lines):
    height = len(lines)
    width = max(len(line) for line in lines)
    cells = []
    start = None
    exit_cell = None
    for y, line in enumerate(lines):
        row = []
        for x in range(width):
            ch = line[x] if x < len(line) else "#"
            if ch == "#":
                cell = Cell(x, y, isWall=True)
            elif ch == "S":
                cell = Cell(x, y, isWall=False, isStart=True)
                start = cell
            elif ch == "E":
                cell = Cell(x, y, isWall=False, isExit=True)
                exit_cell = cell
            elif ch == " " or ch == ".":
                cell = Cell(x, y, isWall=False)
            elif ch.isdigit():
                cell = Cell(x, y, isWall=False, weight=int(ch))
            else:
                raise ValueError(f"Unknown symbol '{ch}' at {x},{y}")
            row.append(cell)
        cells.append(row)
    return Maze(cells, width, height, start, exit_cell)


def generate_empty_maze(width, height):
    lines = [" " * width for _ in range(height)]
    lines = [list(row) for row in lines]
    lines[1][1] = "S"
    lines[height - 2][width - 2] = "E"
    return build_maze_from_symbols(["".join(row) for row in lines])


def generate_simple_maze(width, height):
    grid = [["#" for _ in range(width)] for _ in range(height)]
    for x in range(1, width - 1):
        grid[1][x] = " "
    for y in range(1, height - 1):
        grid[y][width - 2] = " "
    grid[1][1] = "S"
    grid[height - 2][width - 2] = "E"
    return build_maze_from_symbols(["".join(row) for row in grid])


def generate_branching_maze(width, height, seed=42, wall_density=0.30):
    rng = random.Random(seed)
    grid = [["#" for _ in range(width)] for _ in range(height)]
    x, y = 1, 1
    grid[y][x] = "S"
    while (x, y) != (width - 2, height - 2):
        candidates = []
        for dx, dy in [(1, 0), (0, 1)]:
            nx, ny = x + dx, y + dy
            if 1 <= nx < width - 1 and 1 <= ny < height - 1:
                candidates.append((nx, ny))
        if not candidates:
            break
        x, y = rng.choice(candidates)
        grid[y][x] = " "
    grid[height - 2][width - 2] = "E"

    # carve extra corridors and dead ends
    for yy in range(1, height - 1):
        for xx in range(1, width - 1):
            if grid[yy][xx] == "#" and rng.random() > wall_density:
                grid[yy][xx] = " "
    grid[1][1] = "S"
    grid[height - 2][width - 2] = "E"
    return build_maze_from_symbols(["".join(row) for row in grid])


def generate_no_path_maze(width, height):
    grid = [[" " for _ in range(width)] for _ in range(height)]
    for x in range(width):
        grid[height // 2][x] = "#"
    grid[1][1] = "S"
    grid[height - 2][width - 2] = "E"
    return build_maze_from_symbols(["".join(row) for row in grid])


def generate_weighted_maze(width, height, seed=123):
    rng = random.Random(seed)
    grid = [[" " for _ in range(width)] for _ in range(height)]
    for y in range(height):
        for x in range(width):
            r = rng.random()
            if r < 0.12:
                grid[y][x] = "#"
            elif r < 0.25:
                grid[y][x] = "3"
            elif r < 0.40:
                grid[y][x] = "2"
            else:
                grid[y][x] = "1"
    # ensure path-ish
    for x in range(width):
        grid[1][x] = "1"
    for y in range(1, height):
        grid[y][width - 2] = "1"
    grid[1][1] = "S"
    grid[height - 2][width - 2] = "E"
    return build_maze_from_symbols(["".join(row) for row in grid])


def bench_one_maze(maze_name, maze, strategies, repeats=5):
    summary_rows = []
    raw_rows = []
    for strategy_name, strategy_factory in strategies:
        times, visiteds, lengths = [], [], []
        for run in range(1, repeats + 1):
            solver = MazeSolver(maze)
            solver.setStrategy(strategy_factory())
            stats = solver.solve()
            raw_rows.append([maze_name, strategy_name, run, f"{stats.timeMs:.6f}", stats.visitedCells, stats.pathLength])
            times.append(stats.timeMs)
            visiteds.append(stats.visitedCells)
            lengths.append(stats.pathLength)
        summary_rows.append([maze_name, strategy_name, f"{mean(times):.6f}", f"{mean(visiteds):.2f}", f"{mean(lengths):.2f}", repeats])
    return summary_rows, raw_rows


def save_csv(path, rows):
    with open(path, "w", newline="", encoding="utf-8") as f:
        csv.writer(f).writerows(rows)


def plot_summary(summary_rows):
    by_maze = {}
    for row in summary_rows[1:]:
        maze_name, strategy, avg_time, avg_visited, avg_len, runs = row
        by_maze.setdefault(maze_name, []).append((strategy, float(avg_time), float(avg_visited), float(avg_len)))

    for maze_name, items in by_maze.items():
        items.sort(key=lambda t: t[0])
        strategies = [i[0] for i in items]
        x = list(range(len(strategies)))

        plt.figure(figsize=(8, 4))
        plt.bar(x, [i[1] for i in items])
        plt.xticks(x, strategies)
        plt.ylabel("ms")
        plt.title(f"{maze_name} — avg time")
        plt.tight_layout()
        plt.savefig(OUT_DIR / f"{maze_name}_time.png", dpi=150)
        plt.close()

        plt.figure(figsize=(8, 4))
        plt.bar(x, [i[2] for i in items])
        plt.xticks(x, strategies)
        plt.ylabel("cells")
        plt.title(f"{maze_name} — visited cells")
        plt.tight_layout()
        plt.savefig(OUT_DIR / f"{maze_name}_visited.png", dpi=150)
        plt.close()

        plt.figure(figsize=(8, 4))
        plt.bar(x, [i[3] for i in items])
        plt.xticks(x, strategies)
        plt.ylabel("cells")
        plt.title(f"{maze_name} — path length")
        plt.tight_layout()
        plt.savefig(OUT_DIR / f"{maze_name}_length.png", dpi=150)
        plt.close()


def main():
    OUT_DIR.mkdir(exist_ok=True)

    strategies = [
        ("BFS", BFSStrategy),
        ("DFS", DFSStrategy),
        ("A*", AStarStrategy),
        ("Dijkstra", DijkstraStrategy),
    ]

    mazes = [
        ("small_10x10", generate_simple_maze(10, 10)),
        ("medium_50x50", generate_branching_maze(50, 50)),
        ("large_100x100", generate_branching_maze(100, 100, seed=99, wall_density=0.35)),
        ("empty_30x30", generate_empty_maze(30, 30)),
        ("no_path_30x30", generate_no_path_maze(30, 30)),
        ("weighted_30x30", generate_weighted_maze(30, 30)),
    ]

    summary = [["maze", "strategy", "avg_time_ms", "avg_visited_cells", "avg_path_length", "runs"]]
    raw = [["maze", "strategy", "run", "time_ms", "visited_cells", "path_length"]]

    for maze_name, maze in mazes:
        s_rows, r_rows = bench_one_maze(maze_name, maze, strategies, repeats=5)
        summary.extend(s_rows)
        raw.extend(r_rows)

    save_csv(OUT_DIR / "summary.csv", summary)
    save_csv(OUT_DIR / "raw.csv", raw)
    plot_summary(summary)

    print("Saved to", OUT_DIR.resolve())


if __name__ == "__main__":
    main()