import numpy as np
from scipy import ndimage
from sklearn.cluster import MiniBatchKMeans


def vectorized_edge_preserving_quantization(
    voxel_matrix, lego_colors, n_initial_clusters=50
):
    """
    Vectorized version of edge-preserving color quantization

    Parameters:
    voxel_matrix: numpy array of shape (x, y, z, 3) containing RGB values
    lego_colors: list of [R, G, B] values for target LEGO colors
    n_initial_clusters: number of initial color clusters before mapping to LEGO colors
    """
    lego_colors = np.array(lego_colors)
    shape = voxel_matrix.shape

    # Reshape to 2D array of pixels
    pixels = voxel_matrix.reshape(-1, 3)

    # Step 1: Initial color clustering using K-means
    kmeans = MiniBatchKMeans(
        n_clusters=n_initial_clusters, batch_size=1000, random_state=42
    )
    labels = kmeans.fit_predict(pixels)
    cluster_centers = kmeans.cluster_centers_

    # Step 2: Create 3D gradient magnitude
    gradients = np.zeros(shape[:3])

    # Compute gradients along each axis
    for axis in range(3):
        # Forward difference
        forward = np.roll(voxel_matrix, -1, axis=axis)
        # Compute color differences
        diff = np.sqrt(np.sum((forward - voxel_matrix) ** 2, axis=-1))
        # Set boundary differences to 0
        slice_idx = [slice(None)] * 3
        slice_idx[axis] = -1
        diff[tuple(slice_idx)] = 0
        gradients += diff

    # Step 3: Segment using watershed algorithm
    # Reshape labels back to 3D
    labels_3d = labels.reshape(shape[:3])

    # Find local minima in gradient magnitude
    markers = ndimage.label(ndimage.minimum_filter(gradients, size=3) == gradients)[0]

    # Apply watershed segmentation
    segments = ndimage.watershed_ift(gradients.astype(np.uint8), markers)

    # Step 4: Map segments to LEGO colors
    # Get mean color for each segment
    segment_colors = ndimage.mean(
        voxel_matrix.reshape(-1, 3),
        labels=segments.ravel(),
        index=np.arange(segments.max() + 1),
    )

    # Find nearest LEGO color for each segment color§
    def find_nearest_lego_colors(colors):
        # Reshape inputs for broadcasting
        colors = colors[:, np.newaxis, :]
        lego_colors_r = lego_colors[np.newaxis, :, :]

        # Compute distances to all LEGO colors at once
        distances = np.sqrt(np.sum((colors - lego_colors_r) ** 2, axis=2))

        # Find index of minimum distance for each color
        nearest_indices = np.argmin(distances, axis=1)

        return lego_colors[nearest_indices]

    segment_lego_colors = find_nearest_lego_colors(segment_colors)

    # Create output array
    result = np.zeros_like(voxel_matrix)

    # Map segments to final colors
    for i, color in enumerate(segment_lego_colors):
        mask = segments == i
        result[mask] = color

    return result


def analyze_quantization(original, quantized):
    """
    Analyze the results of quantization
    """
    original_colors = np.unique(original.reshape(-1, 3), axis=0)
    quantized_colors = np.unique(quantized.reshape(-1, 3), axis=0)

    stats = {
        "original_colors": len(original_colors),
        "quantized_colors": len(quantized_colors),
        "reduction_ratio": len(quantized_colors) / len(original_colors),
    }

    return stats