import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
import cv2
import progressbar

N_MAX = 4
INPUT_FILE = 'curve.png'
# INPUT_FILE = 'hough-lane-detect-python/canny_output/cannylane_test.jpg'

def forceAspect(ax,aspect):
    xleft, xright = ax.get_xlim()
    ybottom, ytop = ax.get_ylim()
    ax.set_aspect(abs((xright-xleft)/(ybottom-ytop))*aspect)


def plot_accumulator(accumulator):
    # Create the x, y, and z coordinate arrays.  We use
    # numpy's broadcasting to do all the hard work for us.
    # We could shorten this even more by using np.meshgrid.
    x = np.arange(accumulator.shape[0])[:, None, None]
    y = np.arange(accumulator.shape[1])[None, :, None]
    z = np.arange(accumulator.shape[2])[None, None, :]
    x, y, z = np.broadcast_arrays(x, y, z)
    c = np.tile(accumulator.ravel()[:, None], [1, 3])

    ax = plt.axes(projection='3d')
    ax.scatter(x.ravel(), y.ravel(), z.ravel(), c=accumulator.ravel());
    plt.show()


def plot_curve(img,k,beta,v,ax_img):
    h,w = img.shape
    x_vals = np.linspace(-w,w,2*w)
    y_vals = k/x_vals + beta*x_vals + v
    ax_img.plot(x_vals,y_vals,'k',color='firebrick',alpha=0.5)
    return ax_img

def do_hough(img):

    print("-------------------------------------")

    small_to_large_image_size_ratio = 0.3
    img = cv2.resize(img,
                       (0,0), # set fx and fy, not the final size
                       fx=small_to_large_image_size_ratio,
                       fy=small_to_large_image_size_ratio,
                       interpolation=cv2.INTER_LINEAR)

    img = np.flipud(img) # vertical flip

    h,w = img.shape
    diag = np.ceil(np.hypot(h,w))
    print(f"IMG dimensions: {img.shape} max. intensity: {np.max(img)}")

    beta_min = 0
    beta_max = 10
    k_min = -2000
    k_max = 0
    betas = np.linspace(0,3,21)
    # betas = [2,3]
    # print(betas)
    ks = np.linspace(k_min,k_max,101)
    # ks = [-1000,-1200,-1500]

    vs = np.arange(0,2*h+1)
    accumulator = np.zeros((len(betas),len(ks),len(vs)), dtype=np.uint64)

    for y in progressbar.progressbar(range(1,h-1)):
        for x in range(1,w-1):
            if img[y,x] > 0:  # if we're on an edge
                # print(i,j)
                for beta_i in range(len(betas)):
                    for k_i in range(len(ks)):
                        v = np.int64(np.round(y - ks[k_i]/x - betas[beta_i]*x))
                        # print(f"x:{x} | y:{y} | beta:{betas[beta_i]} | k:{ks[k_i]} | v:{v}")
                        if v >= -h and v < h:
                            accumulator[beta_i,k_i,v+h] += 1  # increment accumulator for this coordinate pair



    fig = plt.figure()

    ax_img = fig.add_subplot(111)  # original image
    ax_img.imshow(np.flipud(img),cmap='gray',extent=[0, w, 0, h])


    for i in range(N_MAX):
        max = np.unravel_index(accumulator.argmax(), accumulator.shape)
        print(np.max(accumulator),"at",max)
        print("k",ks[max[1]]," | beta",betas[max[0]]," | v",vs[max[2]])

        # plot_accumulator(accumulator)
        ax_img = plot_curve(np.flipud(img),ks[max[1]],betas[max[0]],max[2]-h,ax_img)

        accumulator[max[0]][max[1]][max[2]] = 0


    ax_img.set_ylim(0,h)
    ax_img.set_xlim(0,w)
    forceAspect(ax_img,0.5)

    # plt.savefig('curve.png')
    # ax_img.invert_yaxis()
    plt.show()


# open test img as grayscale array
img = np.array(Image.open(INPUT_FILE).convert("L"))

# img = np.zeros((250,250))
# img[120,10] = 255
# img[190,20] = 255
# img[210,25] = 255
# img[91,8] = 255
do_hough(img)