import cv2
import numpy as np

def draw_arrows(frame, optical_flow, fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0, margin = 0.8, colormap=cv2.COLORMAP_JET, below_threshold_color=(0, 0, 255), edge_margin=15):
    """
    Draw colored optical flow arrows on the image by magnitude (JET colormap).
    Frame and optical_flow must have the same height/width; arrows are drawn in pixel coordinates.
    """
    image = frame.copy()
    if len(image.shape) == 2 or image.shape[2] == 1:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)

    h, w = image.shape[:2]
    # Arrow grid in pixel coordinates (assume frame and flow same size)
    rr = np.arange(grid_spacing // 2, h, grid_spacing)
    cc = np.arange(grid_spacing // 2, w, grid_spacing)
    rr, cc = np.meshgrid(rr, cc, indexing="ij")
    rr, cc = rr.reshape(-1), cc.reshape(-1)

    # Displacement at each grid point (flow stored as half-displacement)
    disp = 2.0 * optical_flow[rr, cc, :]
    mag = np.linalg.norm(disp, axis=1)
    dx = disp[:, 0] / arrow_scale
    dy = disp[:, 1] / arrow_scale
    length = np.sqrt(dx * dx + dy * dy)
    # Cap arrow length at grid_spacing
    scale = np.where(length > grid_spacing*margin, grid_spacing / np.maximum(length, 1e-9), 1.0)*margin
    dx, dy = dx * scale, dy * scale
    end_x = (cc + dx).astype(np.int32)
    end_y = (rr + dy).astype(np.int32)

    valid = (
        (cc >= edge_margin) & (cc < w - edge_margin) & (rr >= edge_margin) & (rr < h - edge_margin)
        & (end_x >= 0) & (end_x < w) & (end_y >= 0) & (end_y < h)
    )
    mag_range = mag.max() - mag.min()
    norm_mag = ((mag - mag.min()) / mag_range * 255).astype(np.uint8) if mag_range > 0 else np.zeros_like(mag, dtype=np.uint8)
    norm_mag[mag < threshold] = 0
    colors = cv2.applyColorMap(norm_mag.reshape(-1, 1), colormap)

    for i in np.where(valid)[0]:
        pt1, pt2 = (int(cc[i]), int(rr[i])), (int(end_x[i]), int(end_y[i]))
        color = tuple(int(c) for c in colors[i, 0])
        if mag[i] < threshold:
            cv2.circle(image, pt1, radius=1, color=below_threshold_color, thickness=-1)
        else:
            # thickness = max(1, int(2 * mag[i] / mag.max())) if mag.max() > 0 else 2
            thickness = 1
            cv2.arrowedLine(image, pt1, pt2, color, thickness=thickness, tipLength=0.2)
    return image


def draw_contact(image, contour, centroid, point_color=(0, 255, 0), region_color=(0, 255, 255)):
    """
    在图像上绘制扩张焦点和扩张区域。
    
    Args:
        image: 背景图像
        contour: 轮廓
        centroid: 触碰点
        point_color: 触碰点的颜色 (BGR)
        region_color: 轮廓的颜色 (BGR)
        
    Returns:
        np.ndarray: 绘制了触碰点和轮廓的图像
    """
    image = image.copy()
    
    # 绘制扩张区域轮廓
    if contour is not None:
        cv2.drawContours(image, [contour], -1, region_color, 2)
    
    # 绘制触碰点
    if centroid is not None:
        cx, cy = centroid
        cv2.circle(image, (cx, cy), 8, point_color, -1)
        cv2.circle(image, (cx, cy), 10, point_color, 2)
    
    return image

