import cv2
import numpy as np
import matplotlib.pyplot as plt

def information_gravity_filter_and_quiver(
    image_path,
    gravity_percentile=85,
    quiver_step=10,
    arrow_scale=15.0,
    show_original=True
):
    """
    Filter an image by information gravity (gradient magnitude) and show
    directional quiver plot of the gravity field.
    
    Parameters
    ----------
    image_path : str
        Path to input image.
    gravity_percentile : float, optional
        Percentile threshold for information gravity (0-100). Higher = keep only
        strongest gravity regions.
    quiver_step : int, optional
        Plot one arrow every 'quiver_step' pixels (to avoid clutter).
    arrow_scale : float, optional
        Scaling factor for arrow lengths.
    show_original : bool, optional
        Whether to show original image next to filtered/quiver result.
    """
    # Load image and convert to grayscale
    img_bgr = cv2.imread(image_path)
    if img_bgr is None:
        raise FileNotFoundError(f"Could not load image at {image_path}")
    img_gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY).astype(np.float32)

    # ---- Step 1: Compute information gravity (gradient magnitude) ----
    grad_x = cv2.Sobel(img_gray, cv2.CV_32F, 1, 0, ksize=3)
    grad_y = cv2.Sobel(img_gray, cv2.CV_32F, 0, 1, ksize=3)
    gravity = np.hypot(grad_x, grad_y)          # G(x,y)
    
    # Normalize gravity to [0,1] for display
    gravity_norm = gravity / (gravity.max() + 1e-8)
    
    # ---- Step 2: Filter by percentile threshold ----
    threshold = np.percentile(gravity, gravity_percentile)
    mask = gravity >= threshold
    
    # Apply mask to original image (zero out low‑gravity areas)
    filtered_img = img_gray.copy()
    filtered_img[~mask] = 0
    
    # ---- Step 3: Prepare quiver grid (subsampled) ----
    h, w = gravity.shape
    y_grid, x_grid = np.mgrid[0:h:quiver_step, 0:w:quiver_step]
    
    # Sample gradient components at grid points
    u = grad_x[y_grid, x_grid]   # x‑component (dx)
    v = grad_y[y_grid, x_grid]   # y‑component (dy)
    
    # Optional: length scaling
    magnitude = np.hypot(u, v)
    # Avoid division by zero
    u_norm = np.divide(u, magnitude + 1e-8)
    v_norm = np.divide(v, magnitude + 1e-8)
    u_plot = u_norm * arrow_scale
    v_plot = -v_norm * arrow_scale  # Negate y-component to align with inverted y-axis of imshow
    
    # ---- Step 4: Visualisation ----
    # Mask grid points to only show arrows in high‑gravity regions
    mask_sub = mask[y_grid, x_grid]
    x_points = x_grid[mask_sub]
    y_points = y_grid[mask_sub]
    u_points = u_plot[mask_sub]
    v_points = v_plot[mask_sub]
    
    if show_original:
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        axes[0].imshow(img_gray, cmap='gray')
        axes[0].set_title("Original Grayscale")
        axes[0].axis('off')
        
        axes[1].imshow(gravity_norm, cmap='inferno')
        axes[1].set_title(f"Information Gravity (G)\nThreshold at {gravity_percentile}th percentile")
        axes[1].axis('off')
        
        axes[2].imshow(filtered_img, cmap='gray')
        axes[2].quiver(x_points, y_points, u_points, v_points,
                       angles='xy', scale_units='xy', scale=1, alpha=0.8,
                       color='cyan', headwidth=3, headlength=5, headaxislength=4.5)
        axes[2].set_title("Filtered Image + Directional Quiver")
        axes[2].axis('off')
    else:
        plt.figure(figsize=(8, 8))
        plt.imshow(filtered_img, cmap='gray')
        plt.quiver(x_points, y_points, u_points, v_points,
                   angles='xy', scale_units='xy', scale=1, alpha=0.8,
                   color='lime', headwidth=3, headlength=5, headaxislength=4.5)
        plt.title("Filtered Image with Information Gravity Direction")
        plt.axis('off')
    
    plt.tight_layout()
    plt.show()
    
    # Return data for further analysis
    return {
        'gravity': gravity,
        'mask': mask,
        'filtered_image': filtered_img,
        'gradient_x': grad_x,
        'gradient_y': grad_y
    }

# Example usage
if __name__ == "__main__":
    # Replace with your image path
    result = information_gravity_filter_and_quiver(
        "input.png",
        gravity_percentile=85,
        quiver_step=15,
        arrow_scale=15.0,
        show_original=True
    )