import cv2
import matplotlib.pyplot as plt
import numpy as np
import sys


def halftone(image_path, thresh=3, radius=10, savefig = True, halftone_path='halftone.jpeg', BGR=True, dotsize=30):
    image = cv2.imread(image_path)
    if BGR:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # image = cv2.resize(image, (0, 0), fx = 0.5, fy = 0.5)

    if image is None:
        print("Can not find any image. Choose appropriate file")
        sys.exit()
    k = 30
    r = radius
    height, width = image.shape[:2]

    a = r/np.sqrt(2)
    nx, ny = int(width / a) + 1, int(height / a) + 1

    coords_list = [(ix, iy) for ix in range(nx) for iy in range(ny)]
    cells = {coords: None for coords in coords_list}
    def get_cell_coords(pt):
        return int(pt[0] // a), int(pt[1] // a)

    def get_neighbours(coords):
        dxdy = [(-1,-2),(0,-2),(1,-2),(-2,-1),(-1,-1),(0,-1),(1,-1),(2,-1),
                (-2,0),(-1,0),(1,0),(2,0),(-2,1),(-1,1),(0,1),(1,1),(2,1),
                (-1,2),(0,2),(1,2),(0,0)]
        neighbours = []
        for dx, dy in dxdy:
            neighbour_coords = coords[0] + dx, coords[1] + dy
            if not (0 <= neighbour_coords[0] < nx and
                    0 <= neighbour_coords[1] < ny):
                # We're off the grid: no neighbours here.
                continue
            neighbour_cell = cells[neighbour_coords]
            if neighbour_cell is not None:
                # This cell is occupied: store this index of the contained point.
                neighbours.append(neighbour_cell)
        return neighbours

    def point_valid(pt):
        cell_coords = get_cell_coords(pt)
        for idx in get_neighbours(cell_coords):
            nearby_pt = samples[idx]
            # Squared distance between or candidate point, pt, and this nearby_pt.
            distance2 = (nearby_pt[0]-pt[0])**2 + (nearby_pt[1]-pt[1])**2
            if distance2 < r**2:
                # The points are too close, so pt is not a candidate.
                return False
        # All points tested: if we're here, pt is valid
        return True

    def get_point(k, refpt):
        i = 0
        while i < k:
            i += 1
            rho = np.sqrt(np.random.uniform(r**2, 4 * r**2))
            theta = np.random.uniform(0, 2*np.pi)
            pt = refpt[0] + rho*np.cos(theta), refpt[1] + rho*np.sin(theta)
            if not (0 <= pt[0] < width and 0 <= pt[1] < height):
                # This point falls outside the domain, so try again.
                continue
            if point_valid(pt):
                return pt

        # We failed to find a suitable point in the vicinity of refpt.
        return False
    ################################################
    ### Processs image to grayscale and get edges
    ### output: edge (think edge image)
    ################################################
    grayScaleImage = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    smoothGrayScale = cv2.medianBlur(grayScaleImage, 5)
    getEdge = cv2.adaptiveThreshold(smoothGrayScale, 255,
      cv2.ADAPTIVE_THRESH_MEAN_C,
      cv2.THRESH_BINARY, 9, 9)
    edge = cv2.medianBlur(getEdge,3)

    ################################################
    ### Make Poisson Disc Sampling
    ### output: samples
    ################################################
    pt = (np.random.uniform(0, width), np.random.uniform(0, height))
    samples = [pt]
    cells[get_cell_coords(pt)] = 0
    active = [0]

    nsamples = 1
    while active:
        idx = np.random.choice(active)
        refpt = samples[idx]
        pt = get_point(k, refpt)
        if pt:
            samples.append(pt)
            nsamples += 1
            active.append(len(samples)-1)
            cells[get_cell_coords(pt)] = len(samples) - 1
        else:
            active.remove(idx)

    ################################################
    ### halftone
    ################################################
    edge_bit = 1.0-(edge/255.0 >= 0.5)*1.0
    w, h = image.shape[:2]
    edge_bit_large = np.zeros((w+4, h+4))
    edge_bit_large[2:2+w,2:2+h] = edge_bit

    area = []
    xs = []
    ys = []

    for i in range(len(samples)):
        x, y = samples[i]
        x_round = round(x)+2
        y_round = round(y)+2
        sum = np.sum(edge_bit_large[height-y_round-2:height-y_round+3,x_round-2:x_round+3])
        if sum > thresh:
            xs.append(x)
            ys.append(y)
            area.append(sum/25)

    area = [i*dotsize for i in area]

    plt.scatter(xs,ys, color='k', alpha=1, lw=0, s=area)
    plt.xlim(0, width)
    plt.ylim(0, height)
    plt.axis('off')
    plt.savefig(halftone_path)
    # if savefig:
    #     if halftone_path == '':
    #         plt.show()
    #         print("Please provide a halftone_path, not saving image")
    #     else:
    #         plt.savefig(halftone_path)
    #         plt.show()
    # else:
    #     plt.show()
    img = cv2.imread(halftone_path)

    ls_points = []
    for i in range(len(xs)):
        ls_points.append([round(xs[i],2), round(ys[i],2), round(area[i],2)])
    print(f"{len(ls_points)} points are generated.")

    sorted_points = sort_points(ls_points)

    return img, sorted_points

def sort_points(ls_points):
    '''Starting from the first point, sort the points by nearest neighbor'''
    sorted_points = []
    sorted_points.append(ls_points[0])
    ls_points.pop(0)
    while len(ls_points) > 0:
        last_point = sorted_points[-1]
        min_dist = 100000000
        min_point = []
        for point in ls_points:
            dist = np.sqrt((point[0]-last_point[0])**2 + (point[1]-last_point[1])**2)
            if dist < min_dist:
                min_dist = dist
                min_point = point
        sorted_points.append(min_point)
        ls_points.remove(min_point)
    return sorted_points

# if __name__ == '__main__':
img, sorted_points = halftone('03-photo/media/lightning.jpeg', radius=5, dotsize=30)
plt.imshow(img)
plt.show()
print(sorted_points)