import sys
import os
from collections import deque
import heapq
import time
import csv
import matplotlib.pyplot as plt
import numpy as np

class GridPoint:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.blocked = False
        self.is_start = False
        self.is_exit = False

    def can_step(self):
        return not self.blocked

class Labyrinth:
    def __init__(self, w, h):
        self.w = w
        self.h = h
        self.grid = [[GridPoint(x, y) for x in range(w)] for y in range(h)]
        self.start_point = None
        self.exit_point = None

    def get_point(self, x, y):
        if 0 <= x < self.w and 0 <= y < self.h:
            return self.grid[y][x]
        return None

    def set_point(self, x, y, typ):
        p = self.get_point(x, y)
        if not p:
            return
        if typ == 'wall':
            p.blocked = True
        elif typ == 'start':
            if self.start_point:
                self.start_point.is_start = False
            p.is_start = True
            p.blocked = False
            self.start_point = p
        elif typ == 'exit':
            if self.exit_point:
                self.exit_point.is_exit = False
            p.is_exit = True
            p.blocked = False
            self.exit_point = p
        elif typ == 'path':
            p.blocked = False

    def neighbors(self, p):
        dirs = [(0, -1), (0, 1), (-1, 0), (1, 0)]
        res = []
        for dx, dy in dirs:
            nx, ny = p.x + dx, p.y + dy
            nb = self.get_point(nx, ny)
            if nb and nb.can_step():
                res.append(nb)
        return res

class MazeLoader:
    def load(self, filename):
        raise NotImplementedError

class TextMazeLoader(MazeLoader):
    def load(self, filename):
        with open(filename, 'r') as f:
            lines = [line.rstrip('\n') for line in f]
        h = len(lines)
        w = max(len(line) for line in lines) if h > 0 else 0
        start_cnt = 0
        exit_cnt = 0
        lab = Labyrinth(w, h)

        for y, line in enumerate(lines):
            for x, ch in enumerate(line):
                if ch == '#':
                    lab.set_point(x, y, 'wall')
                elif ch == 'S':
                    lab.set_point(x, y, 'start')
                    start_cnt += 1
                elif ch == 'E':
                    lab.set_point(x, y, 'exit')
                    exit_cnt += 1
                else:
                    lab.set_point(x, y, 'path')
        if start_cnt != 1 or exit_cnt != 1:
            raise ValueError(f"Need exactly one S and one E. Found S={start_cnt}, E={exit_cnt}")
        return lab

class SearchAlgorithm:
    def find_way(self, lab, start, goal):
        raise NotImplementedError

    def _build_path(self, prev, start, goal):
        path = []
        cur = goal
        while cur:
            path.append(cur)
            cur = prev.get(cur)
        path.reverse()
        return path

    def get_visited(self):
        return getattr(self, '_visited', 0)

class BreadthFirst(SearchAlgorithm):
    def find_way(self, lab, start, goal):
        q = deque([start])
        prev = {start: None}
        seen = {start}
        while q:
            cur = q.popleft()
            if cur == goal:
                self._visited = len(seen)
                return self._build_path(prev, start, goal)
            for nb in lab.neighbors(cur):
                if nb not in seen:
                    seen.add(nb)
                    prev[nb] = cur
                    q.append(nb)
        self._visited = len(seen)
        return []

class DepthFirst(SearchAlgorithm):
    def find_way(self, lab, start, goal):
        stack = [start]
        prev = {start: None}
        seen = {start}
        while stack:
            cur = stack.pop()
            if cur == goal:
                self._visited = len(seen)
                return self._build_path(prev, start, goal)
            for nb in lab.neighbors(cur):
                if nb not in seen:
                    seen.add(nb)
                    prev[nb] = cur
                    stack.append(nb)
        self._visited = len(seen)
        return []

class AStar(SearchAlgorithm):
    def _dist(self, a, b):
        return abs(a.x - b.x) + abs(a.y - b.y)

    def find_way(self, lab, start, goal):
        heap = []
        cnt = 0
        start_f = self._dist(start, goal)
        heapq.heappush(heap, (start_f, cnt, start))
        cnt += 1
        prev = {}
        g = {start: 0}
        f = {start: start_f}
        seen = set()
        while heap:
            cur_f, _, cur = heapq.heappop(heap)
            seen.add(cur)
            if cur == goal:
                self._visited = len(seen)
                return self._build_path(prev, start, goal)
            if cur_f > f.get(cur, float('inf')):
                continue
            for nb in lab.neighbors(cur):
                new_g = g[cur] + 1
                if new_g < g.get(nb, float('inf')):
                    prev[nb] = cur
                    g[nb] = new_g
                    new_f = new_g + self._dist(nb, goal)
                    f[nb] = new_f
                    heapq.heappush(heap, (new_f, cnt, nb))
                    cnt += 1
        self._visited = len(seen)
        return []

