import cv2
import numpy as np
from scipy.interpolate import griddata

def draw_arrows(image, optical_flow,fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0):
    """
    Draw colored optical flow arrows on sensor image based on flow magnitude.
    
    Args:
        threshold: Minimum flow magnitude to display arrow (default: 10.0)
        grid_spacing: Spacing between arrows in pixels (default: 20)
        arrow_scale: Scale factor for arrow length (default: 3.0)
    
    Returns:
        np.ndarray: 400x400 BGR image with colored arrows (JET colormap)
    
    """
    # Create working copies to avoid modifying originals
    
    # Convert grayscale to BGR to enable colored arrow drawing
    if len(image.shape) == 2 or image.shape[2] == 1:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    
    # Calculate scaling factors between image and flow dimensions
    img_height, img_width = image.shape[:2]
    flow_height, flow_width = optical_flow.shape[:2]
    scale_x, scale_y = img_width / flow_width, img_height / flow_height
    
    # Create sampling grid for flow field (used for interpolation source points)
    sample_step = 19  # Spacing between sample points
    flow_start = np.stack(np.meshgrid(range(0, flow_width, sample_step), range(0, flow_height, sample_step)), 2)
    flow_end = (optical_flow[flow_start[:, :, 1], flow_start[:, :, 0], :] * 2.0 + flow_start).astype(np.int32)
    
    # Create regular grid for arrow placement in image coordinates
    grid_x, grid_y = np.mgrid[grid_spacing//2:flow_height:grid_spacing, grid_spacing//2:flow_width:grid_spacing]
    x_axis = (grid_y.reshape(-1) * scale_x).astype(np.int64)  # Scale to image width
    y_axis = (grid_x.reshape(-1) * scale_y).astype(np.int64)  # Scale to image height
    
    # Interpolate flow vectors from sample grid to regular arrow grid using cubic interpolation
    grid_vector = griddata(
        flow_start.reshape(-1, 2), 
        (flow_end - flow_start).reshape(-1, 2),  # Flow from start to end (original to new position)
        (grid_y, grid_x), method='cubic', fill_value=0
    )
    grid_vector[np.isnan(grid_vector)] = 0  # Replace NaN with zero for stability
    grid_flat_scaled = grid_vector.reshape(-1, 2) * np.array([scale_x, scale_y])
    
    # Compute arrow endpoints and ensure they stay within image bounds
    end_x = (grid_flat_scaled[:, 0] / arrow_scale + x_axis).astype(np.int64)
    end_y = (grid_flat_scaled[:, 1] / arrow_scale + y_axis).astype(np.int64)
    valid_mask = (y_axis >= 0) & (y_axis < img_height) & (x_axis >= 0) & (x_axis < img_width)
    valid_mask &= (end_x >= 0) & (end_x < img_width) & (end_y >= 0) & (end_y < img_height)
    
    # Filter arrows by magnitude threshold and compute colors using JET colormap
    magnitudes = np.linalg.norm(grid_flat_scaled, axis=1)
    valid_mask &= (magnitudes >= 0)  # Only show arrows above threshold
    valid_indices = np.where(valid_mask)[0]
    mag_range = magnitudes.max() - magnitudes.min()
    # Normalize magnitudes to 0-255 range for colormap application, but set below-threshold values to 0 (blue)
    norm_magnitudes = ((magnitudes - magnitudes.min()) / mag_range * 255.0).astype(np.uint8) if mag_range > 0 else np.zeros_like(magnitudes, dtype=np.uint8)
    norm_magnitudes[magnitudes < threshold] = 0  # Keep lowest color (blue) for below-threshold magnitudes
    colors_array = cv2.applyColorMap(norm_magnitudes.reshape(-1, 1), cv2.COLORMAP_JET)  # JET: blue (low) -> red (high)
    
    # Draw arrows at each valid grid point with magnitude-based coloring
    for i in valid_indices:
        color = tuple(map(int, colors_array[i, 0]))
        # Scale arrow length and thickness based on threshold
        if magnitudes[i] < threshold:
            # Below threshold: draw as small dot (zero-length arrow with larger tip)
            cv2.circle(image, (x_axis[i], y_axis[i]), radius=1, color=color, thickness=-1)
        else:
            # Above threshold: draw full arrow with thickness proportional to magnitude
            # Arrow rooted at original position (x_axis, y_axis), head points to new position (end_x, end_y)
            thickness = max(1, int(2 * magnitudes[i] / magnitudes.max())) if magnitudes.max() > 0 else 2
            cv2.arrowedLine(image, (x_axis[i], y_axis[i]), (end_x[i], end_y[i]), color, thickness=thickness, tipLength=0.3)
    return image

def draw_magnitude_map(image, flow_field, threshold=3.0, colormap=cv2.COLORMAP_JET):
    """
    Draw color map of pointwise magnitude of velocity field.
    
    Args:
        info_type: Type of flow field to visualize (VRAW, VNORMAL, VSHEAR, or None for total)
        threshold: Minimum magnitude to display (below threshold shows as zero/blue)
        colormap: OpenCV colormap to use (default: COLORMAP_JET)
    
    Returns:
        np.ndarray: 400x400 BGR image with colored magnitude map
    """
    # Get the appropriate flow field
    # if info_type is None or info_type == self.info.VRAW:
    #     flow_field = self.flow.copy() if self.flow is not None else None
    # elif info_type == self.info.VNORMAL:
    #     flow_field = self.normal.copy() if self.normal is not None else None
    # elif info_type == self.info.VSHEAR:
    #     flow_field = self.shear.copy() if self.shear is not None else None
    # else:
    #     # For combined normal + shear magnitude
    #     if self.normal is not None and self.shear is not None:
    #         flow_field = self.normal + self.shear
    #     else:
    #         flow_field = None
    
    # if flow_field is None:
    #     # Return blank image if no flow data
    #     return np.zeros((self.config_resolution_stitch[0], self.config_resolution_stitch[1], 3), dtype=np.uint8)
    
    # Compute magnitude at each point
    magnitude = np.sqrt(flow_field[:, :, 0]**2 + flow_field[:, :, 1]**2)
    
    # Smooth by downsampling and interpolating back
    downsample_factor = 10
    h, w = magnitude.shape
    magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor]
    magnitude_smoothed = cv2.resize(magnitude_downsampled, (w, h), interpolation=cv2.INTER_CUBIC)
    
    # Apply threshold - set below-threshold values to zero
    magnitude_thresholded = magnitude_smoothed.copy()
    magnitude_thresholded[magnitude_smoothed < threshold] = 0
    
    # Normalize to 0-255 range
    mag_min, mag_max = magnitude_thresholded.min(), magnitude_thresholded.max()
    if mag_max > mag_min:
        magnitude_norm = ((magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255).astype(np.uint8)
    else:
        magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8)
    
    # Apply colormap
    magnitude_colored = cv2.applyColorMap(magnitude_norm, colormap)
    
    # Resize to output size
    # out_img = cv2.resize(magnitude_colored, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), interpolation=cv2.INTER_CUBIC)
    
    return magnitude_colored

def draw_fps(image, fps):
    """
    Draw FPS counter on the top left corner of the image.
    """
    font_scale=0.5
    thickness=2 
    color=(0, 255, 0)
    image = image.copy()
    
    # Draw FPS text
    fps_text = f"FPS: {fps:.1f}"
    position = (10, 30)  # Top left corner with margin
    font = cv2.FONT_HERSHEY_SIMPLEX
    
    # Draw FPS text
    cv2.putText(image, fps_text, position, font, font_scale, color, thickness, cv2.LINE_AA)
    
    return image