import numpy as np
import cv2 as cv
import glob
import os

# Here is a litle help:
# https://temugeb.github.io/opencv/python/2021/02/02/stereo-camera-calibration-and-triangulation.html

def display_depth_map(mtx1, dist1, mtx2, dist2, R, T):
    # Load stereo rectification parameters
    R1, R2, P1, P2, Q, _, _ = cv.stereoRectify(mtx1, dist1, mtx2, dist2, (640, 480), R, T)
    map1x, map1y = cv.initUndistortRectifyMap(mtx1, dist1, R1, P1, (640, 480), cv.CV_32FC1)
    map2x, map2y = cv.initUndistortRectifyMap(mtx2, dist2, R2, P2, (640, 480), cv.CV_32FC1)

    # Initialize the stereo block matching algorithm
    stereo = cv.StereoBM_create(numDisparities=16, blockSize=5)

    # Initialize the stereo video capture
    cap1 = cv.VideoCapture(1)
    cap2 = cv.VideoCapture(2)

    while True:
        # Capture stereo images
        ret1, frame1 = cap1.read()
        ret2, frame2 = cap2.read()

        if not (ret1 and ret2):
            break

        # Rectify stereo images
        rectified_frame1 = cv.remap(frame1, map1x, map1y, cv.INTER_LINEAR)
        rectified_frame2 = cv.remap(frame2, map2x, map2y, cv.INTER_LINEAR)

        # Convert stereo images to grayscale
        gray1 = cv.cvtColor(rectified_frame1, cv.COLOR_BGR2GRAY)
        gray2 = cv.cvtColor(rectified_frame2, cv.COLOR_BGR2GRAY)

        # Compute the disparity map
        disparity = stereo.compute(gray1, gray2)

        # Convert disparity map to depth map
        depth_map = cv.convertScaleAbs(disparity, alpha=1.0/8)

        # Display depth map
        cv.imshow('Depth Map', depth_map)
        cv.waitKey(1000)
        # Break the loop when 'q' is pressed
        if cv.waitKey(1) & 0xFF == ord('q'):
            break

    # Release video capture and close windows
    cap1.release()
    cap2.release()
    cv.destroyAllWindows()
def stereo_line(mtx1, dist1, mtx2, dist2):
    cap1 = cv.VideoCapture(1)
    cap2 = cv.VideoCapture(2)
    # Create SIFT detector
    sift = cv.SIFT_create()

    while True:
        # Capture stereo images
        ret1, frame1 = cap1.read()
        ret2, frame2 = cap2.read()
        frame1 = cv.undistort(frame1, mtx1, dist1)
        frame2 = cv.undistort(frame2, mtx2, dist2)

        # Convert stereo images to grayscale
        frame1 = cv.cvtColor(frame1, cv.COLOR_BGR2GRAY)
        frame2 = cv.cvtColor(frame2, cv.COLOR_BGR2GRAY)

        # Find keypoints and descriptors
        kp1, des1 = sift.detectAndCompute(frame1, None)
        kp2, des2 = sift.detectAndCompute(frame2, None)

        # FLANN parameters
        FLANN_INDEX_KDTREE = 1
        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        search_params = dict(checks=50)

        # Create FLANN matcher
        flann = cv.FlannBasedMatcher(index_params, search_params)

        # Match descriptors
        matches = flann.knnMatch(des1, des2, k=2)

        # Apply ratio test as per Lowe's paper
        pts1 = []
        pts2 = []
        for m, n in matches:
            if m.distance < 0.8 * n.distance:
                pts2.append(kp2[m.trainIdx].pt)
                pts1.append(kp1[m.queryIdx].pt)

        # Convert points to numpy arrays
        pts1 = np.int32(pts1)
        pts2 = np.int32(pts2)

        # Find Fundamental Matrix using RANSAC
        F, mask = cv.findFundamentalMat(pts1, pts2, cv.FM_RANSAC)

        # Select only inlier points
        pts1 = pts1[mask.ravel() == 1]
        pts2 = pts2[mask.ravel() == 1]

        # Find epilines corresponding to points in the right image (second image)
        lines1 = cv.computeCorrespondEpilines(pts2.reshape(-1, 1, 2), 2, F)
        lines1 = lines1.reshape(-1, 3)

        # Find epilines corresponding to points in the left image (first image)
        lines2 = cv.computeCorrespondEpilines(pts1.reshape(-1, 1, 2), 1, F)
        lines2 = lines2.reshape(-1, 3)

        # Draw epilines on both images
        img5, img6 = drawlines(frame1, frame2, lines1, pts1, pts2)
        img3, img4 = drawlines(frame2, frame1, lines2, pts2, pts1)

        # Display images with epilines
        cv.imshow('Epilines Left', img5)
        cv.imshow('Epilines Right', img3)
        cv.waitKey(1000)

        # Break the loop on 'q' key press
        if cv.waitKey(1) & 0xFF == ord('q'):
            break

    # Release video capture and close windows
    cap1.release()
    cap2.release()
    cv.destroyAllWindows()
