import numpy as np
import cv2
import os

def generate_point_cloud(img_left_path, img_right_path, output_path='out.ply'):
    # Load images
    print('loading images...')
    imgL = cv2.pyrDown(cv2.imread(img_left_path))  # downscale images for faster processing
    imgR = cv2.pyrDown(cv2.imread(img_right_path))

    # Disparity parameters tuned for 'aloe' image pair
    window_size = 3
    min_disp = 16
    num_disp = 112 - min_disp
    stereo = cv2.StereoSGBM(
        minDisparity=min_disp,
        numDisparities=num_disp,
        SADWindowSize=window_size,
        uniquenessRatio=10,
        speckleWindowSize=100,
        speckleRange=32,
        disp12MaxDiff=1,
        P1=8 * 3 * window_size**2,
        P2=32 * 3 * window_size**2,
        fullDP=False
    )

    # Compute disparity
    print('computing disparity...')
    disp = stereo.compute(imgL, imgR).astype(np.float32) / 16.0

    # Generate 3D point cloud
    print('generating 3d point cloud...')
    h, w = imgL.shape[:2]
    f = 0.8 * w  # Guess for focal length
    Q = np.float32([[1, 0, 0, -0.5 * w],
                    [0, -1, 0, 0.5 * h],  # Turn points 180 deg around x-axis
                    [0, 0, 0, -f],  # so that y-axis looks up
                    [0, 0, 1, 0]])
    points = cv2.reprojectImageTo3D(disp, Q)
    colors = cv2.cvtColor(imgL, cv2.COLOR_BGR2RGB)
    mask = disp > disp.min()
    out_points = points[mask]
    out_colors = colors[mask]

    # Write point cloud to file
    ply_header = '''ply
    format ascii 1.0
    element vertex %(vert_num)d
    property float x
    property float y
    property float z
    property uchar red
    property uchar green
    property uchar blue
    end_header
    '''

    def write_ply(fn, verts, colors):
        verts = verts.reshape(-1, 3)
        colors = colors.reshape(-1, 3)
        verts = np.hstack([verts, colors])
        with open(fn, 'w') as f:
            f.write(ply_header % dict(vert_num=len(verts)))
            np.savetxt(f, verts, '%f %f %f %d %d %d')

    write_ply(output_path, out_points, out_colors)
    print('%s saved' % output_path)

    # Display results
    cv2.imshow('left', imgL)
    cv2.imshow('disparity', (disp - min_disp) / num_disp)
    cv2.waitKey()
    cv2.destroyAllWindows()

def capture_stereo_images(cam1_id, cam2_id, num_images=10, output_folder='captured_images'):
    # Initialize camera objects
    cap1 = cv2.VideoCapture(cam1_id)
    cap2 = cv2.VideoCapture(cam2_id)

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

    # Create output folder if it doesn't exist
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Capture stereo image pairs
    for i in range(num_images):
        ret1, frame1 = cap1.read()
        ret2, frame2 = cap2.read()

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

        # Save images
        cv2.imwrite(f"{output_folder}/left_{i}.jpg", frame1)
        cv2.imwrite(f"{output_folder}/right_{i}.jpg", frame2)

        print(f"Captured pair {i + 1}/{num_images}")

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


capture_stereo_images(1, 2, num_images=10, output_folder='captured_images')

generate_point_cloud('path_to_left_img.jpg', 'path_to_right_img.jpg')