import cv2
import numpy as np
from collections import deque
from queue import PriorityQueue


class A_starSearch:
    def __init__(self, src_img, s, e,wall_distance_weight=10,min_safe_distance=15):
        assert len(src_img.shape) == 2
        self.h, self.w = src_img.shape
        self.src_img = src_img
        self.mark = np.zeros_like(src_img)
        self.edge_to = np.full((self.h, self.w, 2), -1, dtype=np.int32)
        self.s = s
        self.e = e
        self.wall_distance = self.calculate_wall_distance(src_img)
        self.wall_distance_weight = wall_distance_weight
        self.min_safe_distance = min_safe_distance
        self.a_star()


    def calculate_wall_distance(self,img):
        dist_transform = cv2.distanceTransform(img, cv2.DIST_L2, 5)
        return dist_transform

    def get_neighbors(self, px, py):
        nbs = []
        directions = [(-1, -1), (-1, 0), (-1, 1),
                      (0, -1), (0, 1),
                      (1, -1), (1, 0), (1, 1)]

        for dx, dy in directions:
            nx, ny = px + dx, py + dy
            if 0 <= nx < self.w and 0 <= ny < self.h:
                nbs.append((nx, ny))
        return nbs

    def path_to(self, e=None):
        if e is None:
            e = self.e
        e_x, e_y = e

        queue_path = deque()
        while (e_x, e_y) != self.s:
            queue_path.append((e_x, e_y))

            if not self.mark[e_y,e_x]:
                print(f'can not get to {e_x, e_y}')
                return []
            e_x, e_y = self.edge_to[e_y,e_x][0],self.edge_to[e_y,e_x][1]

        queue_path.append((e_x, e_y))
        return list(reversed(queue_path))

    def a_star(self):
        assert self.s != self.e
        pq = PriorityQueue()
        initial_heuristic = self.heuristic(self.s, self.e)
        pq.put((initial_heuristic,0, self.s))
        cost_so_far = {self.s: 0}
        self.mark[self.s[1],self.s[0]] = 1

        while not pq.empty():
            _,current_cost, (px, py) = pq.get()

            if (px, py) == self.e:
                break

            for nb_x, nb_y in self.get_neighbors(px, py):

                if self.src_img[nb_y,nb_x] == 0:
                    continue
                dx, dy = abs(nb_x - px), abs(nb_y - py)
                if dx == 1 and dy == 1:
                    base_cost = np.sqrt(2)
                else:
                    base_cost = 1.0

                distance_to_wall = self.wall_distance[nb_y, nb_x]
                if distance_to_wall > self.min_safe_distance:
                    distance_cost = self.wall_distance_weight*0.01
                else:
                    distance_cost = self.wall_distance_weight * (1.0 / (distance_to_wall + 0.1))

                new_cost = current_cost + base_cost + distance_cost

                heuristic = self.heuristic((nb_x, nb_y), self.e)
                priority = new_cost + heuristic

                if (nb_x, nb_y) not in cost_so_far or new_cost < cost_so_far[(nb_x, nb_y)]:
                    cost_so_far[(nb_x, nb_y)] = new_cost
                    self.edge_to[nb_y,nb_x] = [px, py]
                    self.mark[nb_y,nb_x] = 1
                    pq.put((priority,new_cost,(nb_x, nb_y)))

    def heuristic(self, current, goal):
        dx = abs(current[0] - goal[0])
        dy = abs(current[1] - goal[1])
        return 1 * (dx + dy) + (np.sqrt(2) - 2 * 1) * min(dx, dy)


dest_img_copy = None
start_p = None
end_p = None


def mouse_callback(event, x, y, flags, param):
    global dest_img_copy, start_p, end_p

    if event == cv2.EVENT_LBUTTONDOWN:
        print(dest_img_copy[y, x])

        if dest_img_copy[y, x].tolist() == [0, 0,0]:
            raise ValueError(f'{x, y} is black!')

        if start_p is None:
            start_p = (x, y)
            print(f'start_p: {x, y}')
        elif end_p is None:
            end_p = (x, y)
            print(f'end_p: {x, y}')



if __name__ == '__main__':
    cap = cv2.VideoCapture(0)
    while True:
        ret, frame = cap.read()
        frame = cv2.flip(frame, 1)
        cv2.imshow('video',  frame)
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        k = cv2.waitKey(100)
        if k == 13:
            break
    cap.release()
    cv2.destroyWindow('video')

    _, img = cv2.threshold(img, 50, 255, cv2.THRESH_BINARY)

    cv2.namedWindow('img', cv2.WINDOW_NORMAL)
    cv2.setMouseCallback("img", mouse_callback)
    inverted_img = cv2.bitwise_not(img)
    dest_img_ = cv2.morphologyEx(inverted_img, cv2.MORPH_CLOSE, np.ones((8, 8), dtype=np.uint8))
    dest_img= cv2.bitwise_not(dest_img_)

    dest_img_copy = dest_img.copy()[:, :, None] * np.ones((1, 1, 3), dtype=np.uint8)

    cv2.imshow('img', dest_img_copy)

    while start_p is None:
        cv2.waitKey(1)

    while end_p is None:
        cv2.waitKey(1)

    search = A_starSearch(dest_img, start_p, end_p)
    path_to_end = search.path_to()
    if path_to_end:
        show_img = img[:, :, None] * np.ones((1, 1, 3), dtype=np.uint8)
        print(f'path_to_end: {path_to_end}')
        for x, y in path_to_end:
            cv2.circle(show_img, (x, y), 1, (0, 0, 255), 1, cv2.LINE_AA)

        cv2.namedWindow('show_img', cv2.WINDOW_NORMAL)
        cv2.imshow('show_img', show_img)
        cv2.waitKey()

    else:
        raise ValueError('cat not find the path')

    cap.release()
    cv2.destroyAllWindows()