# Name: *** Enter your name here ***
#
# Student-ID: *** Enter your student ID here ***
#
# Rename the file to <yourname>03.py
#
# Submit your solutions via moodle
#
import numpy as np
import matplotlib.pylab as plt

# Please do not change the code above this line

# Exercise 1: Conversion of RGB to grayscale images
#
def exercise1():

    # close all figures
    plt.close('all')

    # Task 1A: Read the image `FluorescentCells01.jpg` and print some
    # information about it such as the *shape* and *data type* of the image, as
    # well as the *minimum* and *maximum* pixel value. 

    rgb_image = plt.imread('FluorescentCells01.jpg')
    
    print('shape:', np.shape(rgb_image))
    print('dtype:', rgb_image.dtype)
    print('  min:', np.min(rgb_image.reshape(-1, 3), axis=0))
    print('  max:', np.max(rgb_image.reshape(-1, 3), axis=0))
    
    # Task 1B: Convert the image to a grayscale image by averaging the pixel
    # values of the individual color layers. The data type of the grayscale
    # image should be identical to the data type of the input image. 

    gray_image = rgb_image.mean(axis=2).astype(rgb_image.dtype)
    
    # Task 1C: To account for differences in the sensitivity of the
    # photoreceptors in our retina, a weighted sum of the R, G, and B
    # components should be formed where the weights for the color channels are
    # 0.2989, 0.5870, and 0.1140, respectively. 

    weights = [0.2989, 0.5870, 0.1140]
    weighted_image = rgb_image.dot(weights).astype(rgb_image.dtype)

    # Task 1D: Show the original RGB image and both grayscale images

    fig, ax = plt.subplots(1, 3, figsize=(12, 3), sharex=True, sharey=True)
    for a in ax.flat:
        a.axis('off')
    ax[0].set_title('RGB image')
    ax[0].imshow(rgb_image)
    
    ax[1].set_title('Unweighted grayscale image')
    ax[1].imshow(gray_image, cmap='gray')

    ax[2].set_title('Weighted grayscale image')
    ax[2].imshow(weighted_image, cmap='gray')
    fig.tight_layout()


# Exercise 2: Showing under- and overflow in a color image
#
def exercise2():

    # close all figures
    plt.close('all')

    # Task 2A: Read the image and show the individual layers as grayscale
    # images. 

    colors = ('red', 'green', 'blue')
    image = plt.imread('FluorescentCells03.jpg')
    layers = np.moveaxis(image, 2, 0)
    
    fig, axes = plt.subplots(1, 3, figsize=(12, 3))
    for ax, layer, color in zip(axes, layers, colors):
        ax.set_title('{} channel'.format(color))
        ax.imshow(layer, cmap='gray')
        ax.axis('off')
    fig.tight_layout()

    # Task 2B: For each channel, show underflowing pixels in blue, overflowing
    # pixels in red, and all other pixels in gray

    info = np.iinfo(image.dtype)    
    lower, upper = info.max * np.array([0.01, 0.99])
    fig, axes = plt.subplots(1, 3, figsize=(12, 3))
    for ax, layer, color in zip(axes, layers, colors):
        ax.set_title('{} channel'.format(color))
        image2 = np.dstack([layer] * 3)
        image2[layer < lower] = [info.min, info.min, info.max]
        image2[layer > upper] = [info.max, info.min, info.min]
        ax.imshow(image2)
        ax.axis('off')
    fig.tight_layout()
    
    
# Exercise 3: Image histograms 
#
def exercise3():

    # close all figures
    plt.close('all')

    colors = ('red', 'green', 'blue')
    image = plt.imread('FluorescentCells03.jpg')
    layers = np.rollaxis(image, 2, 0)
    
    # Task 3A: Read the image and show image histograms for the individual
    # layers using a logarithmic scale for the counts. 

    info = np.iinfo(image.dtype)
    hist_kw = dict(color='k', alpha=0.3, bins=50)
    plt.rc('font', size=16)
    fig, axes = plt.subplots(1, 3, figsize=(12, 4), sharex=True, sharey=True)
    for ax, layer, color in zip(axes, layers, colors):
        ax.set_title('{} channel'.format(color))
        ax.hist(layer.flat, **hist_kw)
        ax.semilogy()
        ax.set_xlim(info.min, info.max)
        ax.set_xlabel('pixel value')
    axes[0].set_ylabel('counts')
    fig.tight_layout()

    # Task 3B: Show image histograms for the individual layers excluding the
    # intensities of pixels suffering from under- or overflow

    fig, axes = plt.subplots(1, 3, figsize=(12, 4), sharex=True, sharey=True)
    for ax, layer, color in zip(axes, layers, colors):
        ax.set_title('{} channel'.format(color))
        selection = (0.01 * info.max < layer) & (layer < 0.99 * info.max)
        ax.hist(layer[selection], **hist_kw)
        ax.set_xlim(info.min, info.max)
        ax.set_xlabel('pixel value')
        ax.yaxis.get_major_formatter().set_powerlimits((0, 0))
    axes[0].set_ylabel('counts')
    fig.tight_layout()


# Test code - you shouldn't change this part of the template

if __name__ == '__main__':

    exercises = {'1': (exercise1, ),
                 '2': (exercise2, ),
                 '3': (exercise3, ), 
                 }

    choice = input('Please enter the number of the exercise: ').upper()
    if choice not in exercises:
        print('Please choose among {0}'.format(
            list(exercises.keys())))
    else:
        func, *args = exercises[choice]
        print('Task {0}, running "{1}"'.format(
            choice, func.__name__))
        result = func(*args)
        if result is not None:
            print('Result:', result)