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

def find_camera(find_flag):
    if find_flag:
        cam_available = []
        for i in range(10):  # Try indices from 0 to 9
            cap = cv.VideoCapture(i)
            if cap.isOpened():
                print(f"Camera found at index {i}")
                cam_available.append(i)
                cap.release()
            if len(cam_available) > 2:
                break

        if len(cam_available) > 2 and cam_available[0] == 0:
            cam1 = cam_available[1]
            cam2 = cam_available[2]
        else:
            cam1 = cam_available[0]
            cam2 = cam_available[1]
    else:
        cam1 = 1
        cam2 = 2
    print(f"Cameras number used : {cam1} & {cam2}")
    return cam1, cam2
def img_capture(camera_num):
    # Create a directory to save captured images
    output_dir = f"camera{camera_num}_images"
    os.makedirs(output_dir, exist_ok=True)

    # Initialize the camera
    cap = cv.VideoCapture(camera_num)

    # Check if the camera is opened successfully
    if not cap.isOpened():
        print(f"Error: Could not open camera {camera_num}")
        exit()
    i = 0
    # Capture and save 12 images
    while i < 6:
        # Capture a frame from the camera
        ret, frame = cap.read()

        # Check if the frame is captured successfully
        if not ret:
            print("Error: Could not read frame")
            break

        # Display the captured image
        cv.imshow('Capture Image', frame)

        # Save the captured image if the 's' key is pressed
        key = cv.waitKey(5) & 0xFF
        if key == ord('s'):
            img_path = os.path.join(output_dir, f'image_{i+1}.jpg')
            cv.imwrite(img_path, frame)
            print(f"Image {i+1} saved: {img_path}")
            i += 1
        # If 'q' key is pressed, exit the loop
        elif key == ord('q'): break
    # Release the camera and close all OpenCV windows
    cap.release()
    cv.destroyAllWindows()
    print("Image capture complete.")
    return
def single_calibration(camera_num, img_cap):
    if img_cap: img_capture(camera_num)
    # Termination criteria
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    # Prepare object points, assuming a chessboard with 9 by 6 squares of 30mm
    square_size = 30  # in millimeters
    row = 9
    col = 6
    objp = np.zeros((row * col, 3), np.float32)
    objp[:, :2] = np.mgrid[0:row, 0:col].T.reshape(-1, 2) * square_size

    # 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(f'camera{camera_num}_images/*.jpg')
    
    for frame in images:
        img = cv.imread(frame)
        gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        # Find the chessboard corners
        ret, corners = cv.findChessboardCorners(gray, (row, col), None)

        # If found, add object points, image points (after refining them)
        if ret == True:
            objpoints.append(objp)
            corners2 = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
            imgpoints.append(corners2)
            # Draw and display the corners
            cv.drawChessboardCorners(img, (row, col), corners2, ret)
            cv.imshow('img', img)
            cv.waitKey(1)

    cv.destroyAllWindows()
    ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(objpoints, imgpoints, (gray.shape[1], gray.shape[0]), None, None)
    return mtx, dist

def stereo_capture():
    # Open two video capture objects for each camera
    cap_left = cv.VideoCapture(cam1)  # Adjust the index if needed
    cap_right = cv.VideoCapture(cam2)  # Adjust the index if needed

    # Check if the cameras opened successfully
    if not cap_left.isOpened() or not cap_right.isOpened():
        print("Error: Couldn't open one or both cameras.")
        exit()
    
    # Create a directory to save images
    output_dir = 'stereo_images'
    os.makedirs(output_dir, exist_ok=True)

    frame_counter = 0
    while frame_counter < 6:
        # Read frames from both cameras
        ret_left, frame_left = cap_left.read()
        ret_right, frame_right = cap_right.read()

        # Break the loop if either of the cameras fails to read a frame
        if not ret_left or not ret_right:
            print("Error: Couldn't read frames from one or both cameras.")
            break

        # Display the frames side by side for stereo effect
        stereo_frame = cv.hconcat([frame_left, frame_right])
        cv.imshow('Stereo Camera Feed', stereo_frame)

        key = cv.waitKey(5) & 0xFF
        # Save the captured image if the 's' key is pressed
        if key == ord('s'):
            # Save the frames from both cameras
            frame_counter += 1
            img_path_left = os.path.join(output_dir, f'{frame_counter}_left_image.jpg')
            img_path_right = os.path.join(output_dir, f'{frame_counter}_right_image.jpg')
            cv.imwrite(img_path_left, frame_left)
            cv.imwrite(img_path_right, frame_right)
            print(f"Image {frame_counter} saved")
        # Break the loop if 'q' key is pressed
        if key == ord('q'): break
    # Release the video capture objects and close the OpenCV window
    cap_left.release()
    cap_right.release()
    cv.destroyAllWindows()
    return
