"""
Main module - set of functions that simulate or correct colorblindness for images
as described in
https://www.researchgate.net/publication/326626897_Smartphone_Based_Image_Color_Correction_for_Color_Blindness
"""
import numpy as np
import cv2

## LMS Daltonization
def rgb_to_lms(img):
    lms_matrix = np.array(
        [[17.8824, 43.5161, 4.11935],
        [3.45565, 27.1554, 3.86714],
        [0.0299566, 0.184309, 1.46709]
        ]
        )
    return np.tensordot(img, lms_matrix, axes=[[2], [1]])

def lms_to_rgb(img):
    rgb_matrix = np.array(
        [[0.0809444479, -0.130504409, 0.116721066],
        [0.113614708, -0.0102485335, 0.0540193266],
        [-0.000365296938, -0.00412161469, 0.693511405]
        ]
        )
    return np.tensordot(img, rgb_matrix, axes=[[2], [1]])

def simulate_colorblindness(img, colorblind_type):
    lms_img = rgb_to_lms(img)
    if colorblind_type.lower() == 'protanopia':
        sim_matrix = np.array(
            [[0.0, 2.02344, -2.52582],
            [0.0, 1.0, 0.0],
            [0.0, 0.0, 1.0]
            ]
            )
    elif colorblind_type.lower() == 'duteranopia':
        sim_matrix = np.array(
            [[1.0, 0.0, 0.0],
            [0.49421, 0.0, 1.24827],
            [0.0, 0.0, 1.0]
            ]
            )
    elif colorblind_type.lower() == 'tritanopia':
        sim_matrix = np.array(
            [[1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [-0.395913, 0.801109, 0.0]
            ]
            )
    else:
        raise ValueError('{} is an unrecognized colorblindness type.'.format(colorblind_type))
    lms_img = np.tensordot(lms_img, sim_matrix, axes=[[2], [1]])
    rgb_img = lms_to_rgb(lms_img)
    return rgb_img

def daltonize_correct(img, colorblind_type):
    colorblind_img = simulate_colorblindness(img, colorblind_type=colorblind_type)
    error_matrix = img - colorblind_img
    if colorblind_type.lower() == 'protanopia':
        correction_matrix = np.array(
            [[0.0, 0.0, 0.0],
            [0.7, 1.0, 0.0],
            [0.7, 0.0, 1.0]
            ]
            )
    elif colorblind_type.lower() == 'duteranopia':
        correction_matrix = np.array(
            [[1.0, 0.7, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.7, 1.0]
            ]
            )
    elif colorblind_type.lower() == 'tritanopia':
        correction_matrix = np.array(
            [[1.0, 0.0, 0.7],
            [0.0, 1.0, 0.7],
            [0.0, 0.0, 0.0]
            ]
            )
    else:
        raise ValueError('{} is an unrecognized colorblindness type.'.format(colorblind_type))
    corrected_error_matrix = np.tensordot(error_matrix, correction_matrix, axes=[[2], [1]])
    return img + corrected_error_matrix

## color-blind filter service (CBFS) algorithm


## LAB color correction
def lab_correct(img, l_shift=15, a_shift=15, b_shift=15):
    lab_img = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
    for i, shift in enumerate(l_shift, a_shift, b_shift):
        lab_img[..., i] = np.where(
            lab_img[..., i] > 127,
            lab_img[..., i] + shift,
            lab_img[..., i] - shift
            )
    lab_img = np.clip(lab_img, a_min=0.0, a_max=255)
    rgb_img = cv2.cvtColor(img, cv2.COLOR_LAB2RGB)
    return rgb_img

## HSV Color Shifting Algorithm

# hsv and rgb conversions from
# https://stackoverflow.com/questions/27041559/rgb-to-hsv-python-change-hue-continuously

def rgb_to_hsv(rgb):
    # Translated from source of colorsys.rgb_to_hsv
    # r,g,b should be a numpy arrays with values between 0 and 255
    # rgb_to_hsv returns an array of floats between 0.0 and 1.0.
    rgb = rgb.astype('float')
    hsv = np.zeros_like(rgb)
    # in case an RGBA array was passed, just copy the A channel
    hsv[..., 3:] = rgb[..., 3:]
    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
    maxc = np.max(rgb[..., :3], axis=-1)
    minc = np.min(rgb[..., :3], axis=-1)
    hsv[..., 2] = maxc
    mask = maxc != minc
    hsv[mask, 1] = (maxc - minc)[mask] / maxc[mask]
    rc = np.zeros_like(r)
    gc = np.zeros_like(g)
    bc = np.zeros_like(b)
    rc[mask] = (maxc - r)[mask] / (maxc - minc)[mask]
    gc[mask] = (maxc - g)[mask] / (maxc - minc)[mask]
    bc[mask] = (maxc - b)[mask] / (maxc - minc)[mask]
    hsv[..., 0] = np.select(
        [r == maxc, g == maxc], [bc - gc, 2.0 + rc - bc], default=4.0 + gc - rc)
    hsv[..., 0] = (hsv[..., 0] / 6.0) % 1.0
    return hsv

def hsv_to_rgb(hsv):
    # Translated from source of colorsys.hsv_to_rgb
    # h,s should be a numpy arrays with values between 0.0 and 1.0
    # v should be a numpy array with values between 0.0 and 255.0
    # hsv_to_rgb returns an array of uints between 0 and 255.
    rgb = np.empty_like(hsv)
    rgb[..., 3:] = hsv[..., 3:]
    h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
    i = (h * 6.0).astype('uint8')
    f = (h * 6.0) - i
    p = v * (1.0 - s)
    q = v * (1.0 - s * f)
    t = v * (1.0 - s * (1.0 - f))
    i = i % 6
    conditions = [s == 0.0, i == 1, i == 2, i == 3, i == 4, i == 5]
    rgb[..., 0] = np.select(conditions, [v, q, p, p, t, v], default=v)
    rgb[..., 1] = np.select(conditions, [v, v, v, q, p, p], default=t)
    rgb[..., 2] = np.select(conditions, [v, p, t, v, v, q], default=p)
    return rgb.astype('uint8')

def color_shift_correct(img, h=0.3):
    hsv_img = rgb_to_hsv(img)
    hsv_img[..., 0] = hsv_img[..., 0] + h
    hsv_img[..., 0] = np.where(
        hsv_img[..., 0] > 1.0,
        hsv_img[..., 0] - 1.0,
        hsv_img[..., 0]
        )
    hsv_img[..., 2] = hsv_img[..., 2] * 255
    rgb_img = hsv_to_rgb(hsv_img)
    return rgb_img
