import sys

import matplotlib.pyplot as plt

import numpy as np

import scipy.ndimage as filt

import torch.nn as nn

from utils import mask_generator

class BezierPolypExtender(nn.Module):

    def __init__(self, num_nodes, degree, minimum_distance=50, maximum_distance=100):

        super(BezierPolypExtender, self).__init__()

        self.num_nodes = num_nodes

        self.degree = degree

        self.minimum_distance = minimum_distance

        self.maximum_distance = maximum_distance

        # recursion depth is often exceeded despite low memory usage.

        sys.setrecursionlimit(10000)

    def get_distances_along_edge_from_seed(self, binary_edge_image, current_coord, out, iter=0):

        if iter > self.maximum_distance + self.minimum_distance:

            return out

        for xoff in (-1, 0, 1):

            for yoff in (-1, 0, 1):

                try:

                    if binary_edge_image[current_coord[0] + xoff, current_coord[1] + yoff] != 0 and out[

                        current_coord[0] + xoff, current_coord[1] + yoff] == 0:

                        iter += 1

                        out[current_coord[0] + xoff, current_coord[1] + yoff] = iter

                        self.get_distances_along_edge_from_seed(binary_edge_image,

                                                                [current_coord[0] + xoff, current_coord[1] + yoff], out,

                                                                iter)

                except IndexError:

                    print("continuing...")

                    continue

        return out

    def forward(self, original_mask):

        # select seed from edge pixels

        edges = ((filt.sobel(original_mask, axis=-1) ** 2 + filt.sobel(original_mask, axis=-2) ** 2) != 0).astype(int)

        edge_indexes = np.argwhere(edges == 1)

        plt.imshow(edges)

        plt.title(np.unique(edges))

        plt.savefig("edges")

        seed = edge_indexes[np.random.choice(range(edge_indexes.shape[0]), 1)][0]

        # generate edge image with proximity to seed

        proximity_image = np.zeros_like(edges)  # todo change to distance via contour

        proximity_image = self.get_distances_along_edge_from_seed(edges, seed, proximity_image)

        # i = 0

        # current_coords = seed.copy()

        # # todo perform thining ahead of iteration over contour to prevent gettign stuck

        # retrace_node = seed.copy()

        # while i < self.minimum_distance + self.maximum_distance:

        #     found_new = False

        #     print(f"{i}/{self.minimum_distance + self.maximum_distance}")

        #     for xoff in (-1, 0, 1):

        #         for yoff in (-1, 0, 1):

        #             if proximity_image[current_coords[0] + xoff, current_coords[1] + yoff] == 0 and edges[

        #                 current_coords[0] + xoff, current_coords[1] + yoff] != 0:

        #                 i += 1

        #                 proximity_image[current_coords[0] + xoff, current_coords[1] + yoff] = i

        #                 retrace_node = current_coords

        #                 current_coords = [current_coords[0] + xoff, current_coords[1] + yoff]

        #                 found_new = True

        #     if not found_new:

        #         print("failed to find new!")

        #         current_coords = retrace_node

        #         plt.imshow(edges, alpha=0.5)

        #         plt.imshow(proximity_image, alpha=0.5)

        #         plt.savefig("wtf.png")

        #         plt.show()

        plt.imshow(proximity_image)

        plt.show()  # for x in np.arange(proximity_image.shape[0]):

        #     for y in np.arange(proximity_image.shape[1]):

        #         if edges[x, y] == 1:

        #             proximity_image[x, y] = np.linalg.norm(seed - np.array([[x, y]]))

        # convert to pdf

        pdf = (np.max(proximity_image) - proximity_image)

        pdf[proximity_image < self.minimum_distance] = 0  # controls minimum abberation size

        pdf = pdf / np.sum(pdf)  # normalize

        ac_idx = np.argwhere(pdf != 0)

        probs = pdf[ac_idx[:, 0], ac_idx[:, 1]]

        anchorpoint = ac_idx[np.random.choice(range(ac_idx.shape[0]), 1, p=probs)]

        plt.imshow(pdf)

        plt.colorbar()

        plt.scatter(y=seed[:, 0], x=seed[:, 1], marker="X")

        plt.scatter(y=anchorpoint[:, 0], x=anchorpoint[:, 1], marker="o")

        plt.show()

class RandomDraw(nn.Module):

    def __init__(self):

        super(RandomDraw, self).__init__()

    def forward(self, rad):

        return mask_generator.generate_a_mask(rad=rad)