import cv2
import numpy as np


def _prepare_draw_mask(valid_mask, shape):
    """Resize a boolean mask to (height, width) for arrow culling."""
    if valid_mask is None:
        return None
    h, w = shape[:2]
    mask = np.asarray(valid_mask, dtype=bool)
    if mask.shape != (h, w):
        mask_u8 = mask.astype(np.uint8)
        if mask_u8.ndim > 2:
            mask_u8 = mask_u8[:, :, 0]
        mask = cv2.resize(mask_u8, (w, h), interpolation=cv2.INTER_NEAREST).astype(bool)
    return mask


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,
    valid_mask=None,
):
    """
    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.

    Parameters
    ----------
    valid_mask : array-like, optional
        Boolean (H, W) mask. Grid points outside the mask are skipped (no arrows or dots).
    """
    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]
    mask = _prepare_draw_mask(valid_mask, (h, w))

    # 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)
    )
    if mask is not None:
        valid &= mask[rr, cc]
        # Arrow tips may fall outside the frame (the length cap is applied
        # before the bounds check above), so clamp before indexing: those points
        # are already False in ``valid``, and in-bounds indices are unchanged.
        valid &= mask[np.clip(end_y, 0, h - 1), np.clip(end_x, 0, w - 1)]

    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 = 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