import UZ_utils as uz
import numpy as np
import numpy.typing as npt
from matplotlib import pyplot as plt
import random
from PIL import Image

#######################################
# EXCERCISE 1: Basic image processing #
#######################################

def excercise_one() -> None:
    image = one_a()
    #one_b(image)
    #one_c(image)
    #one_d(100, 200, 50, 200, image)
    one_e()
    
def one_a() -> npt.NDArray[np.float64]:
    """
    Read the image from the file umbrellas.jpg and display it
    """
    image = uz.imread('./images/umbrellas.jpg')
    uz.imshow(image, 'Umbrellas')
    return image

def one_b(image: npt.NDArray[np.float64]) -> None:
    """
    Convert the loaded image to grayscale. A very simple way of doing this is summing
    up the color channels and dividing the result by 3, effectively averaging the values.
    The issue, however, is that the sum easily reaches beyond the np.uint8 range. We
    can avoid that by casting the data to a floating point type. You can access a specific
    image channel using the indexing syntax like red = I[:,:,0].
    """
    grayscale_image = np.zeros(image.shape[:2])
    
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            grayscale_image[i, j] = (image[i, j, 0] + image[i,j, 1] + image[i, j, 2]) / 3

    uz.imshow(grayscale_image, 'Umbrellas grayscale')

def one_c(image: npt.NDArray[np.float64]) -> None:
    """
    Cut and display a specific part of the loaded image. Extract only one of the channels
    so you get a grayscale image. You can do this by indexing along the first two axes,
    for instance: cutout=I[130:260, 240:450, 1]. You can display multiple images in
    a single window using plt.subplot().
    Grayscale images can be displayed using different mappings (on a RGB monitor,
    every value needs to be mapped to a RGB triplet). Pyplot defaults to a color map
    named viridis, but often it is preferable to use a grayscale color map. This can be set
    with an additional argument to plt.imshow, like plt.imshow(I, cmap=’gray’).
    Question: Why would you use different color maps?
    Answer:
    """
    uz.imshow(image[50:200, 100:400, 2], "Just one piece of umbrellas")

def one_d(startx: int, endx: int, starty: int, endy: int, image:npt.NDArray[np.float64]) -> None:
    """
    (height, width, color)
    (x     , y    , color)
    y ->
    ################# x
    #               # |
    #               # v 
    #               #
    #               #
    #               #
    #################
    You can also replace only a part of the image using indexing. Write a script that
    inverts a rectangular part of the image. This can be done pixel by pixel in a loop or
    by using indexing.
    Question: How is inverting a grayscale value defined for uint8 ?
    Answer:
    """
    inverted_image = image.copy() 

    for i in range(startx, endx):
        for j in range(starty, endy):
            inverted_image[i, j, 0] = 1 - image[i, j, 0]
            inverted_image[i, j, 1] = 1 - image[i, j, 1]
            inverted_image[i, j, 2] = 1 - image[i, j, 2]

    fig, (ax0, ax1) = plt.subplots(1, 2)
    fig.suptitle("Lomberlini")


    ax0.imshow(image, cmap="gray") 
    ax1.imshow(inverted_image, vmax=255, cmap="gray")

    ax0.set(title="Original image")
    ax1.set(title="Inverted image")

    plt.show()

def one_e() -> None:
    """
    Perform a reduction of grayscale levels in the image. First read the image from
    umbrellas.jpg and convert it to grayscale. You can write your own function for
    grayscale conversion or use the function in UZ_utils.py.
    Convert the grayscale image to floating point type. Then, rescale the image values
    so that the largest possible value is 63. Convert the image back to uint8 and display
    both the original and the modified image. Notice that both look the same. Pyplot
    tries to maximize the contrast in displayed images by checking their values and
    scaling them to cover the entire uint8 interval. If you want to avoid this, you need
    to set the maximum expected value when using plt.imshow(), like plt.imshow(I,
    vmax=255. Use this to display the resulting image so the change is visible.
    """
    grayscale_image = uz.imread_gray("./images/umbrellas.jpg", "float64")
    upscaled_grayscale_image = (grayscale_image.copy() * 63).astype(np.uint8)

    fig, (ax0, ax1) = plt.subplots(1, 2)
    fig.suptitle("Lomberlini")


    ax0.imshow(grayscale_image, cmap="gray") 
    ax1.imshow(upscaled_grayscale_image, vmax=255, cmap="gray")

    ax0.set(title="Original grayscale image")
    ax1.set(title="Upscaled grayscale image")

    plt.show()

    
