import numpy as np
import cv2
import argparse
import math

def intersect(P0,P1):
    """P0 and P1 are NxD arrays defining N lines.
    D is the dimension of the space. This function
    returns the least squares intersection of the N
    lines from the system given by eq. 13 in
    http://cal.cs.illinois.edu/~johannes/research/LS_line_intersect.pdf.
    """
    # generate all line direction vectors
    n = (P1-P0)/np.linalg.norm(P1-P0,axis=1)[:,np.newaxis] # normalized

    # generate the array of all projectors
    projs = np.eye(n.shape[1]) - n[:,:,np.newaxis]*n[:,np.newaxis]  # I - n*n.T
    # see fig. 1

    # generate R matrix and q vector
    R = projs.sum(axis=0)
    q = (projs @ P0[:,:,np.newaxis]).sum(axis=0)

    # solve the least squares problem for the
    # intersection point p: Rp = q
    p = np.linalg.lstsq(R,q,rcond=None)[0]

    return p

width = 960
height = 540

K = np.eye(3)
K[0,0] = K[1,1] = 1000
K[0:2, 2] = (width-1)/2.0, (height-1)/2.0



ap = argparse.ArgumentParser()
ap.add_argument('-i', '--input', required=True, help = 'path to input image')
ap.add_argument('-o', '--output', required=False, help = 'path to output video')
ap.add_argument('-f', '--factor', required=False, help = 'downscale factor')
ap.add_argument('-s', '--show', required=False, action='store_true', help = 'Shows output if present')
args = ap.parse_args()

if args.show:
    win_name = "Stream"
    cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)

cap = cv2.VideoCapture(args.input)
writer = None

if args.factor is not None:
    factor = int(args.factor)
else:
    factor = 1

while cap.isOpened():
    # read the next frame from the file
    grabbed, img = cap.read()
    if not grabbed:
        break

    #img = cv2.undistort(img, K, np.float32([-0.2,  0., 0., 0., 0.]))
    img = cv2.resize(img, (int(img.shape[1]/factor), int(img.shape[0]/factor)))
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    lsd = cv2.createLineSegmentDetector(0)
    lines = lsd.detect(gray)[0] #Position 0 of the returned tuple are the detected lines
    #drawn_img = lsd.drawSegments(img,lines)
    #plt.imshow(drawn_img)

    count = 0
    a = []
    b = []

    for line in lines:
        #print(line)
        min_theta = 0.35
        x1 = line[0][0]
        y1 = line[0][1]
        x2 = line[0][2]
        y2 = line[0][3]
        theta = math.atan((y2-y1)/(x2-x1))
        dist = np.linalg.norm((x1-x2,y1-y2))
        if dist > 100 and abs(theta) > min_theta and abs(theta) < math.pi/2 - min_theta:
            a.append([x1, y1])
            b.append([x2, y2])
            count += 1
            vec_x = x2-x1
            vec_y = y2-y1
            fact = 10000
            #cv2.line(img,(x1,y1),(x2,y2),(0,255,0),2)
            cv2.line(img,(int(x1+fact*vec_x), int(y1+fact*vec_y)),(int(x2-fact*vec_x),int(y2-fact*vec_y)),(0,255,0),2)

    if(len(a) > 0):
        point = intersect(np.array(a),np.array(b))
        cv2.circle(img, (int(point[0][0]), int(point[1][0])), 10, (255,0,0), 2)


    if args.show:
        cv2.imshow(win_name, img)
        cv2.resizeWindow(win_name, 1080, 720)
        if cv2.waitKey(1) == ord('q'):
            break

    if args.output is not None:
        if writer is None:
            # initialize our video writer
            fourcc = cv2.VideoWriter_fourcc(*"MJPG")
            writer = cv2.VideoWriter(args.output, fourcc, 30,
                    (img.shape[1], img.shape[0]), True)
        # write the output frame to disk
        writer.write(img)


# release the file pointers
print("[INFO] cleaning up...")
writer.release()
cap.release()
if args.show:
    cv2.destroyAllWindows()
