from scipy import ndimage
from skimage import filters,feature
import sys
import matplotlib.pyplot as plt
import math
import numpy as np
import pdb

def calculate_edges(img):
    # TODO fix the performance - added ad hoc
    (n,m) = img.shape
    edge = np.zeros((n,m), dtype=np.bool)
#    pdb.set_trace()
    for i in range(0, n):
        for j in range(0, m):
            edge_point = False
            for k in range(-1, 2):
                if edge_point:
                    break
                for l in range(-1,2):
                    if edge_point:
                        break

                    y = (i+k)%n
                    x = (j+l)%m

                    if img[y,x] != img[i,j]:
                        edge_point = True
                        break
            edge[i,j] = edge_point
    return edge

def calc_shape_corners(img_mask2, range_ext=8):
    (my,mx) = ndimage.center_of_mass(img_mask2)
#    edge = feature.canny(img_mask2)
    edge = calculate_edges(img_mask2)

    (n,m) = img_mask2.shape

    print("Center of mass {}x{}".format(mx,my))
    distances = np.zeros((n,m), dtype=np.float)
    distances[:,:] = 0
    for i in range(0, n):
        for j in range(0, m):
            if edge[i,j]:
                dx = j - mx
                dy = i - my
                d = math.sqrt(dx*dx+dy*dy)
                distances[i,j] = d

    max_distances = ndimage.maximum_filter(distances, range_ext)
    diff = np.abs(max_distances - distances)
    corners = edge * (diff < 1e-3)
    return corners

if __name__ == "__main__":
    img = ndimage.imread(sys.argv[1])
    img_gray = 0.21 * img[:,:,0] + 0.72 * img[:,:,1] + 0.07 * img[:,:,2]
    otsu_lvl = filters.threshold_otsu(img_gray)
    img_mask = img_gray <= otsu_lvl
    img_mask2 = ndimage.median_filter(img_mask, 3)
    (n,m) = img_mask2.shape


    corners = calc_shape_corners(img_mask2)

    mat_x = np.tile(np.arange(0,m),n).reshape((n,m))
    mat_y = np.tile(np.arange(0,n),m).reshape((m,n)).T

    xs = mat_x[corners]
    ys = mat_y[corners]

    plt.imshow(img_gray, cmap="Greys_r")
    plt.plot(xs, ys, "r+")
    plt.show()