############################################
# EXCERCISE 2: Thresholding and histograms #
############################################

def excercise_two() -> None:
    """
    Thresholding an image is an operation that produces a binary image (mask) of the same
    size where the value of pixels is determined by whether the value of the corresponding
    pixels in the source image is greater or lower than the given threshold.
    """
    #two_a()
    two_b('./images/bird.jpg', 100, 20)
    two_d()
    two_e(uz.imread_gray("./images/bird.jpg", "float64"))


def two_a() -> tuple[npt.NDArray[np.float64], npt.NDArray[np.uint8]]:
    """
    Create a binary mask from a grayscale image. The binary mask is a matrix the same
    size as the image which contains 1 where some condition holds and 0 everywhere
    else. In this case the condition is simply the original image intensity. Use the image
    bird.jpg. Display both the image and the mask.
    """
    random_number = random.random()
    TRESHOLD = 0.4

    image = uz.imread_gray("./images/bird.jpg", "float64")
    binary_mask = image.copy()
 
    if random_number < 0.5:
        binary_mask[binary_mask < TRESHOLD] = 0
        binary_mask[binary_mask >= TRESHOLD] = 1
    else:
        binary_mask = np.where(binary_mask < TRESHOLD, 0, 1)
    
    binary_mask = uz.convert_float64_array_to_uint8_array(binary_mask) 

    fig, (ax0, ax1) = plt.subplots(1, 2)
    fig.suptitle("Birdie and its mask")


    ax0.imshow(image, cmap="gray") 
    ax1.imshow(binary_mask, cmap="gray")

    ax0.set(title="Original image")
    ax1.set(title="Mask of birdie")

    plt.show()
    return (image, binary_mask)


def my_hist_for_loop(image: npt.NDArray[np.float64], number_of_bins: int) -> npt.NDArray[np.float64]:
    bin_restrictions = np.arange(0, 1, 1 / number_of_bins)
    bins = np.zeros(number_of_bins).astype(np.float64)

    for pixel in image.reshape(-1):
        # https://stackoverflow.com/a/16244044
        bins[np.argmax(bin_restrictions > pixel)] += 1

    return bins / np.sum(bins) 

# Much faster implementation than for loop
def my_hist(image: npt.NDArray[np.float64], number_of_bins: int) -> npt.NDArray[np.float64]:
    bins = np.arange(0, 1, 1 / number_of_bins)
    # Put pixels into classes
    # ex. binsize = 10 then 0.4 would map into 4
    binarray = np.digitize(image.reshape(-1), bins).astype(np.uint8)

    # Now count those values 
    binarray = np.unique(binarray, return_counts=True)

    counts = binarray[1].astype(np.float64) # Get the counts out of tuple

    # Check if there is any empty bin
    empty_bins = []
    bins = binarray[0]
    for i in range(1, number_of_bins + 1):
        if i not in bins:
            empty_bins.append(i)
    
    # Add empty bins with zeros
    if empty_bins != []:
        for i in empty_bins:
            counts = np.insert(counts, i - 1, 0)

    return counts / np.sum(counts)
    

def two_b(image_path: str, number_of_bins_first: int, number_of_bins_second: int) -> None:
    """
    Write a function myhist that accepts a grayscale image and the number of bins that
    will be used in building a histogram. The function should return a 1D array that
    represents the image histogram (the size should be equal to the number of bins, of
    course).
    The histogram is simply a count of pixels with same (or similar) intensity for all
    bins. You can assume the values of the image are within the interval [0,255]. If you
    use fewer than 255 bins, intensities will have to be grouped together, e.g. if using 10
    bins, all values on the interval [0,25] will fall into bin 0. 

    Write a script that calculates and displays histograms for different numbers of bins
    using bird.jpg
    """
    image = uz.imread_gray(image_path, "float64")


    H1 = my_hist(image, number_of_bins_first)
    H2 = my_hist(image, number_of_bins_second)

    fig, (ax0, ax1, ax2) = plt.subplots(1, 3)
    fig.suptitle("Birdie and histgrams")

    ax0.imshow(image, cmap="gray") 
    ax0.set(title="Birdie image")
    ax1.bar(np.arange(number_of_bins_first), H1)
    ax1.set(title="100 bins")
    ax2.bar(np.arange(number_of_bins_second), H2)
    ax2.set(title="20 bins")

    plt.show()

