import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from scipy.ndimage import convolve
from scipy.linalg import lstsq

# Step 1: Preprocess images (keeping them in [0, 1] range)
def preprocess_image(image_path):
    """Load and preprocess image (grayscale and normalized)."""
    image = Image.open(image_path).convert('L')  # Convert to grayscale
    return np.array(image) / 255.0  # Normalize to [0, 1]

# Step 2: Compute Optical Flow using Lucas-Kanade Method
def opticalFlow(I1, I2, windowSize, tau):
    """
    Computes optical flow using Lucas-Kanade method.
    
    Parameters:
        I1, I2 : numpy arrays
            Grayscale images at t=1 and t=2, normalized to [0, 1].
        windowSize : int
            Size of the window for local calculations.
        tau : float
            Threshold for minimum eigenvalue of ATA.
    
    Returns:
        u, v : numpy arrays
            Optical flow components in x and y directions.
        hitMap : numpy array
            Binary image indicating valid corner regions.
    """
    assert windowSize % 2 == 1, "Window size must be odd."
    half_win = windowSize // 2
    
    # Compute image gradients using Sobel filter
    kernel_x = np.array([[-1, 1], [-1, 1]]) * 0.5
    kernel_y = np.array([[-1, -1], [1, 1]]) * 0.5
    kernel_t = np.ones((2, 2)) * 0.25
    
    Ix = convolve(I1, kernel_x) + convolve(I2, kernel_x)
    Iy = convolve(I1, kernel_y) + convolve(I2, kernel_y)
    It = convolve(I2, kernel_t) - convolve(I1, kernel_t)
    
    u = np.zeros_like(I1)
    v = np.zeros_like(I1)
    hitMap = np.zeros_like(I1, dtype=bool)
    
    for i in range(half_win, I1.shape[0] - half_win):
        for j in range(half_win, I1.shape[1] - half_win):
            Ix_win = Ix[i-half_win:i+half_win+1, j-half_win:j+half_win+1].flatten()
            Iy_win = Iy[i-half_win:i+half_win+1, j-half_win:j+half_win+1].flatten()
            It_win = It[i-half_win:i+half_win+1, j-half_win:j+half_win+1].flatten()
            
            A = np.stack((Ix_win, Iy_win), axis=1)
            b = -It_win
            
            ATA = A.T @ A
            eigenvalues = np.linalg.eigvals(ATA)
            
            if min(eigenvalues) >= tau:
                flow = lstsq(ATA, A.T @ b, cond=None)[0]
                u[i, j], v[i, j] = flow
                hitMap[i, j] = True
    
    return u, v, hitMap

# Step 3: Display Optical Flow Results and hitMap for each window size
def display_results(u, v, hitMap, I1, window_size):
    """Display optical flow results and hitmap with correctly scaled flow vectors."""
    
    # Create a white canvas to display the flow vectors (white background)
    flow_image = np.ones_like(I1)  # White background with values [1.0] (in [0, 1] range)
    
    # Create a new figure for each window size
    plt.figure(figsize=(8, 8))
    plt.imshow(flow_image, cmap='gray')  # White canvas for flow vectors
    plt.title(f"Optical Flow Vectors (Window Size {window_size})")
    
    # Plot the flow vectors on the white canvas in blue color
    plt.quiver(np.arange(0, u.shape[1], step=5), np.arange(0, v.shape[0], step=5),
               u[::5, ::5], v[::5, ::5], color='blue', scale=50, width=0.005)  # Adjust scale and density
    plt.gca().invert_yaxis()  # To match the image coordinate system
    plt.show()

    # Display hitMap as binary (black and white) image
    plt.figure(figsize=(6, 6))
    plt.imshow(hitMap, cmap='gray')  # Black and white image (binary)
    plt.title(f"Valid Areas (HitMap) - Window Size {window_size}")
    plt.axis('off')  # Hide axes
    plt.show()

# Example usage
I1_path = "/content/coridor/hw4/data/corridor/bt.000.png"
I2_path = "/content/coridor/hw4/data/corridor/bt.001.png"

# Load images
I1 = preprocess_image(I1_path)
I2 = preprocess_image(I2_path)

# Set parameters
tau = 0.01  # Threshold for minimum eigenvalue of ATA
window_sizes = [15, 31, 101]  # Three different window sizes (updated to be odd)

# For each window size, compute and display results
for window_size in window_sizes:
    u, v, hitMap = opticalFlow(I1, I2, window_size, tau)
    display_results(u, v, hitMap, I1, window_size)