class LabyrinthSolver:
    def __init__(self, lab):
        self.lab = lab
        self.algorithm = None

    def set_algorithm(self, algo):
        self.algorithm = algo

    def solve(self):
        if not self.algorithm:
            return None
        t0 = time.perf_counter()
        path = self.algorithm.find_way(self.lab, self.lab.start_point, self.lab.exit_point)
        t1 = time.perf_counter()
        ms = (t1 - t0) * 1000
        return ms, self.algorithm.get_visited(), len(path)

class Player:
    def __init__(self, start, lab):
        self.current = start
        self.last = None
        self.lab = lab

    def move(self, cell):
        if cell and cell.can_step():
            self.last = self.current
            self.current = cell
            return True
        return False

    def undo(self):
        if self.last:
            self.current, self.last = self.last, None
            return True
        return False

class Command:
    def do(self):
        raise NotImplementedError
    def revert(self):
        raise NotImplementedError

class MoveCommand(Command):
    def __init__(self, player, dx, dy, lab):
        self.player = player
        self.dx = dx
        self.dy = dy
        self.lab = lab
        self.done = False

    def do(self):
        nx = self.player.current.x + self.dx
        ny = self.player.current.y + self.dy
        target = self.lab.get_point(nx, ny)
        if target and target.can_step():
            self.player.move(target)
            self.done = True
            return True
        return False

    def revert(self):
        if self.done:
            self.player.undo()
            self.done = False
            return True
        return False

class InteractiveView:
    def __init__(self, lab, player):
        self.lab = lab
        self.player = player

    def render(self):
        os.system('cls' if os.name == 'nt' else 'clear')
        print("=" * (self.lab.w * 2 + 4))
        print("   LABYRINTH (P = player)")
        print("=" * (self.lab.w * 2 + 4))
        for y in range(self.lab.h):
            print("  ", end='')
            for x in range(self.lab.w):
                p = self.lab.get_point(x, y)
                if self.player.current == p:
                    print('P', end=' ')
                elif p == self.lab.start_point:
                    print('S', end=' ')
                elif p == self.lab.exit_point:
                    print('E', end=' ')
                elif p.blocked:
                    print('#', end=' ')
                else:
                    print('.', end=' ')
            print()
        print("=" * (self.lab.w * 2 + 4))
        print(f"  Position: ({self.player.current.x},{self.player.current.y})")
        print("  Controls: h(left) j(down) k(up) l(right)  u=undo  q=quit")
        print("  Auto-search: b=BFS  d=DFS  a=A*")

def run_experiment(maze_file, algo, runs=5):
    loader = TextMazeLoader()
    lab = loader.load(maze_file)
    total_ms = 0
    total_visited = 0
    total_len = 0
    for _ in range(runs):
        solver = LabyrinthSolver(lab)
        solver.set_algorithm(algo)
        stats = solver.solve()
        if stats:
            ms, vis, plen = stats
            total_ms += ms
            total_visited += vis
            total_len += plen
    return total_ms / runs, total_visited / runs, total_len / runs