def two_c() -> None:
    print("to be implemented", "to be implemented")

def two_d() -> None:
    """
    Test myhist function on images (three or more) of the same scene in
    different lighting conditions. One way to do this is to capture several images using
    your web camera and change the lighting of the room. Visualize the histograms for
    all images for different number of bins and interpret the results.
    """
    light = uz.imread_gray("./images/ROOM_LIGHTS_ON.jpg", "float64")
    darker = uz.imread_gray("./images/ONE_ROOM_LIGH_ON.jpg", "float64")
    dark = uz.imread_gray("./images/DARK.jpg", "float64")

    H10 = my_hist(light, 20)
    H11 = my_hist(light, 60)
    H12 = my_hist(light, 100)
    H20 = my_hist(darker, 20)
    H21 = my_hist(darker, 60)
    H22 = my_hist(darker, 100)
    H30 = my_hist(dark, 20)
    H31 = my_hist(dark, 60) 
    H32 = my_hist(dark, 100)
    
    fig, axs = plt.subplots(3, 4)
    fig.suptitle("spanskiduh and histgrams")

    axs[0, 0].imshow(light, cmap="gray") 
    axs[0, 0].set(title="Image in light conditions")
    axs[0, 1].bar(np.arange(20), H10)
    axs[0, 1].set(title="Using 20 bins")
    axs[0, 2].bar(np.arange(60), H11)
    axs[0, 2].set(title="Using 60 bins")
    axs[0, 3].bar(np.arange(100), H12)
    axs[0, 3].set(title="Using 100 bins")

    axs[1, 0].imshow(darker, cmap="gray") 
    axs[1, 0].set(title="Image in darker conditions")
    axs[1, 1].bar(np.arange(20), H20)
    axs[1, 1].set(title="Using 20 bins")
    axs[1, 2].bar(np.arange(60), H21)
    axs[1, 2].set(title="Using 60 bins")
    axs[1, 3].bar(np.arange(100), H22)
    axs[1, 3].set(title="Using 100 bins")

    axs[2, 0].imshow(dark, cmap="gray") 
    axs[2, 0].set(title="Image in dark conditions")
    axs[2, 1].bar(np.arange(20), H30)
    axs[2, 1].set(title="Using 20 bins")
    axs[2, 2].bar(np.arange(60), H31)
    axs[2, 2].set(title="Using 60 bins")
    axs[2, 3].bar(np.arange(100), H32)
    axs[2, 3].set(title="Using 100 bins")
    
    plt.show()

def two_e(image: npt.NDArray[np.float64]) -> None:
    """
    Implement Otsu’s method for automatic threshold calculation. It
    should accept a grayscale image and return the optimal threshold. Using normalized
    histograms, the probabilities of both classes are easy to calculate. Write a script that
    shows the algorithm’s results on different images.
    References: https://en.wikipedia.org/wiki/Otsu%27s_method
    """
    treshold_range = np.arange(np.max(image) + 1) 
    criterias = []

    for treshold in treshold_range:
        # create the thresholded image
        thresholded_im = np.zeros(image.shape)

        thresholded_im[image >= treshold] = 1

        # compute weights
        nb_pixels = image.size
        nb_pixels1 = np.count_nonzero(thresholded_im)
        weight1 = nb_pixels1 / nb_pixels
        weight0 = 1 - weight1

        # if one the classes is empty, eg all pixels are below or above the threshold, that threshold will not be considered
        # in the search for the best threshold
        if weight1 == 0 or weight0 == 0:
            continue 

        # find all pixels belonging to each class
        val_pixels1 = image[thresholded_im == 1]
        val_pixels0 = image[thresholded_im == 0]

        # compute variance of these classes
        var0 = np.var(val_pixels0) if len(val_pixels0) > 0 else 0
        var1 = np.var(val_pixels1) if len(val_pixels1) > 0 else 0

        criterias.append( weight0 * var0 + weight1 * var1)

    best_threshold = treshold_range[np.argmin(criterias)]
    print(best_threshold)
    print("h")

def main() -> None:
    #excercise_one()
    excercise_two()

if __name__ == "__main__":
    main()