import numpy as np


def patch_count(height: int, width: int, patch_size: int) -> int:
    """
    Compute the number of non-overlapping patches produced from an image.

    Args:
        height:     Image height H (must be divisible by patch_size)
        width:      Image width W (must be divisible by patch_size)
        patch_size: Patch size P

    Returns:
        Number of patches N = (H // P) * (W // P)
    """
    # Fill in


def extract_patches(image: np.ndarray, patch_size: int) -> np.ndarray:
    """
    Extract non-overlapping patches from an image in raster-scan order
    (left-to-right, top-to-bottom), following the ViT tokenization approach
    from Dosovitskiy et al. (2020).

    Args:
        image:      Array of shape (C, H, W) — channels-first format
        patch_size: Patch height and width P (H and W are divisible by P)

    Returns:
        patches: Array of shape (N, C * P * P) where N = (H // P) * (W // P).
                 Each row is a flattened patch with all channels concatenated
                 in channel-first, row-major order.
    """
    # Fill in


def add_cls_token(patches: np.ndarray, cls_token: np.ndarray) -> np.ndarray:
    """
    Prepend a learnable CLS token to a sequence of patch embeddings.

    In ViT, a single learnable embedding is prepended to the patch sequence.
    After passing through the Transformer encoder, the output at position 0
    (the CLS token) is used as the image-level feature for classification.

    Args:
        patches:   Array of shape (N, D)
        cls_token: Array of shape (1, D)

    Returns:
        Array of shape (N + 1, D) with cls_token at index 0
    """
    # Fill in


def linear_projection(patches: np.ndarray, weight: np.ndarray) -> np.ndarray:
    """
    Apply a linear projection to map flattened patch vectors into the
    Transformer embedding space.

    This corresponds to the 'Linear Projection of Flattened Patches' step in
    ViT, which projects each patch from dimension C*P*P to the model's
    embedding dimension D.

    Args:
        patches: Array of shape (N, C * P * P)
        weight:  Projection weight matrix of shape (C * P * P, D)

    Returns:
        Array of shape (N, D)
    """
    # Fill in


def add_positional_embedding(
    tokens: np.ndarray, pos_embedding: np.ndarray
) -> np.ndarray:
    """
    Add positional embeddings to a token sequence element-wise.

    Since the Transformer encoder is permutation-invariant, positional
    embeddings encode the spatial location of each patch token so that
    spatial structure is preserved.

    Args:
        tokens:        Array of shape (N + 1, D)
        pos_embedding: Array of shape (N + 1, D)

    Returns:
        Array of shape (N + 1, D)
    """
    # Fill in


def vit_preprocess(
    image: np.ndarray,
    patch_size: int,
    projection_weight: np.ndarray,
    cls_token: np.ndarray,
    pos_embedding: np.ndarray,
) -> np.ndarray:
    """
    Run the complete ViT preprocessing pipeline on a single image.

    Applies all preprocessing steps in sequence:
        1. extract_patches           — tokenize the image into patch vectors
        2. linear_projection         — project patches into the embedding space
        3. add_cls_token             — prepend the learnable classification token
        4. add_positional_embedding  — add spatial position information

    Args:
        image:             Array of shape (C, H, W)
        patch_size:        Patch size P (H and W must be divisible by P)
        projection_weight: Shape (C * P * P, D)
        cls_token:         Shape (1, D)
        pos_embedding:     Shape (N + 1, D) where N = patch_count(H, W, P)

    Returns:
        Array of shape (N + 1, D) — token sequence ready for Transformer input
    """
    # Fill in


if __name__ == "__main__":
    C, H, W, P = 3, 32, 32, 8

    image = np.ones((C, H, W), dtype=np.float32)

    N = patch_count(H, W, P)
    print(f"Image: ({C}, {H}, {W}), patch_size={P}")
    print(f"  patch_count:              {N}")           # Expected: 16

    patches = extract_patches(image, P)
    print(f"  extract_patches shape:    {patches.shape}")  # Expected: (16, 192)

    D = 64
    W_proj = np.eye(C * P * P, D, dtype=np.float32)
    embeddings = linear_projection(patches, W_proj)     # (N, D)
    print(f"  linear_projection shape:  {embeddings.shape}")  # Expected: (16, 64)

    cls_token = np.zeros((1, D))
    tokens = add_cls_token(embeddings, cls_token)
    print(f"  add_cls_token shape:      {tokens.shape}")   # Expected: (17, 64)

    pos_embedding = np.zeros((N + 1, D))
    output = add_positional_embedding(tokens, pos_embedding)
    print(f"  add_positional_embedding: {output.shape}")   # Expected: (17, 64)

    output2 = vit_preprocess(image, P, W_proj, cls_token, pos_embedding)
    print(f"  vit_preprocess shape:     {output2.shape}")  # Expected: (17, 64)
