import numpy as np
import cv2
from collections import Counter


COLOR_NAME = ['black', 'brown', 'blue', 'gray', 'green', 'orange', 'pink', 'purple', 'red', 'white', 'yellow']

def im2c(im, w2c):
    """
    Convert an image to color name representation using a color-name matrix.
    
    Returns:
        numpy.ndarray: Processed image based on color parameter.
    """
    # Define color name mappings
    color_values = np.array([[  0,   0,   0], 
                                [165,  81,  43], 
                                [  0,   0, 255], 
                                [127, 127, 127],
                                [  0, 255,   0], 
                                [255, 127,   0], 
                                [255, 165, 216], 
                                [191,   0, 191],
                                [255,   0,   0], 
                                [255, 255, 255], 
                                [255, 255,   0]], dtype=np.uint8)
    
    
    # Extract RGB channels
    # RR, GG, BB = im[:, :, 0].flatten(), im[:, :, 1].flatten(), im[:, :, 2].flatten()
    
    # Compute index for w2c lookup
    index_im = ((im[:, :, 0].flatten() // 8) + 32 * (im[:, :, 1].flatten()// 8) + 32 * 32 * (im[:, :, 2].flatten() // 8)).astype(np.int32)
    
    # w2cM = np.argmax(w2c, axis=1)
    # name_idx_img = w2cM[index_im].reshape(im.shape[:2])
    
    # max_prob = w2c[np.arange(w2c.shape[0]), w2cM]
    # max_prob_map = max_prob[index_im].reshape(im.shape[:2])
    
    prob_map = w2c[index_im, :].reshape((im.shape[0], im.shape[1], w2c.shape[1]))
    max_prob_map = np.max(prob_map, axis=2)
    name_idx_img = np.argmax(prob_map, axis=2)
    
    color_img = np.zeros_like(im).astype(np.uint8)
        
    for jj in range(im.shape[0]):
        for ii in range(im.shape[1]):
            color_img[jj, ii, :] = np.array(color_values[name_idx_img[jj, ii]])

    return prob_map, max_prob_map, name_idx_img, color_img
    

if __name__ == "__main__":
    # Load w2c matrix
    w2c = np.load('w2c11_j.npy').astype(np.float16)  # Assuming 'w2c.mat' was converted to 'w2c.npy'


    image_path = './test.jpg'
    img = cv2.imread(image_path).astype(np.float32)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    prob_map, max_prob_img, name_idx_img, color_img = im2c(img, w2c)


    filtered_counts = Counter(name_idx_img[name_idx_img <= 10])
    sorted_counts = sorted(filtered_counts.items(), key=lambda x: x[1], reverse=True)
    top_3_values = [num for num, count in sorted_counts[:3]]
    top_3_colors = [COLOR_NAME[i] for i in top_3_values]

    print("Top 3 colors:", top_3_colors)

    cv2.imwrite('colormap_j.jpg', cv2.cvtColor(color_img, cv2.COLOR_BGR2RGB))