def drawlines(img1, img2, lines, pts1, pts2):
    ''' Draw epilines on the images '''
    r, c = img1.shape[:2]
    img1_color = cv.cvtColor(img1, cv.COLOR_GRAY2BGR)
    img2_color = cv.cvtColor(img2, cv.COLOR_GRAY2BGR)
    for r, pt1, pt2 in zip(lines, pts1, pts2):
        color = tuple(np.random.randint(0, 255, 3).tolist())
        x0, y0 = map(int, [0, -r[2] / r[1]])
        x1, y1 = map(int, [c, -(r[2] + r[0] * c) / r[1]])
        img1_color = cv.line(img1_color, (x0, y0), (x1, y1), color, 1)
        img1_color = cv.circle(img1_color, tuple(pt1), 5, color, -1)
        img2_color = cv.circle(img2_color, tuple(pt2), 5, color, -1)
    return img1_color, img2_color
def finding_gripper(mtx1, dist1):
    def draw(img, corners, imgpts):
        corner = tuple(corners[0].ravel().astype(int))
        imgpt_0 = tuple(imgpts[0].ravel().astype(int))
        imgpt_1 = tuple(imgpts[1].ravel().astype(int))
        imgpt_2 = tuple(imgpts[2].ravel().astype(int))
        img = cv.line(img, corner, imgpt_0, (255,0,0), 5)
        img = cv.line(img, corner, imgpt_1, (0,255,0), 5)
        img = cv.line(img, corner, imgpt_2, (0,0,255), 5)
        return img
    
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    row = 3
    col = 3
    objp = np.zeros((row*col,3), np.float32)
    objp[:,:2] = np.mgrid[0:col,0:row].T.reshape(-1,2)
    
    axis = np.float32([[3,0,0], [0,3,0], [0,0,-3]]).reshape(-1,3)

    cap = cv.VideoCapture(1)
    while True:
        ret, img = cap.read()  # Read a frame from the camera
        if not ret:
            print("Failed to capture frame")
            break

        gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
        ret, corners = cv.findChessboardCorners(gray, (col,row),None)
        
        if ret == True:
            corners2 = cv.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
    
            # Find the rotation and translation vectors.
            ret,rvecs, tvecs = cv.solvePnP(objp, corners2, mtx1, dist1)
    
            # project 3D points to image plane
            imgpts, jac = cv.projectPoints(axis, rvecs, tvecs, mtx1, dist1)
            img = draw(img,corners2,imgpts)
            cv.imshow('img',img)
            if cv.waitKey(1) & 0xFF == ord('q'):
                break
    
    cv.destroyAllWindows()
