import numpy as np
import cv2
from matplotlib import pyplot as plt
import orisys

sensor = orisys.Sensor(vid_src='stereo_1.mp4',config_name='./config/stereo.json',mirror=False)
print(sensor.boxes)
sensor.get_img()
tiles = sensor.stitcher.tile_list
print(tiles)

imgL = tiles[0]
imgR = tiles[1]

 
# imgL = cv.imread('tsukuba_l.png', cv.IMREAD_GRAYSCALE)
# imgR = cv.imread('tsukuba_r.png', cv.IMREAD_GRAYSCALE)

plt.subplot(1,2,1)
plt.imshow(imgL,'gray')
plt.subplot(1,2,2)
plt.imshow(imgR,'gray')
plt.show()
 
# Feature matching for stereo computation
def match_features(img1, img2, method='SIFT', max_matches=50):
    """
    Detect and match features between two images.
    
    Args:
        img1: Left image (grayscale)
        img2: Right image (grayscale)
        method: 'SIFT' or 'ORB'
        max_matches: Maximum number of matches to show
    
    Returns:
        img_matches: Image with matches drawn
        matches: List of matches
        kp1, kp2: Keypoints in img1 and img2
    """
    # Convert to grayscale if needed
    if len(img1.shape) == 3:
        img1_gray = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    else:
        img1_gray = img1.copy()
    
    if len(img2.shape) == 3:
        img2_gray = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    else:
        img2_gray = img2.copy()
    
    # Initialize detector
    if method == 'SIFT':
        detector = cv2.SIFT_create()
    elif method == 'ORB':
        detector = cv2.ORB_create()
    else:
        raise ValueError("method must be 'SIFT' or 'ORB'")
    
    # Detect keypoints and descriptors
    kp1, des1 = detector.detectAndCompute(img1_gray, None)
    kp2, des2 = detector.detectAndCompute(img2_gray, None)
    
    if des1 is None or des2 is None or len(des1) == 0 or len(des2) == 0:
        print(f"Warning: No features detected with {method}")
        return None, [], kp1, kp2
    
    # Match features
    if method == 'SIFT':
        # Check if we have enough descriptors for knnMatch (need at least k=2)
        min_descriptors = min(len(des1), len(des2))
        if min_descriptors < 2:
            print(f"Warning: Not enough descriptors ({min_descriptors}) for knnMatch, using regular matching")
            # Fallback to BFMatcher
            matcher = cv2.BFMatcher(cv2.NORM_L2, crossCheck=False)
            matches = matcher.match(des1, des2)
            matches = sorted(matches, key=lambda x: x.distance)[:max_matches]
        else:
            # Use FLANN matcher for SIFT
            try:
                FLANN_INDEX_KDTREE = 1
                index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
                search_params = dict(checks=50)
                matcher = cv2.FlannBasedMatcher(index_params, search_params)
                matches = matcher.knnMatch(des1, des2, k=2)
                # Apply Lowe's ratio test
                good_matches = []
                for match_pair in matches:
                    if len(match_pair) == 2:
                        m, n = match_pair
                        if m.distance < 0.7 * n.distance:
                            good_matches.append(m)
                matches = sorted(good_matches, key=lambda x: x.distance)[:max_matches]
            except cv2.error:
                # Fallback to BFMatcher if FLANN fails
                print("Warning: FLANN matcher failed, using BFMatcher instead")
                matcher = cv2.BFMatcher(cv2.NORM_L2, crossCheck=False)
                matches = matcher.match(des1, des2)
                matches = sorted(matches, key=lambda x: x.distance)[:max_matches]
    else:  # ORB
        # Use BF matcher for ORB
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        matches = matcher.match(des1, des2)
        matches = sorted(matches, key=lambda x: x.distance)[:max_matches]
    
    # Draw matches
    img_matches = cv2.drawMatches(
        img1_gray, kp1, img2_gray, kp2, matches, None,
        flags=cv2.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS
    )
    
    print(f"Found {len(matches)} good matches using {method}")
    return img_matches, matches, kp1, kp2

# Perform feature matching
print("\n=== Feature Matching ===")
img_matches, matches, kp1, kp2 = match_features(imgL, imgR, method='SIFT', max_matches=50)

if img_matches is not None:
    plt.figure(figsize=(15, 6))
    plt.imshow(img_matches)
    plt.title(f'Feature Matches (SIFT) - {len(matches)} matches')
    plt.axis('off')
    plt.tight_layout()
    plt.show()
    
    # Also create a side-by-side visualization with lines connecting matches
    h1, w1 = imgL.shape[:2]
    h2, w2 = imgR.shape[:2]
    h_max = max(h1, h2)
    combined = np.zeros((h_max, w1 + w2), dtype=np.uint8)
    combined[:h1, :w1] = imgL
    combined[:h2, w1:w1+w2] = imgR 
    combined_color = combined
    
    # Draw lines connecting matched features
    for match in matches[:30]:  # Show top 30 matches
        pt1 = tuple(map(int, kp1[match.queryIdx].pt))
        pt2 = tuple(map(int, kp2[match.trainIdx].pt))
        pt2 = (pt2[0] + w1, pt2[1])
        cv2.line(combined_color, pt1, pt2, (0, 255, 0), 1)
        cv2.circle(combined_color, pt1, 3, (0, 0, 255), -1)
        cv2.circle(combined_color, pt2, 3, (255, 0, 0), -1)
    
    plt.figure(figsize=(15, 6))
    plt.imshow(cv2.cvtColor(combined_color, cv2.COLOR_BGR2RGB))
    plt.title(f'Stereo Feature Matching - {len(matches[:30])} matches shown')
    plt.axis('off')
    plt.tight_layout()
    plt.show()

# Stereo disparity computation
print("\n=== Stereo Disparity Computation ===")
stereo = cv2.StereoBM.create(numDisparities=32, blockSize=15)
disparity = stereo.compute(imgL, imgR)
plt.figure(figsize=(12, 5))
plt.imshow(disparity, cmap='jet')
plt.colorbar(label='Disparity')
plt.title('Stereo Disparity Map')
plt.axis('off')
plt.tight_layout()
plt.show()