def generate_plots(results):
    mazes = list(set([r['maze'] for r in results]))
    strategies = ['BFS', 'DFS', 'AStar']
    
    fig, axes = plt.subplots(1, 3, figsize=(15, 5))
    x = np.arange(len(mazes))
    width = 0.25

    for i, strat in enumerate(strategies):
        times = []
        for maze in mazes:
            val = next((r['time_ms'] for r in results if r['maze'] == maze and r['strategy'] == strat), 0)
            times.append(val)
        axes[0].bar(x + i*width, times, width, label=strat)
    axes[0].set_xlabel('Maze')
    axes[0].set_ylabel('Time (ms)')
    axes[0].set_title('Execution Time')
    axes[0].set_xticks(x + width)
    axes[0].set_xticklabels(mazes, rotation=45, ha='right')
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)

    for i, strat in enumerate(strategies):
        visited = []
        for maze in mazes:
            val = next((r['visited_cells'] for r in results if r['maze'] == maze and r['strategy'] == strat), 0)
            visited.append(val)
        axes[1].bar(x + i*width, visited, width, label=strat)
    axes[1].set_xlabel('Maze')
    axes[1].set_ylabel('Visited Cells')
    axes[1].set_title('Visited Cells')
    axes[1].set_xticks(x + width)
    axes[1].set_xticklabels(mazes, rotation=45, ha='right')
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)

    for i, strat in enumerate(strategies):
        lengths = []
        for maze in mazes:
            val = next((r['path_length'] for r in results if r['maze'] == maze and r['strategy'] == strat), 0)
            lengths.append(val)
        axes[2].bar(x + i*width, lengths, width, label=strat)
    axes[2].set_xlabel('Maze')
    axes[2].set_ylabel('Path Length')
    axes[2].set_title('Path Length')
    axes[2].set_xticks(x + width)
    axes[2].set_xticklabels(mazes, rotation=45, ha='right')
    axes[2].legend()
    axes[2].grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig('performance_comparison.png', dpi=150, bbox_inches='tight')
    plt.show()

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == 'experiment':
        print("Running experiments on all mazes...")
        maze_files = [
            ("maze/maze1.txt", "Small 10x6"),
            ("maze/maze10x10.txt", "Medium 10x10"),
            ("maze/maze20x20.txt", "Large 20x20"),
            ("maze/maze_empty.txt", "Empty 15x15"),
            ("maze/maze_no_exit.txt", "No exit 10x10")
        ]
        algorithms = [
            ("BFS", BreadthFirst()),
            ("DFS", DepthFirst()),
            ("AStar", AStar())
        ]
        results = []
        for fname, label in maze_files:
            print(f"Testing {label}...")
            for aname, algo in algorithms:
                try:
                    avg_t, avg_v, avg_l = run_experiment(fname, algo, runs=3)
                    results.append({
                        'maze': label,
                        'strategy': aname,
                        'time_ms': avg_t,
                        'visited_cells': avg_v,
                        'path_length': avg_l
                    })
                    print(f"  {aname}: time={avg_t:.3f}ms visited={avg_v:.0f} length={avg_l:.0f}")
                except Exception as e:
                    print(f"  {aname}: ERROR {e}")
        # save csv
        with open('experiment_results.csv', 'w', newline='', encoding='utf-8') as f:
            writer = csv.DictWriter(f, fieldnames=['maze', 'strategy', 'time_ms', 'visited_cells', 'path_length'])
            writer.writeheader()
            writer.writerows(results)
        generate_plots(results)
        print("Done. Results saved to experiment_results.csv and performance_comparison.png")
        sys.exit(0)

    # else interactive mode
    loader = TextMazeLoader()
    lab = loader.load("maze/maze1.txt")
    player = Player(lab.start_point, lab)
    view = InteractiveView(lab, player)
    view.render()

    solver = LabyrinthSolver(lab)
    history = []

    while True:
        key = input("\n  > ").lower()
        if key == 'q':
            print("Goodbye!")
            break
        elif key == 'b':
            solver.set_algorithm(BreadthFirst())
            ms, vis, plen = solver.solve()
            print(f"BFS: {ms:.3f}ms, visited={vis}, length={plen}")
        elif key == 'd':
            solver.set_algorithm(DepthFirst())
            ms, vis, plen = solver.solve()
            print(f"DFS: {ms:.3f}ms, visited={vis}, length={plen}")
        elif key == 'a':
            solver.set_algorithm(AStar())
            ms, vis, plen = solver.solve()
            print(f"A*: {ms:.3f}ms, visited={vis}, length={plen}")
        elif key in ('h','j','k','l'):
            moves = {'h': (-1,0), 'l': (1,0), 'k': (0,-1), 'j': (0,1)}
            dx, dy = moves[key]
            cmd = MoveCommand(player, dx, dy, lab)
            if cmd.do():
                history.append(cmd)
                view.render()
                if player.current == lab.exit_point:
                    print("\n*** YOU REACHED THE EXIT! ***")
                    print(f"Total moves: {len(history)}")
                    break
            else:
                print("Can't go there - wall!")
        elif key == 'u':
            if history:
                cmd = history.pop()
                cmd.revert()
                view.render()
                print("Undo last move")
            else:
                print("Nothing to undo")
        else:
            print("Unknown command")