#H#######################################################################
# FILENAME :        transformation.py             
#
# DESCRIPTION :
#       This script implements basic image transformations.
#
# AUTHOR :          Jinhang <jinhang.d.zhu@gmail.com>
# START DATE :      07 Mar. 2019
# LAST UPDATE :     12 Mar. 2019
# 
#H#

import numpy as np

#   Transforming an image into a grayscale one
def rgb2gray(source):
    if len(np.shape(source)) == 3:
        return np.dot(source[...,:3], [0.299, 0.587, 0.114])
    elif len(np.shape(source)) == 2:
        print('[olala]  The image is already grayscale!')
        return source

#   Making downsampled image
def downsample_image(gray):
    odd_even = gray[0::2,1::2] # Extract odd rows and even columns
    return odd_even

#   Binarizing the image with a threshold input
def gray2bin(gray, threshold):
    binary = gray[:,:] > threshold
    return binary

#   Shuffle the pixels of an image
def shuffle_image(source):
    im_shuffle = source.ravel().copy()
    np.random.shuffle(im_shuffle)
    im_shuffle=np.reshape(im_shuffle, source.shape)
    return im_shuffle

#   Invert the image grayscale
def invert_image(source):
    a = -1
    b = 255
    invert_image = a*source.copy() + b
    return invert_image

#   Linearize the image grayscale
def linear(source, a, b):
    linearized = a*source.copy() + b
    linearized[linearized > 255] = 255
    linearized[linearized < 0] = 0
    return linearized

#   VERSION 1.5
# #   Calculate the histogram of the image
# def calchist(source):
#     gray = rgb2gray(source)
#     (height,width) = np.shape(gray)
#     hist = np.zeros(256, dtype = int)
#     axis_height = np.arange(0,height,1)
#     axis_width = np.arange(0,width,1)
#     for i in axis_height:
#         for j in axis_width:
#             hist[np.uint8(gray[i,j])] = hist[np.uint(gray[i,j])] + 1
#     hist = hist/np.size(gray)
#     return hist

# #   Calculate the transfer function of equalization
# #   Calculate the cumulative function of the function
# def calcequaltf(source):
#     axis_range = np.arange(0,256)
#     hist = calchist(source)
#     cumu_trans = np.zeros(256)
#     for i in axis_range:
#         if i == 0:
#             cumu_trans[i] = hist[i]
#         else:
#             cumu_trans[i] = cumu_trans[i-1] + hist[i]
#     cumu_trans = cumu_trans * 255
#     return cumu_trans

# #   Histogram Equalization
# def equalization(source):
#     gray = rgb2gray(source)
#     (height,width) = np.shape(gray)
#     axis_height = np.arange(0,height,1)
#     axis_width = np.arange(0,width,1)
#     cumu_trans = calcequaltf(source)
#     equalized = source.copy()
#     for i in axis_height:
#         for j in axis_width:
#             equalized[i,j] = cumu_trans[np.uint8(equalized[i,j])]
#     return equalized

# #   Calculate the transfer function of specification
# def calcmatchtf(source,dest):
#     cumu_source = calcequaltf(source)
#     cumu_dest = calcequaltf(dest)
#     axis_range = np.arange(0,256)
#     # cumu_dest = cumu_dest.astype(np.uint8)
#     # cumu_source = cumu_source.astype(np.uint8)

#     matchtf = np.zeros(256)
#     for i in axis_range:
#         min_diff = abs(cumu_dest[0] - cumu_source[i])
#         for j in axis_range:
#             if abs(cumu_dest[j] - cumu_source[i]) < min_diff:
#                 min_diff = abs(cumu_dest[j] - cumu_source[i])
#                 matchtf[i] = j
#     return matchtf

# #   Histogram specification or matching
# def matching(source,dest):
#     gray = rgb2gray(source)
#     (height,width) = np.shape(gray)
#     axis_height = np.arange(0,height,1)
#     axis_width = np.arange(0,width,1)
#     matchtf = calcmatchtf(source,dest)
#     specified = source.copy()
#     for i in axis_height:
#         for j in axis_width:
#             specified[i,j] = matchtf[np.uint8(specified[i,j])]
#     return specified

