import numpy as np
import cv2

def make_triangle(img, num_triangles):
    print(img.shape)
    y,x = (img.shape[0]//2, img.shape[1]//2)
    angles = 2 * np.pi/num_triangles
    print(angles/2)
    w,h,d = img.shape
    print(np.tan(angles/2))
    z = int(np.tan(angles/2) * (h/2))
    print(z)
    print(h)
    u = (x + z, y + h/2)
    v = (x - z, y + h/2)
    mask = np.zeros((w,h,d))
    
    pts = np.array([v,(x,y),u], np.int32)
    pts = pts.reshape((-1,1,2))
    mask = cv2.fillPoly(mask, [pts], (255,0,0))

    # With mask, get the triangle from the original image.
    img[:,:,0] = np.where(mask[:,:,0] == 255, img[:,:,0], 0)
    img[:,:,1] = np.where(mask[:,:,0] == 255, img[:,:,1], 0)
    img[:,:,2] = np.where(mask[:,:,0] == 255, img[:,:,2], 0)
    return img

def rotate(im, rotation):
    M = cv2.getRotationMatrix2D((im.shape[1]/2,im.shape[0]/2), rotation, 1)
    im[:,:,0] = cv2.warpAffine(im[:,:,0],M,(im.shape[1],im.shape[0]))
    im[:,:,1] = cv2.warpAffine(im[:,:,1],M,(im.shape[1],im.shape[0]))
    im[:,:,2] = cv2.warpAffine(im[:,:,2],M,(im.shape[1],im.shape[0]))
    return im

def _stitch(img, to_stitch):
    img[:,:,0] = np.where((img[:,:,0] == 0) & (to_stitch[:,:,0] != 0), to_stitch[:,:,0], img[:,:,0])
    img[:,:,1] = np.where((img[:,:,1] == 0) & (to_stitch[:,:,1] != 0), to_stitch[:,:,1], img[:,:,1])
    img[:,:,2] = np.where((img[:,:,2] == 0) & (to_stitch[:,:,2] != 0), to_stitch[:,:,2], img[:,:,2])

def make_kaleidoscope(img, num):
    triangle = make_triangle(img, num)
    iters = num
    while iters > 0:
        new_triangle = np.copy(triangle)
        new_triangle = cv2.flip(new_triangle, 1) if iters % 2 != 0 else new_triangle
        rotate(new_triangle, 360/num * iters)
        _stitch(triangle, new_triangle)
        iters -= 1
    return triangle

if __name__ == "__main__":
    img = cv2.imread("/Users/piv/Documents/Projects/car/GestureRecognition/IMG_0818.png")
    img = cv2.resize(img, None, fx=0.3, fy=0.3, interpolation = cv2.INTER_AREA)
    num = 12
    kaleid = make_kaleidoscope(img, num)
    cv2.imshow("", kaleid)
    cv2.waitKey(0)
    cv2.destroyAllWindows()