#!/usr/bin/env python
"""
Library handling projection operations of a standard camera model.
"""
from __future__ import division, print_function
import numpy as np
import cv2


def stretch_image_contrast(img, clim, tile_grid_size=(8, 8)):
    if clim == 0:
        return img

    img = img.astype(np.float32)
    img[img < 0] = 0
    img /= np.percentile(img.ravel(), 99.9)
    img[img > 1] = 1
    img = np.round(img*255).astype(np.uint8)

    clahe = cv2.createCLAHE(clipLimit=clim, tileGridSize=tile_grid_size)

    if img.ndim == 3:
        hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
        hsv[:, :, 2] = clahe.apply(hsv[:, :, 2])

        return cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
    else:
        return clahe.apply(img)


def srgb_gamma_encode(image0):
    ind = image0 > 0.0031308
    image = np.zeros_like(image0)
    image[ind] = 1.055*image0[ind]**(1/2.4) - 0.055
    image[~ind] = 12.92*image0[~ind]
    return image


def srgb_gamma_decode(image):
    image = to_float(image)

    ind = image > 0.04045
    image[ind] = ((image[ind] + 0.055)/1.055)**2.4
    image[~ind] = image[~ind]/12.92

    return image


def to_float(image):
    if image.dtype in (np.float16, np.float32, np.float64, float):
        return image
    if image.dtype == np.uint8:
        return uint8_to_float(image)
    elif image.dtype == np.uint16:
        return uint16_to_float(image)


def uint8_to_float(image):
    # We break up the 0-1 float range into 256 bins and take the bin-center
    # value.
    return image.astype(float)/256 + 1/512


def uint16_to_float(image):
    # We break up the 0-1 float range into 65536 bins and take the bin-center
    # value.
    return image.astype(float)/65536 + 1/131072


def float_to_uint8(image):
    image = np.round(image*256 - 0.5)
    image = np.clip(image, 0, 255)

    return image.astype(np.uint8)


def float_to_uint16(image):
    image = np.round(image*65536 - 0.5)
    image = np.clip(image, 0, 65535)

    return image.astype(np.uint16)



lut_8_to_16 = float_to_uint16(to_float(np.arange(256, dtype=np.uint8)))

decode_lut_8bit_ = srgb_gamma_decode(np.arange(256, dtype=np.uint8))
decode_lut_16bit_ = srgb_gamma_decode(np.arange(65536, dtype=np.uint16))
def srgb_gamma_decode_lut(image):
    """Use lookup table to decode gamme-encoded 8-bit sRGB to linear sRGB.
    """
    if image.dtype == np.uint8:
        return np.take(decode_lut_8bit_, image)
    elif image.dtype == np.uint16:
        return np.take(decode_lut_16bit_, image)
    else:
        raise Exception()


v_in = np.arange(65536, dtype=np.uint16)
v_out0 = srgb_gamma_encode(uint16_to_float(v_in))
encode_lut_8bit_ = float_to_uint8(v_out0)
encode_lut_16bit_ = float_to_uint16(v_out0)
def srgb_gamma_encode_lut(image, outbits=8):
    """Use lookup table to decode gamme-encoded 8-bit sRGB to linear sRGB.
    """
    # First convert to 16 bit
    if image.dtype in (np.float16, np.float32, np.float64, float):
        image = float_to_uint16(image)
    elif image.dtype == np.uint16:
        pass
    elif image.dtype == np.uint8:
        image = np.take(lut_8_to_16, image)
    else:
        raise Exception()

    if outbits == 16:
        return np.take(encode_lut_16bit_, image)
    elif outbits == 8:
        return np.take(encode_lut_8bit_, image)
    else:
        raise Exception('\'outbits\' must be 8 or 16')


def calibrate_device_rgb(image, gain, color_cal_mat, stretch=99.9):
    """Correct a raw device RGB image to calibrated sRGB.

    :param image: Device RGB image to process.
    :type image: Numpy image

    :param gain: Gain image to multiply the raw image by.
    :type gain: Numpy image

    :param color_cal_mat: Color calibration matrix taking linear device RGB and
        returning linear sRGB.
    :type color_cal_mat: 3x3

    :param stretch: Percentile of the highest value in 'image' to set to the
        saturation value during linear stretching.
    :type stretch: float

    """
    #image = np.ones((100,100, 3), dtype=np.uint8)*255
    if image.dtype == np.uint8:
        satval = uint8_to_float(np.array([255]))
    elif image.dtype == np.uint16:
        satval = uint16_to_float(np.array([65535]))
    else:
        satval = 1

    h, w = image.shape[0], image.shape[1]

    image = srgb_gamma_decode_lut(image)
    if gain is not None:
        for i in range(3):
            image[:, :, i] = image[:, :, i]*gain


    image = np.reshape(image, (h*w, 3)).T

    if False:
        # Need to clamp at output white. Otherwise, as one of the device RGB
        # channels saturates, it could change the output calibrated color.
        maxv = np.dot(np.linalg.inv(color_cal_mat), [1, 1, 1])
        maxv /= maxv.max()

        #i = 1; plt.imshow(np.reshape(image[0].T, (h, w)) > maxv[i])

        for i in range(3):
            if maxv[i] < 1:
                image[i] = np.clip(image[i], 0, maxv[i])

    image = np.dot(color_cal_mat, image)

    v = np.array([np.percentile(image[i], stretch) for i in range(3)])

    # TODO add Von Kries's method for proper white balancing.
    # Auto white balance.
    ind = image[0] > 0
    wbm = 1/np.array([1, np.median(image[1][ind]/image[0][ind]),
                    np.median(image[2][ind]/image[0][ind])])

    wbm /= wbm.max()

    s = min(1/wbm/v)
    wbm *= s

    for i in range(3):
        image[i] = np.clip(image[i]*wbm[i], 0, 1)

    # Now no pixel will be mapped to a color "beyond" white.

    image = srgb_gamma_encode_lut(image)
    image = np.reshape(image.T, (h, w, 3))
    return image