def threeD_cube(mtx1, dist1, mtx2, dist2, R, T):
    # Define cube vertices in 3D space
    cube_vertices = np.float32([[0, 0, 0], [0, 3, 0], [3, 3, 0], [3, 0, 0], [0, 0, -3], [0, 3, -3], [3, 3, -3], [3, 0, -3]])
    sq_s = 30
    square_vertices = np.float32([[-sq_s/2, sq_s/2, 0], [sq_s/2, sq_s/2, 0], [sq_s/2, -sq_s/2, 0], [-sq_s/2, -sq_s/2, 0]])
    row = 8
    col = 5
    objp = np.zeros((row*col,3), np.float32)
    objp[:,:2] = np.mgrid[0:col,0:row].T.reshape(-1,2)
    # Initialize video capture for both cameras
    cap1 = cv.VideoCapture(1)  # Assuming camera1 is connected at index 0
    cap2 = cv.VideoCapture(2)  # Assuming camera2 is connected at index 1

    while True:
        # Capture frame from camera1
        ret1, frame1 = cap1.read()
        if not ret1:
            break

        # Undistort frame from camera1
        frame1_undistorted = cv.undistort(frame1, mtx1, dist1)
        
        # Detect corners of the chessboard
        ret_corners, corners = cv.findChessboardCorners(frame1_undistorted, (5, 8), None)
        cv.drawChessboardCorners(frame1_undistorted, (5,8), corners, ret_corners)
        if ret_corners:
            # Estimate pose of the chessboard
            ret_pose, rvecs, tvecs = cv.solvePnP(objp, corners, mtx1, dist1)

            # Project cube vertices onto image plane of camera1
            cube_vertices_img1, _ = cv.projectPoints(cube_vertices, rvecs, tvecs, mtx1, dist1)

            # Draw cube on frame from camera1
            for i in range(4):
                frame1_undistorted = cv.line(frame1_undistorted, tuple(cube_vertices_img1[i].ravel().astype(int)), tuple(cube_vertices_img1[(i+1) % 4].ravel().astype(int)), (0, 255, 0), 3)
                frame1_undistorted = cv.line(frame1_undistorted, tuple(cube_vertices_img1[i].ravel().astype(int)), tuple(cube_vertices_img1[i+4].ravel().astype(int)), (0, 255, 0), 3)
                frame1_undistorted = cv.line(frame1_undistorted, tuple(cube_vertices_img1[i+4].ravel().astype(int)), tuple(cube_vertices_img1[(i+1) % 4 + 4].ravel().astype(int)), (0, 255, 0), 3)

            # Display frame from camera1 with cube
            cv.imshow('Camera 1', frame1_undistorted)
            """
            # Capture frame from camera2
            ret2, frame2 = cap2.read()
            if not ret2:
                break

            # Apply transformation to cube vertices for camera2
            transformed_vertices = np.dot(cube_vertices, R.T) + T

            # Project transformed cube vertices onto image plane of camera2
            cube_vertices_img2, _ = cv.projectPoints(transformed_vertices, rvecs, tvecs, mtx1, dist1)

            # Draw transformed cube on frame from camera2
            for i in range(4):
                frame2 = cv.line(frame2, tuple(cube_vertices_img2[i].ravel().astype(int)), tuple(cube_vertices_img2[(i+1) % 4].ravel().astype(int)), (0, 255, 0), 3)
                frame2 = cv.line(frame2, tuple(cube_vertices_img2[i].ravel().astype(int)), tuple(cube_vertices_img2[i+4].ravel().astype(int)), (0, 255, 0), 3)
                frame2 = cv.line(frame2, tuple(cube_vertices_img2[i+4].ravel().astype(int)), tuple(cube_vertices_img2[(i+1) % 4 + 4].ravel().astype(int)), (0, 255, 0), 3)

            # Display frame from camera2 with transformed cube
            cv.imshow('Camera 2', frame2)"""

            # Exit on 'q' key press
            if cv.waitKey(1) & 0xFF == ord('q'):
                break

    # Release video capture and close windows
    cap1.release()
    cap2.release()
    cv.destroyAllWindows()
def three_axis(mtx1, dist1, mtx2, dist2, R, T):
    def draw(img, corners, imgpts):
        corner = tuple(corners[0].ravel().astype(int))
        imgpt_0 = tuple(imgpts[0].ravel().astype(int))
        imgpt_1 = tuple(imgpts[1].ravel().astype(int))
        imgpt_2 = tuple(imgpts[2].ravel().astype(int))
        img = cv.line(img, corner, imgpt_0, (255,0,0), 5)
        img = cv.line(img, corner, imgpt_1, (0,255,0), 5)
        img = cv.line(img, corner, imgpt_2, (0,0,255), 5)
        return img
    
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
    row = 8
    col = 5
    objp = np.zeros((row*col,3), np.float32)
    objp[:,:2] = np.mgrid[0:col,0:row].T.reshape(-1,2)
    
    axis = np.float32([[3,0,0], [0,3,0], [0,0,-3]]).reshape(-1,3)

    cap = cv.VideoCapture(1)
    while True:
        ret, img = cap.read()  # Read a frame from the camera
        if not ret:
            print("Failed to capture frame")
            break

        gray = cv.cvtColor(img,cv.COLOR_BGR2GRAY)
        ret, corners = cv.findChessboardCorners(gray, (col,row),None)
        
        if ret == True:
            corners2 = cv.cornerSubPix(gray,corners,(11,11),(-1,-1),criteria)
    
            # Find the rotation and translation vectors.
            ret,rvecs, tvecs = cv.solvePnP(objp, corners2, mtx1, dist1)
    
            # project 3D points to image plane
            imgpts, jac = cv.projectPoints(axis, rvecs, tvecs, mtx1, dist1)
            img = draw(img,corners2,imgpts)
            cv.imshow('img',img)
            if cv.waitKey(1) & 0xFF == ord('q'):
                break
    
    cv.destroyAllWindows()

#display_depth_map(mtx1, dist1, mtx2, dist2, R, T)
#stereo_line(mtx1, dist1, mtx2, dist2, R, T)
#finding_gripper(mtx1, dist1)
#threeD_cube(mtx1, dist1, mtx2, dist2, R, T)
#three_axis(mtx1, dist1, mtx2, dist2, R, T)