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

class Dist:
    
    '''input: vec1 , vec2 均为列向量'''
    
    @staticmethod
    def euclidean_dist(vec1, vec2):
        """欧氏距离:
        我们现实所说的直线距离"""
        if type(vec1)==list:
            vec1, vec2 = np.array(vec1), np.array(vec2)
        assert vec1.shape == vec2.shape 

        return np.sqrt((vec2-vec1).T @ (vec2-vec1))
    
    @staticmethod
    def manhattan_dist(vec1, vec2):
        """曼哈顿距离:
        城市距离"""
        if type(vec1)==list:
            vec1, vec2 = np.array(vec1), np.array(vec2)
        assert vec1.shape == vec2.shape 
        
        return sum(abs(vec1 - vec2))

def get_gradient(coord, img, delta = 2):
    '''coord: [x,y] numpy array, img: gray image, dtype: float, if not :img/255
    get image gradient at specified coordinates 
    有没有可能这里的delta应该等于R
    '''
    H, W= img.shape
    if coord[0] < W-delta and coord[1] < H-delta:
        dh = img[coord[1], coord[0]+delta] - img[coord[1], coord[0]]
        dw = img[coord[1]+delta, coord[0]] - img[coord[1], coord[0]]
        mag = np.sqrt(dw*dw + dh*dh)
        return dh/delta,dw/delta, mag
    else:
        return 0,0,0

def get_randomPoint(refImage):
    h,w = refImage.shape[:2]
    x,y = np.random.randint(0,h), np.random.randint(0,w)
    return x,y

def get_point_color(coord, image):
    '''image: bgra channel
        coord: w,h array
    return  bgr color array'''
    bgr = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
    h,w = image.shape[:2]
    coord[0] = np.clip(coord[0], 0, h-1)
    coord[1] = np.clip(coord[1], 0, w-1)
    b,g,r = bgr[coord[0], coord[1], 0], bgr[coord[0], coord[1], 1], bgr[coord[0], coord[1], 2]
    return np.array([b, g, r])


def makeSplineStroke(x0, y0, R, refImage, canvas):
    '''refImage: bgra channel
    取落笔时的颜色，绘制一个stroke
    '''
    strokeColor = get_point_color([y0, x0], refImage)
    K = []
    K.append([x0, y0])
    [x,y] = [x0, y0]
    [lastDx, lastDy] = [1e-6,1e-6]
    
    minStrokeLength, maxStrokeLength = 10, 30
    gray_img = cv2.cvtColor(refImage,cv2.COLOR_BGR2GRAY)
    for i in range(maxStrokeLength):

        refColor = get_point_color([y,x], refImage)
#         diff_1 = Dist.euclidean_dist(refColor, get_point_color([y,x], canvas))
#         diff_2 = Dist.euclidean_dist(refColor, strokeColor)
        diff_1 = Dist.manhattan_dist(refColor, get_point_color([y,x], canvas))
        diff_2 = Dist.manhattan_dist(refColor, strokeColor)

        if i > minStrokeLength and (diff_1 < diff_2):
            return K
        
        gx,gy, gradientMag = get_gradient([x,y], gray_img/255)
        # detect vanishing gradient
#         print('gradientMag', gradientMag)
        if gradientMag==0:
            return K
        # compute a normal direction
        dx, dy = -gy, gx
        # if necessary, reverse direction
        if lastDx * dx +lastDy * dy < 0:
            dx, dy = -dx, dy
        # filter the stroke direction
        fc = 1.
        dx = fc*dx + (1-fc)*lastDx 
        dy = fc*dy + (1-fc)*lastDy
       
        dx, dy = dx/np.sqrt(dx*dx + dy*dy), dy/np.sqrt(dx*dx + dy*dy)
        x,y = int(x+R*dx), int(y+R*dy)
        
        lastDx, lastDy = dx, dy
        K.append([x, y])
        
    return K


from IPython import display

def test():
    refImage = cv2.imread('temp/nuannuan.png')

    img = cv2.blur(refImage, (10,10))
    canvas = np.zeros(refImage.shape).astype('uint8')
    R = refImage.shape[0]/70 # -----------------------------
    for i in range(1500):
        x0, y0 = get_randomPoint(img)
        K = makeSplineStroke(x0, y0,R,img, canvas)

        color = get_point_color([y0,x0], img)
        b,g,r = color[0],color[1],color[2]

        b,g,r = int(b), int(g), int(r)

        for coord in K:
            circle = cv2.circle(canvas,(coord[0], coord[1]), 12, (r,g,b),-1)
        
        if i%50==0:
            display.clear_output(wait=True)
            plt.figure(figsize=(6,6))
            plt.imshow(circle);plt.show()
            
    return circle


if __name__ == '__main__':

    circle = test()
    