def stereo_calibration(mtx1, dist1, mtx2, dist2, frames_folder, stereo_capture_flag):
    if stereo_capture_flag: stereo_capture()
    # Read the synched frames
    images_names = glob.glob(frames_folder)
    images_names = sorted(images_names)
    c1_images_names = images_names[0::2]
    c2_images_names = images_names[1::2]
 
    c1_images = []
    c2_images = []
    for im1, im2 in zip(c1_images_names, c2_images_names):
        _im = cv.imread(im1, 1)
        c1_images.append(_im)
 
        _im = cv.imread(im2, 1)
        c2_images.append(_im)
 
    #change this if stereo calibration not good.
    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.0001)
 
    rows = 6 #number of checkerboard rows.
    columns = 9 #number of checkerboard columns.
    world_scaling = 30 #change this to the real world square size. Or not.
 
    #coordinates of squares in the checkerboard world space
    objp = np.zeros((rows*columns,3), np.float32)
    objp[:,:2] = np.mgrid[0:rows,0:columns].T.reshape(-1,2)
    objp = world_scaling* objp
 
    #frame dimensions. Frames should be the same size.
    width = c1_images[0].shape[1]
    height = c1_images[0].shape[0]
 
    #Pixel coordinates of checkerboards
    imgpoints_left = [] # 2d points in image plane.
    imgpoints_right = []
 
    #coordinates of the checkerboard in checkerboard world space.
    objpoints = [] # 3d point in real world space
 
    for frame1, frame2 in zip(c1_images, c2_images):
        gray1 = cv.cvtColor(frame1, cv.COLOR_BGR2GRAY)
        gray2 = cv.cvtColor(frame2, cv.COLOR_BGR2GRAY)
        c_ret1, corners1 = cv.findChessboardCorners(gray1, (rows, columns), None)
        c_ret2, corners2 = cv.findChessboardCorners(gray2, (rows, columns), None)
 
        if c_ret1 == True and c_ret2 == True:
            corners1 = cv.cornerSubPix(gray1, corners1, (11, 11), (-1, -1), criteria)
            corners2 = cv.cornerSubPix(gray2, corners2, (11, 11), (-1, -1), criteria)
 
            cv.drawChessboardCorners(frame1, (rows, columns), corners1, c_ret1)
            #cv.imshow('img', frame1)
 
            cv.drawChessboardCorners(frame2, (rows, columns), corners2, c_ret2)
            #cv.imshow('img2', frame2)
            stereo_chess = cv.hconcat([frame1, frame2])
            cv.imshow('stereo', stereo_chess)
            cv.waitKey(1)
 
            objpoints.append(objp)
            imgpoints_left.append(corners1)
            imgpoints_right.append(corners2)
 
    stereocalibration_flags = cv.CALIB_FIX_INTRINSIC
    ret, CM1, dist1_bis, CM2, dist2_bis, R, T, E, F = cv.stereoCalibrate(objpoints, imgpoints_left, imgpoints_right, mtx1, dist1, mtx2, dist2, (width, height), criteria = criteria, flags = stereocalibration_flags)
    cv.destroyAllWindows()
    return R, T

