#!/usr/bin/python
# -*- coding: utf-8 -*-
import numpy as np
import cv2

sub_paths = []
visited = set()
current_path = []
sk = None

DO_PRINT_DEBUG = False


def clear():
    sub_paths.clear()
    visited.clear()
    current_path.clear()


def get_neighbors(point):
    """Get the neighboring points of a given point in the skeleton."""
    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 < sk.shape[0] and 0 <= ny < sk.shape[1]:
                if sk[nx, ny] == 255:  # Check if it's part of the skeleton
                    neighbors.append((nx, ny))
    return neighbors


def add_a_point(p):
    if p in visited:
        return
    if DO_PRINT_DEBUG:
        print(f'Adding: {p}')
    visited.add(p)
    current_path.append(p)
    neighbors = [pix for pix in get_neighbors(p) if pix not in visited]
    branch_count = len(neighbors)
    # print(f'LN: {len(neighbors)} {neighbors}')
    while branch_count == 1:
        new_point = neighbors[0]
        visited.add(new_point)
        current_path.append(new_point)
        neighbors = [pix for pix in get_neighbors(new_point) if pix not in visited]
        branch_count = len(neighbors)
        # print(f'LN: {len(neighbors)} {neighbors}')

    sub_paths.append(current_path.copy())
    current_path.clear()
    for nei in neighbors:
        add_a_point(nei)


def compute_all_the_sub_paths(skeleton):
    global sk
    sk = skeleton
    clear()
    skeleton_points = np.argwhere(skeleton == 255)  # Get all skeleton points
    for skp in skeleton_points:
        pto = tuple(skp)
        add_a_point(pto)
    return sub_paths


def draw_paths(image_shape, sub_paths):
    """Draw the remaining paths on a blank image."""
    drawn_image = np.zeros(image_shape, dtype=np.uint8)  # Create a blank image
    for path in sub_paths:
        for point in path:
            drawn_image[point] = 255  # Draw the path in white
    return drawn_image


def give_me_the_main_loop_only(input_skeleton):
    # Extract sub-paths
    sub_paths = compute_all_the_sub_paths(input_skeleton)
    end_points = {}

    def is_close_enough(p1, p2):
        x, y = p1
        px, py = p2
        dx = abs(x - px)
        dy = abs(y - py)
        if dx <= 3 and dy <= 3:
            return True
        return False

    def find_close_point(pto):
        for p in end_points.keys():
            if is_close_enough(pto, p):
                return p
        return None

    def add_endpoint(pto, idx):
        where = find_close_point(pto)
        if where is None:
            # no close point found, add it
            end_points[pto] = []
            where = pto
        end_points[where].append(idx)

    # Print the results
    to_be_removed = set()
    for i, path in enumerate(sub_paths):
        # if len(path) < 2:
        #     continue
        begin = path[0]
        end = path[-1]
        if is_close_enough(begin, end):
            to_be_removed.add(i)
            continue
        add_endpoint(begin, i)
        add_endpoint(end, i)
        if DO_PRINT_DEBUG:
            print(f"Sub-path {i}: {len(path):4}  {path[0]},{path[-1]}")

    # Draw the remaining paths on a blank image
    # to_be_removed = {0, 7, 12, 13}
    # to_be_removed = set()
    for n, (cc, ep) in enumerate(end_points.items()):
        pto = int(cc[0]), int(cc[1])
        if DO_PRINT_DEBUG:
            print(f'Endpoint #{n} {pto}: {len(ep)}, {ep}')
        if len(ep) == 1:
            # this is a dead branch
            to_be_removed.add(ep[0])
    if DO_PRINT_DEBUG:
        print(f'{to_be_removed=}')
    sub_paths_remaining = [sp for n, sp in enumerate(sub_paths) if n not in to_be_removed]

    return draw_paths(input_skeleton.shape, sub_paths_remaining)


if __name__ == '__main__':
    skeleton_image_path = 'sk_bonk.png'  # Update with your image path
    skeleton = cv2.imread(skeleton_image_path, cv2.IMREAD_GRAYSCALE)

    # Threshold the image to ensure it's binary
    _, skeleton = cv2.threshold(skeleton, 127, 255, cv2.THRESH_BINARY)

    cv2.imshow('Input', skeleton)

    # sp = compute_all_the_sub_paths(skeleton=skeleton)
    # print(len(sp))
    # for n, p in enumerate(sp):
    #     print(f'{n:3} {len(p)}')

    loop_only = give_me_the_main_loop_only(skeleton)
    cv2.imshow('Output', loop_only)


    cv2.waitKey(0)
    cv2.destroyAllWindows()