import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
import cv2

def individual_channel_image(img_arr, channel= 'r', ax=None):
  img_arr = img_arr[:,:,0:3]
  
  if channel in ['r','red','Red']:
    plot_arr = img_arr[:,:,0]
    channel_name = 'Red'
    cmap = 'Reds'
  if channel in ['g','green','Green']:
    plot_arr = img_arr[:,:,1]
    channel_name = 'Green'
    cmap = 'Greens'
  if channel in ['b','blue','Blue']:
    plot_arr = img_arr[:,:,2]
    channel_name = 'Blue'
    cmap = 'Blues'
  
  if channel not in ['r','red','Red','g','green','Green','b','blue','Blue']:
    plot_arr = img_arr 
    channel_name = 'Original'

  if ax is None:
    if channel_name == 'Original':
      # plt.imshow(cv2.cvtColor(img_arr,cv2.COLOR_BGR2RGB))
      plt.imshow(cv2.cvtColor(np.flip(img_arr,axis=-1),cv2.COLOR_BGR2RGB))
    else:
      plt.imshow(plot_arr, cmap = cmap)
      plt.colorbar(orientation= 'vertical', shrink = 0.7, pad = 0.01, fraction = 0.046,)
      plt.title('Image in the {} channel'.format(channel_name))
      plt.show()

  if ax is not None:
    if channel_name == 'Original':
      # ax.imshow(cv2.cvtColor(img_arr,cv2.COLOR_BGR2RGB))
      ax.imshow(cv2.cvtColor(np.flip(img_arr, axis=-1),cv2.COLOR_BGR2RGB))
      # ax.imshow(cv2.cvtColor(img_arr,cv2.COLOR_RGB2BGR))
    else:
      plot = ax.imshow(plot_arr, cmap = cmap)
      plt.colorbar(plot, orientation= 'vertical', ax = ax, fraction = 0.046,)
      # plt.colorbar(orientiation= 'vertical', shrink = 0.7, pad = 0.1)
      ax.set_title('Image in the {} channel'.format(channel_name))


def individual_channel_image_final(img_arr, channel='Red'):
    if channel in ['Red','Green','Blue']:
        fig, ax = plt.subplots(figsize = (15,10))
        individual_channel_image(img_arr, channel= channel)
        plt.tight_layout()
        plt.show()
        fig.canvas.draw()
        image_array = np.array(fig.canvas.renderer.buffer_rgba())
        return image_array

    if channel in ['All']:
        fig, ax = plt.subplots(2,2, figsize = (12,10))
        individual_channel_image(img_arr, channel='r', ax=ax[0,0])
        individual_channel_image(img_arr, channel='g', ax=ax[0,1])
        individual_channel_image(img_arr, channel='b', ax=ax[1,0])
        individual_channel_image(img_arr, channel='full', ax=ax[1,1])
        plt.tight_layout()
        plt.show()
        fig.canvas.draw()
        image_array = np.array(fig.canvas.renderer.buffer_rgba())
        plt.close(fig)
        return image_array

def channel_distribution_plotter(img_array):

    img_array = img_array[:,:,:3] #Not considering the A channel, if it's a RGBA image.

    fig, ax = plt.subplots(figsize=(8,8))
    plt.yticks([])
    plt.xticks([])
    
    plt.subplot(2,2,1)
    plt.hist(img_array[:,:,0].ravel(),bins=256,color='red');
    plt.title("Red Channel")

    plt.subplot(2,2,2)
    plt.hist(img_array[:,:,1].ravel(),bins=256,color='green');
    plt.title("Green Channel")

    plt.subplot(2,2,3)
    plt.hist(img_array[:,:,2].ravel(),bins=256,color='blue');
    plt.title("Blue Channel")

    plt.subplot(2,2,4)
    plt.imshow(cv2.cvtColor(np.flip(img_array, axis=-1),cv2.COLOR_BGR2RGB))
    # plt.imshow(cv2.cvtColor(img_array,cv2.COLOR_RGB2BGR))    
    plt.title("Original Image")

    plt.suptitle("Pixel values distribution in each channel\nx-axis: pixel values, y-axis: number of pixels")
    plt.tight_layout()
    
    plt.show()
    fig.canvas.draw()
    image_array = np.array(fig.canvas.renderer.buffer_rgba())
    return image_array
     


def which_channel_dominates(img_arr, original_image_plot = 'yes', original_image_opacity = 0.3, channel_opacity = 0.7):
    cmap = mcolors.ListedColormap(['red', 'green', 'blue', 'white', 'black', 'gray'])
    img_arr = img_arr[:,:,:3]

    red_channel = img_arr[:,:,0]
    green_channel = img_arr[:,:,1]
    blue_channel = img_arr[:,:,2]


    print("Red channel max. = ", np.max(red_channel), "Green max. = ", np.max(green_channel), "Blue max. = ",np.max(blue_channel))


    which_channel_dominates = np.zeros((img_arr.shape[0],img_arr.shape[1]))

    red_greater_green = np.greater(red_channel,green_channel)
    red_greater_blue = np.greater(red_channel,blue_channel)
    green_greater_blue = np.greater(green_channel,blue_channel)

    #Red is greatest if red is greater than green and blue

    which_channel_dominates[(red_greater_green & red_greater_blue)] = 1
    which_channel_dominates[green_greater_blue & (~red_greater_green)] = 2
    which_channel_dominates[~green_greater_blue & (~red_greater_blue)] = 3

    which_channel_dominates[(red_channel == green_channel) & (red_channel == blue_channel)] = 6
    which_channel_dominates[(red_channel == 255) & (blue_channel == 255) & (green_channel == 255)] = 4
    which_channel_dominates[(red_channel == 0) & (blue_channel == 0) & (green_channel == 0)] = 5

    print("Unique elements of channel dominat array are: - ",np.unique(which_channel_dominates))

    #Map the color code to the image
    fig, ax = plt.subplots(figsize=(8,8))

    if original_image_plot == 'yes':
       plt.imshow(cv2.cvtColor(np.flip(img_arr, axis=-1),cv2.COLOR_BGR2RGB), alpha=original_image_opacity)

    plot = plt.imshow(which_channel_dominates, cmap=cmap, alpha=channel_opacity)

    
    #Customize the ticks of the colorbar
    plt.colorbar(plot, orientation='vertical', 
                ticks = [],
                fraction=0.032, 
                pad=0.04,
                label='Dominant Color Channel'
                )
    
    text = "Which channel dominates in the image below?\nWhite : R=G=B=255, Black : R=G=B=0\nGray : 0 < R=G=B< 255"
    plt.title(text)
    fig.canvas.draw()
    # plt.plot()
    image_array = np.array(fig.canvas.renderer.buffer_rgba())
    return image_array