import numpy as np
import cv2
import glob



def RunAll():
    #Part A, finding how to undistort image

    # termination criteria
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    WIDTH, HEIGHT = 6, 9

    objp = np.zeros((WIDTH*HEIGHT,3), np.float32)
    objp[:,:2] = np.mgrid[0:HEIGHT,0:WIDTH].T.reshape(-1,2)
    #scaling up points
    objp *= 100

    # Arrays to store object points and image points from all the images.
    objpoints = [] # 3d point in real world space
    imgpoints = [] # 2d points in image plane.
    images = glob.glob("images/sample_*.jpg")

    for filename in images:
        img = cv2.imread(filename)
        gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

        # Find the chess board corners
        ret, corners = cv2.findChessboardCorners(gray, (HEIGHT,WIDTH),None)
        # If found, add object points, image points (after refining them)
        if ret:
            objpoints.append(objp)
            cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
            imgpoints.append(corners)

            # Draw and display the corners
            cv2.drawChessboardCorners(img, (HEIGHT,WIDTH), corners,ret)
            cv2.imshow('img', img)
            cv2.waitKey(10)
    cv2.destroyAllWindows()
    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1])

    image = cv2.imread('images/test_001.jpg')
    h, w = image.shape[:2]
    newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1,(w,h))
    # undistort
    dst = cv2.undistort(image, mtx, dist, None, newcameramtx)
    # crop the image
    x,y,w,h = roi
    dst = dst[y:y+h, x:x+w]
    cv2.imwrite('CalibrateResult.png',dst)


    #PartB and C
    cap = cv2.VideoCapture(0)
    image_to_place1 = cv2.imread("Show_me_the_image.jpg")
    #resize to fit on board
    image_to_place = cv2.resize(image_to_place1, ((HEIGHT-1)*100,(WIDTH-1)*100))
    #Whitespace to act as canvas for out image
    white_space1 = cv2.imread("white.jpg")
    #resize to fit on board
    white_space = cv2.resize(white_space1, ((HEIGHT-1)*100,(WIDTH-1)*100))

    while cv2.waitKey(1) & 0xFF != ord('q'):
        #capture a frame
        ret, img = cap.read()
        h,  w = img.shape[:2]
        newcameramtx, roi=cv2.getOptimalNewCameraMatrix(mtx,dist,(w,h),1,(w,h))
        undst = cv2.undistort(img, mtx, dist, None, newcameramtx)
        # Our operations on the frame come here
        gray = cv2.cvtColor(undst, cv2.COLOR_BGR2GRAY)

        # Find the chess board corners
        ret, corners = cv2.findChessboardCorners(gray, (HEIGHT,WIDTH),None)
        # If found, add object points and image points
        if ret:
            cv2.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
            #Draw and display the corners
            cv2.drawChessboardCorners(undst, (HEIGHT,WIDTH), corners,ret)
            projection = CalibrateCamera3D(np.hstack([objp, np.reshape(corners, (54, 2))]))
            image_to_place_warp = cv2.warpPerspective(image_to_place, projection, (w,h))
            white_warp = cv2.bitwise_not(cv2.warpPerspective(white_space, projection, (w,h)));
            undst = cv2.bitwise_and(undst, white_warp) + image_to_place_warp

        cv2.imshow('img',undst)

    # release everything
    cap.release()
    cv2.destroyAllWindows()

#Previous Assignments Camera Calibration
def insertOne(data):
    return np.insert(data, 2, 1, axis=1)

def CalibrateCamera3D(data):
    """ copy of data to work with"""
    data=np.delete(data, 2, 1)
    datacopy= insertOne(data).copy()
    """getting XYZ1, x and y respectively"""
    P, x, y = datacopy[:,:3], datacopy[:,-2], datacopy[:,-1]
    """blank array to fill with the new matrix"""
    totalarr = []
    """looping over each row in P so can get each line of XY10... and 000X... """
    for i, row in enumerate(P):
        arr = np.hstack([row,np.zeros(3),-x[i]*row])
        arr1 = np.hstack([np.zeros(3), row  , -y[i]*row])
        totalarr.append(np.vstack([arr, arr1]))
    totalarr = np.vstack(totalarr)
    """array Transpose by array"""
    AtA = totalarr.T.dot(totalarr)
    """getting the eignvectors"""
    eigvals, eigvectors = np.linalg.eig(AtA)
    """ eigenvectors are columns so Transpose matrix"""
    #make sure the eigenvector is always the smallest, as eigvectors[-1] sometimes is not.
    M = eigvectors.T[np.argmin(eigvals)]
    """reshape to correct shape"""
    return M.reshape((3, 3))

if __name__ == "__main__":
    RunAll()