import cv2 as cv
import numpy as np

def Planer_Calibration(cam, num_images):
    # Initialize camera objects
    cap = cv.VideoCapture(cam)

    # Check if cameras are opened successfully
    if not (cap.isOpened()):
        print("Error: Could not open cameras")
        return

    objpoints = []  # 3D object points
    imgpoints = []  # 2D image points for camera 1
    # Prepare object points, assuming a chessboard with 9 by 6 squares of 30mm
    square_size = 30  # in millimeters
    row, col = 8, 5
    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:
        ret, frame = cap.read()

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

        if not (ret):
            print("Error: Could not read frames")
            break

        gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
        # Detect chessboard corners in both images
        ret, corners = cv.findChessboardCorners(gray, (row, col), None)
        if ret:
            # Refine corner positions
            corners = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
            objpoints.append(objp)
            imgpoints.append(corners)

            i += 1

            # Draw and display corners (optional)
            frame = cv.drawChessboardCorners(frame, (row, col), corners, ret)
            cv.imshow('Calibrationg', frame)
            cv.waitKey(500)
        else:
            cv.imshow('No Board', frame)
            cv.waitKey(1)

    return objpoints, imgpoints, width, height


def findPositions(cam, duration):
    cap = cv.VideoCapture(cam)
    i = 0
    while i < duration:
        ret, frame = cap.read()
        frame = cv.undistort(frame, mtx, dist)
        point = detect_cube(frame, show_flag = True)
        cv.circle(frame, (int(point[0]), int(point[1])), 2, (0, 0, 255), -1)
        cv.imshow('Frame', frame)
        cv.waitKey(1)

        key = cv.waitKey(5) & 0xFF
        if key == ord('s'):
            if i == 0:
                if offset == 0:
                    offset = point
                else:
                    appended_array = np.vstack((offset, point))
                    # Calculate the average along the first axis (axis=0)
                    offset = np.mean(appended_array, axis=0)
                
                i += 1
    return


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, 50, 50])
    upper = np.array([75, 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)

cam = 1

objpoints, imgpoints, width, height = Planer_Calibration(cam, num_images = 20)
ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(objpoints, imgpoints, (width, height), None, None)
print(mtx, dist)

findPositions(cam, duration = 10)