#   VERSION 1.6
#   Calculate the histogram of the image
def calchist(source):
    gray = rgb2gray(source)
    (height,width) = np.shape(gray)
    if len(np.shape(source)) is 2:      #   Grayscale
        hist = np.zeros(256, dtype = int)

        # print('Via grayscale, hist shape: ',np.shape(hist))
        axis_height = np.arange(0,height,1)
        axis_width = np.arange(0,width,1)
        for i in axis_height:
            for j in axis_width:
                hist[int(source[i,j])] = hist[int(source[i,j])] + 1
        hist = hist/np.size(source)
        return hist
    elif len(np.shape(source)) is 3:    #   RGB
        hist = np.zeros((3,256), dtype = int)

        # print('Via RGB, hist shape: ',np.shape(hist))
        axis_height = np.arange(0,height,1)
        axis_width = np.arange(0,width,1)
        for i in axis_height:
            for j in axis_width:
                hist[0,int(source[i,j,0])] = hist[0,int(source[i,j,0])] + 1
                hist[1,int(source[i,j,1])] = hist[1,int(source[i,j,1])] + 1
                hist[2,int(source[i,j,2])] = hist[2,int(source[i,j,2])] + 1
        hist = hist/np.size(gray)
        return hist

#   Calculate the transfer function of equalization
#   Calculate the cumulative function of the function
def calcequaltf(source):
    axis_range = np.arange(0,256)
    hist = calchist(source)
    cumu_trans = np.zeros(np.shape(hist))
    for i in axis_range:
        if i == 0:
            cumu_trans[:,i] = hist[:,i]
        else:
            cumu_trans[:,i] = cumu_trans[:,i-1] + hist[:,i]
    cumu_trans = cumu_trans * 255
    return cumu_trans

#   Histogram Equalization
def equalization(source):
    gray = rgb2gray(source)
    (height,width) = np.shape(gray)
    axis_height = np.arange(0,height,1)
    axis_width = np.arange(0,width,1)
    cumu_trans = calcequaltf(source)
    equalized = source.copy()
    if len(np.shape(source)) is 2:      #   Grayscale
        for i in axis_height:
            for j in axis_width:
                equalized[i,j] = cumu_trans[int(equalized[i,j])]
        return equalized

    elif len(np.shape(source)) is 3:    #   RGB
        for i in axis_height:
            for j in axis_width:
                equalized[i,j,0] = cumu_trans[0,int(equalized[i,j,0])]     #   R
                equalized[i,j,1] = cumu_trans[1,int(equalized[i,j,1])]     #   G
                equalized[i,j,2] = cumu_trans[2,int(equalized[i,j,2])]     #   B
        return equalized

#   Calculate the transfer function of specification
def calcmatchtf(source,dest):
    cumu_source = calcequaltf(source)
    cumu_dest = calcequaltf(dest)
    axis_range = np.arange(0,256)
    # cumu_dest = cumu_dest.astype(np.uint8)
    # cumu_source = cumu_source.astype(np.uint8)

    
    if len(np.shape(source)) is 2:      #   Grayscale
        matchtf = np.zeros(256)
        for i in axis_range:
            min_diff = abs(cumu_dest[0] - cumu_source[i])
            for j in axis_range:
                if abs(cumu_dest[j] - cumu_source[i]) < min_diff:
                    min_diff = abs(cumu_dest[j] - cumu_source[i])
                    matchtf[i] = j
        return matchtf

    elif len(np.shape(source)) is 3:    #   RGB
        matchtf = np.zeros((3,256))
        for i in axis_range:
            min_diff = abs(cumu_dest[:,0] - cumu_source[:,i])
            for j in axis_range:
                if abs(cumu_dest[0,j] - cumu_source[0,i]) < min_diff[0]:    #   R
                    min_diff[0] = abs(cumu_dest[0,j] - cumu_source[0,i])
                    matchtf[0,i] = j
            for j in axis_range:
                if abs(cumu_dest[1,j] - cumu_source[1,i]) < min_diff[1]:    #   G
                    min_diff[1] = abs(cumu_dest[1,j] - cumu_source[1,i])
                    matchtf[1,i] = j
            for j in axis_range:
                if abs(cumu_dest[2,j] - cumu_source[2,i]) < min_diff[2]:    #   B
                    min_diff[2] = abs(cumu_dest[2,j] - cumu_source[2,i])
                    matchtf[2,i] = j
        return matchtf

#   Histogram specification or matching
def matching(source,dest):
    gray = rgb2gray(source)
    (height,width) = np.shape(gray)
    axis_height = np.arange(0,height,1)
    axis_width = np.arange(0,width,1)
    matchtf = calcmatchtf(source,dest)
    specified = source.copy()
    if len(np.shape(source)) is 2:      #   Grayscale
        for i in axis_height:
            for j in axis_width:
                specified[i,j] = matchtf[int(specified[i,j])]
        return specified
    
    elif len(np.shape(source)) is 3:    #   RGB
        for i in axis_height:
            for j in axis_width:
                specified[i,j,0] = matchtf[0,int(specified[i,j,0])]    #   R
                specified[i,j,1] = matchtf[1,int(specified[i,j,1])]    #   G
                specified[i,j,2] = matchtf[2,int(specified[i,j,2])]    #   B
        return specified