import numpy as np
import cv2 as cv
import os
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def capture_images(cam1, cam2, num_images):
    # Initialize camera objects
    cap1 = cv.VideoCapture(cam1)
    cap2 = cv.VideoCapture(cam2)

    # Check if cameras are opened successfully
    if not (cap1.isOpened() and cap2.isOpened()):
        print("Error: Could not open cameras")
        return [], [], [], 0, 0

    objpoints = []  # 3D object points
    imgpoints1 = []  # 2D image points for camera 1
    imgpoints2 = []  # 2D image points for camera 2

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

    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)

    i = 0
    while i < num_images:
        ret1, frame1 = cap1.read()
        ret2, frame2 = cap2.read()

        height, width, _ = frame1.shape  # Get image dimensions

        if not (ret1 and ret2):
            print("Error: Could not read frames")
            break

        gray1 = cv.cvtColor(frame1, cv.COLOR_BGR2GRAY)
        gray2 = cv.cvtColor(frame2, cv.COLOR_BGR2GRAY)
        # Detect chessboard corners in both images
        ret1, corners1 = cv.findChessboardCorners(gray1, (row, col), None)
        ret2, corners2 = cv.findChessboardCorners(gray2, (row, col), None)
        if ret1 and ret2:
            # Refine corner positions
            corners1 = cv.cornerSubPix(gray1, corners1, (11, 11), (-1, -1), criteria)
            corners2 = cv.cornerSubPix(gray2, corners2, (11, 11), (-1, -1), criteria)

            objpoints.append(objp)
            imgpoints1.append(corners1)
            imgpoints2.append(corners2)

            i += 1

            # Draw and display corners (optional)
            frame1 = cv.drawChessboardCorners(frame1, (row, col), corners1, ret1)
            frame2 = cv.drawChessboardCorners(frame2, (row, col), corners2, ret2)
            stereo_chess = cv.hconcat([frame1, frame2])
            cv.imshow('stereo', stereo_chess)
            cv.waitKey(500)

    # Release video captures
    cap1.release()
    cap2.release()
    cv.destroyAllWindows()

    return objpoints, imgpoints1, imgpoints2, width, height


cam1 = 1
cam2 = 2
# Capture images from both cameras
objpoints, imgpoints1, imgpoints2, w, h = capture_images(cam1, cam2, num_images=40)
print('$$$ capture done $$$')
# Perform stereo calibration
ret, mtx1, dist1, mtx2, dist2, R, T, E, F = cv.stereoCalibrate(objpoints, imgpoints1, imgpoints2, None, None, None, None, (w, h))
print('$$$ calibration done $$$')
print(mtx1, dist1, mtx2, dist2, R, T)

def find_3d_position(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
        point_left = detect_cube(frame1, False)
        point_right = detect_cube(frame2, False)

        # Convert 2D points to homogeneous coordinates
        point_left = np.array([point_left[0], point_left[1]])
        point_right = np.array([point_right[0], point_right[1]])

        # Triangulate 3D point
        P1 = np.hstack((np.eye(3), np.zeros((3, 1))))
        P2 = np.hstack((R, T))
        #print(point_left.T, point_right.T)
        points_4d = cv.triangulatePoints(P1, P2, point_left.T, point_right.T)

        # Convert homogeneous coordinates to Cartesian coordinates
        points_3d = points_4d[:3] / points_4d[3]

        cv.circle(frame1, (int(point_left[0]), int(point_left[1])), 2, (0, 0, 255), -1)
        cv.circle(frame2, (int(point_right[0]), int(point_right[1])), 2, (0, 0, 255), -1)
        #print(points_3d)
        plot_3d_points(points_3d)

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

    return

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)
        cv.circle(frame1, (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 plot_3d_points(points_3d):
    global ax  # Use the ax defined outside the function

    # Clear the previous plot
    ax.clear()

    # Set labels
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    # Scatter plot the new points
    ax.scatter(points_3d[0], points_3d[1], points_3d[2])

    # Set axis limits to ensure origin is visible
    ax.set_xlim([-100, 1000])  # Adjust the limits according to your data range
    ax.set_ylim([-100, 1000])  # Adjust the limits according to your data range
    ax.set_zlim([-100, 1000])  # Adjust the limits according to your data range

    # Draw and pause to update the plot
    plt.draw()
    plt.pause(0.01)

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')

# Initialize data
x_data = []
y_data = []
z_data = []

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