import numpy as np
import matplotlib.pyplot as plt
import skimage

plt.close('all')
I_ini = plt.imread('./sign2.png')
plt.figure()
plt.subplot(121); plt.imshow(I_ini)

#Resize the image (to compute faster. skimage.transform.resize)
#I_ini = skimage.transform.resize(I_ini, (200,200))

#Contour detection (skimage.feature.canny)
I = skimage.color.rgb2gray(I_ini)
C = skimage.feature.canny(I)
plt.subplot(122); plt.imshow(C)
plt.show()


#%%
#Hough parameter space (define R range)
#Follow algorithm on the slide
r_range = np.arange(38,58)
H = np.zeros((I.shape[0],I.shape[1],len(r_range)))
h, w = I.shape

ind_r = 0
for r in r_range:
    print(r)
    for i in range(0,h):
    	for j in range(0,w):
    		if (C[i,j] == 1):
    			for a in range(j-r, j+r):
    				if ( (a<0) or (a>=w)):
    					continue
    				b1 = i - int(np.sqrt(r**2 - (j - a)**2))
    				if (b1>=0):
    					H[b1,a,ind_r] += 1
    				b2 = i + int(np.sqrt(r**2 - (j - a)**2))
    				if (b2<h):
    					H[b2,a,ind_r] += 1
    ind_r += 1

#%%			
plt.figure()
for i in range(H.shape[2]):
    plt.imshow(H[:,:,i], vmin=np.min(H), vmax=np.max(H))
    plt.pause(0.25)
    plt.title(f"{r_range[i]}")	
    plt.show()	
    
#%%
#Non maximum suppression (NMS), 3x3x3 window (np.ravel, np.argmax)
#NMS: Run through the parameter space volume.
#If a value is not the maximum in a 3x3x3 windows, set it to 0
H_nms = H*0
l = 2
for r in range(len(r_range)):
    print(r)
    for i in range(0,h):
        for j in range(0,w):
            win = H[max(0,i-l):min(h-1,i+l+1),
                    max(0,j-l):min(w-1,j+l+1),
                    max(0,r-l):min(len(r_range)-1,r+l+1)]
            if (np.max(win) == H[i,j,r]):
                H_nms[i,j,r] = H[i,j,r]


#Find maximum values (np.ravel, np.argsort, np.unravel_index ?)
pos_ = np.argsort(H_nms.ravel())
[b,a,r] = np.unravel_index(pos_[-1], (H.shape))
[b2,a2,r2] = np.unravel_index(pos_[-2], (H.shape))

#%%
#%%			
plt.figure()
for i in range(H.shape[2]):
    plt.imshow(H_nms[:,:,i], vmin=np.min(H_nms), vmax=np.max(H_nms))
    plt.pause(0.5)
    plt.title(f"{r_range[i]}")	
    plt.show()	
    

#%%
#Display circles on image
I_circle = np.copy(I_ini)
if (1):
	X, Y = np.meshgrid(range(0,w), range(0,h))
	epsilon = 50
	R1 = abs(((X-a)**2 + (Y-b)**2 - r_range[r]**2)) < epsilon
	R2 = abs(((X-a2)**2 + (Y-b2)**2 - r_range[r2]**2)) < epsilon
	for i in range(0,h):
		for j in range(0,w):
			if (R1[i,j]):
				I_circle[i,j,0] = 0
				I_circle[i,j,2] = 0
				I_circle[i,j,1] = 1
			if (R2[i,j]):
				I_circle[i,j,0] = 0
				I_circle[i,j,2] = 1
				I_circle[i,j,1] = 0

	plt.figure()
	plt.imshow(I_circle)
	plt.show()