def cub_cordinate(mtx1, dist1, mtx2, dist2, R, T):
    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 not ret2 : break

        frame1 = cv.undistort(frame1, mtx1, dist1)
        frame2 = cv.undistort(frame2, mtx2, dist2)
        
        # Detect red cube in both images
        point1 = detect_cube(frame1, False)
        point2 = detect_cube(frame2, False)
        """
        point1 = np.array(point1)
        point2 = np.array(point2)
        
        #RT matrix for C1 is identity.
        RT1 = np.concatenate([np.eye(3), [[0],[0],[0]]], axis = -1)
        P1 = mtx1 @ RT1 #projection matrix for C1
        
        #RT matrix for C2 is the R and T obtained from stereo calibration.
        RT2 = np.concatenate([R, T], axis = -1)
        P2 = mtx2 @ RT2 #projection matrix for C2

        # Call the triangulatePoints function
        points3d_homogeneous = cv.triangulatePoints(P1, P2, point1, point2)
        # Convert homogeneous coordinates to Euclidean coordinates
        points3d_homogeneous /= points3d_homogeneous[3]
        # Extract the 3D points from the homogeneous coordinates
        points3d = points3d_homogeneous[:3]

        print(points3d_homogeneous)"""
        
        #cal_point2 = project_point_to_camera2(point1, mtx1, R, T, mtx2)
        transform = np.vstack((np.hstack((R, T)), [0, 0, 0, 1]))
        point_homogeneous = np.array([point1[0], point1[1], 1, 1])
        cal_point1_homogeneous = np.dot(transform, point_homogeneous)
        cal_point1 = cal_point1_homogeneous[:2] / cal_point1_homogeneous[3]
        cal_point1_x, cal_point1_y = cal_point1

        cv.circle(frame1, (int(point1[0]), int(point1[1])), 2, (0, 0, 255), -1)
        cv.circle(frame2, (int(point2[0]), int(point2[1])), 2, (0, 0, 255), -1)
        cv.circle(frame2, (int(cal_point1_x), int(cal_point1_y)), 2, (255, 0, 0), -1)
        print(point2, cal_point1)

        stereo_frame = cv.hconcat([frame1, frame2])
        cv.imshow('Stereo Frames', stereo_frame)
        cv.waitKey(1)

        # 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 detect_cube(image, show_flag):
    # Convert image to HSV color space
    hsv = cv.cvtColor(image, cv.COLOR_BGR2HSV)

    # Define lower and upper bounds for red color in HSV
    # Red range
    #lower = np.array([0, 100, 100])
    #upper = np.array([5, 255, 255])
    # Yellow range
    #lower = np.array([25, 100, 100])
    #upper = np.array([35, 255, 255])
    # Green range
    #lower = np.array([40, 80, 80])
    #upper = np.array([60, 255, 255])
    # Blue range
    lower = np.array([100, 100, 100])
    upper = np.array([110, 255, 255])

    # Threshold the HSV image to get only red colors
    mask = cv.inRange(hsv, lower, upper)
    # Find non-zero pixel coordinates
    non_zero_pixels = cv.findNonZero(mask)

    # Check if non-zero pixels are found
    if non_zero_pixels is not None:
        # Calculate the average position and extract x and y coordinates of the average position
        average_position = np.mean(non_zero_pixels, axis=0)
        avg_x, avg_y = average_position[0]
    else: avg_x, avg_y = 0, 0

    if show_flag :
        # Apply the mask to the original image
        masked_image = cv.bitwise_and(image, image, mask=mask)
        cv.circle(masked_image, (int(avg_x), int(avg_y)), 2, (0, 0, 255), -1)
        cv.imshow('Remaining Image', masked_image)
        cv.waitKey(1)

    if 0: # Calculate the average value for each channel (Hue, Saturation, Value) across non-zero pixels
        non_zero_indices = np.nonzero(mask)
        non_zero_pixel_values = hsv[non_zero_indices]
        avg = np.mean(non_zero_pixel_values, axis=0)
        print(avg)
    return (avg_x, avg_y)

def triangulate(mtx1, mtx2, R, T):
 
    uvs1 = [[458, 86]]
    uvs2 = [[540, 311]]
 
    uvs1 = np.array(uvs1)
    uvs2 = np.array(uvs2)
 
    #RT matrix for C1 is identity.
    RT1 = np.concatenate([np.eye(3), [[0],[0],[0]]], axis = -1)
    P1 = mtx1 @ RT1 #projection matrix for C1
 
    #RT matrix for C2 is the R and T obtained from stereo calibration.
    RT2 = np.concatenate([R, T], axis = -1)
    P2 = mtx2 @ RT2 #projection matrix for C2

def project_point_to_camera2(point_cam1, mtx1, R, T, mtx2):
    # Step 1: Convert point coordinates to world coordinates in camera 1
    point_world = np.dot(np.linalg.inv(mtx1), np.append(point_cam1, 1))
    # Step 2: Transform world coordinates to camera 2 coordinate system
    point_world_cam2 = np.dot(R, point_world) + T
    # Step 3: Project world coordinates onto image plane of camera 2
    point_cam2_homogeneous = np.dot(mtx2, point_world_cam2)
    point_cam2_homogeneous /= point_cam2_homogeneous[2]  # Convert to homogeneous coordinates
    point_cam2 = point_cam2_homogeneous[:2]  # Extract (x, y) coordinates
    return point_cam2

cam1, cam2 = find_camera(find_flag = False)
mtx1, dist1 = single_calibration(camera_num = cam1, img_cap = True)
mtx2, dist2 = single_calibration(camera_num = cam2, img_cap = True)
R, T = stereo_calibration(mtx1, dist1, mtx2, dist2, 'stereo_images/*', stereo_capture_flag = True)

cub_cordinate(mtx1, dist1, mtx2, dist2, R, T)

print("$$$ Code Done $$$")