import heapq
import numpy as np
import math
from collections import defaultdict
import matplotlib.pyplot as plt
import seaborn as sns


def euclidean_distance(point1, point2):
    return math.sqrt((point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2 + (point1[2] - point2[2]) ** 2)


def chebyshev_distance(point1, point2):
    return max(abs(point1[0] - point2[0]), abs(point1[1] - point2[1]))


def a_star(start, end, grid, neighbors, distance_fn, heuristic_fn):
    open_list = []
    heapq.heappush(open_list, (0, start))
    came_from = dict()
    g_score = defaultdict(lambda: float("inf"))
    g_score[start] = 0
    f_score = defaultdict(lambda: float("inf"))
    f_score[start] = heuristic_fn(start, end)

    while open_list:
        current = heapq.heappop(open_list)[1]
        if current == end:
            return reconstruct_path(came_from, current)

        for neighbor in neighbors(current, grid):
            tentative_g_score = g_score[current] + distance_fn(current, neighbor)
            if tentative_g_score < g_score[neighbor]:
                came_from[neighbor] = current
                g_score[neighbor] = tentative_g_score
                f_score[neighbor] = g_score[neighbor] + heuristic_fn(neighbor, end)
                heapq.heappush(open_list, (f_score[neighbor], neighbor))

    return []


def reconstruct_path(came_from, current):
    path = [current]
    while current in came_from:
        current = came_from[current]
        path.append(current)
    return path[::-1]


def neighbors_3d(point, grid):
    x, y, _ = point
    neighbors = []
    for dx in (-1, 0, 1):
        for dy in (-1, 0, 1):
            if dx == 0 and dy == 0:
                continue
            nx, ny = x + dx, y + dy
            if 0 <= nx < len(grid) and 0 <= ny < len(grid[0]) and grid[nx][ny][3] != 1:
                neighbors.append((nx, ny, grid[ny][nx][2]))
    return neighbors


def read_file(file_name):
    with open(file_name, "r") as file:
        lines = file.readlines()
    return [tuple(map(float, line.strip().split())) for line in lines]


def main():
    surface_data = read_file("surface_100x100.txt")
    end_points_data = read_file("surface_100x100.end_points.txt")

    grid = np.reshape(surface_data, (100, 100, 4))
    start_xy, end_xy = end_points_data[0][:2], end_points_data[1][:2]

    start = (int(start_xy[0]), int(start_xy[1]), grid[int(start_xy[1])][int(start_xy[0])][2])
    end = (int(end_xy[0]), int(end_xy[1]), grid[int(end_xy[1])][int(end_xy[0])][2])

    path_length = a_star(start, end, grid, neighbors_3d, euclidean_distance, euclidean_distance)
    path_steps = a_star(start, end, grid, neighbors_3d, chebyshev_distance, chebyshev_distance)

    with open("output.txt", "w") as file:
        file.write("Minimum Length Path:\n")
        file.write("Cost: {:.2f}\n".format(sum(euclidean_distance(path_length[i], path_length[i + 1])
                                           for i in range(len(path_length) - 1))))
        file.write("Path: {}\n".format(" ".join(str(point[:2]) for point in path_length)))

        file.write("\nMinimum Steps Path:\n")
        file.write("Cost: {}\n".format(len(path_steps) - 1))
        file.write("Path: {}\n".format(" ".join(str(point[:2]) for point in path_steps)))

    # 3D plot
    fig = plt.figure()
    ax = fig.add_subplot(111, projection="3d")
    xs, ys, zs, bs = zip(*surface_data)
    ax.scatter(xs, ys, zs, c=["red" if b == 1 else "blue" for b in bs], marker="o")

    # Euclidean
    x_path_length, y_path_length, z_path_length = zip(*path_length)
    ax.plot(x_path_length, y_path_length, z_path_length, c="lime", linewidth=2, zorder=10)

    # Distance
    x_path_steps, y_path_steps, z_path_steps = zip(*path_steps)
    ax.plot(x_path_steps, y_path_steps, z_path_steps, c="yellow", linewidth=2, zorder=10)

    ax.set_xlabel("X")
    ax.set_ylabel("Y")
    ax.set_zlabel("Z")
    plt.show()

    sns.heatmap([[grid[x][y][3] for y in range(100)] for x in range(100)], cmap="coolwarm", annot=False, cbar=True)
    plt.show()

    fig, ax = plt.subplots()
    ax.imshow([[grid[x][y][3] for y in range(100)] for x in range(100)], cmap="coolwarm", origin="lower")

    # Euclidean
    x_path_length, y_path_length, _ = zip(*path_length)
    ax.plot(y_path_length, x_path_length, c="lime", linewidth=2)

    # Distance
    x_path_steps, y_path_steps, _ = zip(*path_steps)
    ax.plot(y_path_steps, x_path_steps, c="yellow", linewidth=2)

    plt.show()


if __name__ == "__main__":
    main()
