import cv2
import numpy as np
import matplotlib.pyplot as plt
import sys


def DerivGauss(im, sigma=0.5):
    gaussian_die_off = 0.000001
    var = sigma ** 2
    # compute filter width
    width = None
    for i in range(1, 51):
        if np.exp(-(i ** 2) / (2 * var)) > gaussian_die_off:
            width = i
    if width is None:
        width = 1

    # create filter (derivative of Gaussian filter)
    x = np.arange(-width, width + 1)
    y = np.arange(-width, width + 1)
    coordinates = np.meshgrid(x, y)
    x = coordinates[0]
    y = coordinates[1]
    derivate_gaussian_2D = -x * \
        np.exp(-(x * x + y * y) / (2 * var)) / (var * np.pi)

    # apply filter and return magnitude
    ax = cv2.filter2D(im, -1, derivate_gaussian_2D)
    ay = cv2.filter2D(im, -1, np.transpose(derivate_gaussian_2D))
    magnitude = np.sqrt((ax ** 2) + (ay ** 2))
    return magnitude


def GPconstancy_GI(im, gray_pixels, delta_th=10**(-4)):

    # mask saturated pixels and mask very dark pixels
    mask = np.logical_or(np.max(im, axis=2) >= 0.95,
                         np.sum(im, axis=2) <= 0.0315)

    # remove noise with mean filter
    # mean_kernel = np.ones((7, 7), np.float32) / 7**2
    # im = cv2.filter2D(im, -1, mean_kernel)

    # decompose rgb values
    r = im[:, :, 0]
    g = im[:, :, 1]
    b = im[:, :, 2]

    # mask 0 elements
    mask = np.logical_or.reduce((mask, r == 0, g == 0, b == 0))

    # replace 0 values with machine epsilon
    eps = np.finfo(np.float32).eps
    r[r == 0] = eps
    g[g == 0] = eps
    b[b == 0] = eps
    norm = r + g + b

    # mask low contrast pixels
    delta_r = DerivGauss(r)
    delta_g = DerivGauss(g)
    delta_b = DerivGauss(b)
    mask = np.logical_or(mask, np.logical_and.reduce(
        (delta_r <= delta_th, delta_g <= delta_th, delta_b <= delta_th)))

    # compute colors in log domain, only red and blue
    log_r = np.log(r) - np.log(norm)
    log_b = np.log(b) - np.log(norm)

    # mask low contrast pixels in the log domain
    delta_log_r = DerivGauss(log_r)
    delta_log_b = DerivGauss(log_b)
    mask = np.logical_or.reduce(
        (mask, delta_log_r == np.inf, delta_log_b == np.inf))

    # normalize each channel in log domain
    data = np.concatenate(
        (np.reshape(delta_log_r, (-1, 1)), np.reshape(delta_log_b, (-1, 1))), axis=1)
    mink_norm = 2
    norm2_data = np.sum(data ** mink_norm, axis=1) ** (1 / mink_norm)
    map_uniquelight = np.reshape(norm2_data, delta_log_r.shape)

    # make masked pixels to max value
    map_uniquelight[mask] = np.max(map_uniquelight)

    # denoise
    # map_uniquelight = cv2.filter2D(map_uniquelight, -1, mean_kernel)

    # filter using map_uniquelight
    gray_index_unique = map_uniquelight
    sort_unique = np.sort(gray_index_unique.flatten())
    gindex_unique = np.full(gray_index_unique.shape, False, dtype=bool)
    gindex_unique[gray_index_unique <= sort_unique[gray_pixels - 1]] = True
    choosen_pixels = im[gindex_unique]
    mean = np.mean(choosen_pixels, axis=0)
    result = mean / np.apply_along_axis(np.linalg.norm, 0, mean)
    return result


if __name__ == "__main__":

    # read image and convert to 0 1
    im_path = 'example.png'
    im = cv2.cvtColor(cv2.imread(im_path), cv2.COLOR_BGR2RGB)
    im = np.float32(im) / 255

    tot_pixels = im.shape[0] * im.shape[1]
    # compute number of gray pixels
    n = 0.1  # 0.01%
    num_gray_pixels = int(np.floor(n * tot_pixels / 100))

    # compute global illuminant values
    gt = np.array([0.6918, 0.6635, 0.2850])
    lumTriplet = GPconstancy_GI(im, num_gray_pixels, 10**(-4))

    # show results and angular error w.r.t gt
    print(lumTriplet)
    print(np.arccos(np.dot(lumTriplet, gt)) * 180 / np.pi)