import cv2
import numpy as np

from fractions import Fraction
from exifread.utils import Ratio
from PIL import Image
from skimage.color import rgb2hsv, hsv2rgb
from skimage.exposure import rescale_intensity
from skimage.filters import gaussian as sk_gaussian
from skimage.restoration import denoise_tv_bregman
from scipy import signal

from colour_demosaicing import demosaicing_CFA_Bayer_Menon2007

from utils.misc import *
from utils.color import *
from ips.wb import illumination_parameters_estimation


def normalize(raw_image, black_level, white_level):
    if isinstance(black_level, list) and len(black_level) == 1:
        black_level = float(black_level[0])
    if isinstance(white_level, list) and len(white_level) == 1:
        white_level = float(white_level[0])
    black_level_mask = black_level
    if type(black_level) is list and len(black_level) == 4:
        if type(black_level[0]) is Ratio:
            black_level = ratios2floats(black_level)
        if type(black_level[0]) is Fraction:
            black_level = fractions2floats(black_level)
        black_level_mask = np.zeros(raw_image.shape)
        idx2by2 = [[0, 0], [0, 1], [1, 0], [1, 1]]
        step2 = 2
        for i, idx in enumerate(idx2by2):
            black_level_mask[idx[0]::step2, idx[1]::step2] = black_level[i]
    normalized_image = raw_image.astype(np.float32) - black_level_mask
    # if some values were smaller than black level
    normalized_image[normalized_image < 0] = 0
    normalized_image = normalized_image / (white_level - black_level_mask)
    return normalized_image


def demosaic(norm_image, cfa_pattern):
    return demosaicing_CFA_Bayer_Menon2007(norm_image, decode_cfa_pattern(cfa_pattern))


def denoise(demosaiced_image, y_noise_profile, cc_noise_profile):
    ycc_demosaiced = rgb2ycc(demosaiced_image[:, :, ::-1])
    y_demosaiced = ycc_demosaiced[:, :, 0]
    cc_demosaiced = ycc_demosaiced[:, :, 1:]
    current_image_y = y_demosaiced
    current_image_cc = gaussian(cc_demosaiced, sigma=cc_noise_profile)
    current_image_ycc = np.concatenate([
        np.expand_dims(current_image_y, -1), 
        current_image_cc
    ], axis=-1)
    return ycc2rgb(current_image_ycc)[:, :, ::-1]


def raw_color_denoise(demosaiced_image, cc_noise_profile):
    ycc_demosaiced = rgb2ycc(demosaiced_image[:, :, ::-1])
    cc_demosaiced = ycc_demosaiced[:, :, 1:]
    cc_demosaiced_denoised = sk_gaussian(cc_demosaiced, sigma=cc_noise_profile)
    ycc_demosaiced[:, :, 1:] = cc_demosaiced_denoised
    return ycc2rgb(ycc_demosaiced)[:, :, ::-1]


def luminance_denoise(tone_mapped_image, weight=20.0):
    ycc_tone_mapped = rgb2ycc(tone_mapped_image[:, :, ::-1])
    y_tone_mapped = ycc_tone_mapped[:, :, 0]
    y_tone_mapped_denoised = denoise_tv_bregman(y_tone_mapped, weight=weight)
    ycc_tone_mapped[:, :, 0] = np.clip(y_tone_mapped_denoised, 1e-4, 0.999)
    return ycc2rgb(ycc_tone_mapped)[:, :, ::-1]


def white_balance(denoised_image, metadata, max_repeat_limit=10000):
    if metadata["wb_estimation"] is not None:
        as_shot_neutral = np.array(metadata["wb_estimation"])
        white_balanced_image = np.dot(denoised_image, as_shot_neutral.T)
        return np.clip(white_balanced_image, 0.0, 1.0)
    illumuniation_estimation_algorithm = metadata["wb_method"]
    as_shot_neutral = illumination_parameters_estimation(denoised_image, illumuniation_estimation_algorithm)    
    
    if isinstance(as_shot_neutral[0], Ratio):
        as_shot_neutral = ratios2floats(as_shot_neutral)

    as_shot_neutral = np.asarray(as_shot_neutral)
    # transform vector into matrix
    if as_shot_neutral.shape == (3,):
        as_shot_neutral = np.diag(1./as_shot_neutral)

    assert as_shot_neutral.shape == (3, 3)
    repeat_count = 0
    while (as_shot_neutral[0, 0] < 2.3 and as_shot_neutral[2, 2] < 2.3) or (as_shot_neutral[0, 0] < 2.02 or as_shot_neutral[2, 2] < 1.92):
        if repeat_count < max_repeat_limit:
            as_shot_neutral = illumination_parameters_estimation(denoised_image, illumuniation_estimation_algorithm)
            if isinstance(as_shot_neutral[0], Ratio):
                as_shot_neutral = ratios2floats(as_shot_neutral)

            as_shot_neutral = np.asarray(as_shot_neutral)
            # transform vector into matrix
            if as_shot_neutral.shape == (3,):
                as_shot_neutral = np.diag(1./as_shot_neutral)

            assert as_shot_neutral.shape == (3, 3)
        else:
            print(f"WARNING! Invalid range for illumination matrix and repeated to estimate by '{illumuniation_estimation_algorithm}' so many times. Using 'gw' for illumination estimation now...")
            as_shot_neutral = illumination_parameters_estimation(denoised_image, "gw")  
            if isinstance(as_shot_neutral[0], Ratio):
                as_shot_neutral = ratios2floats(as_shot_neutral)

            as_shot_neutral = np.asarray(as_shot_neutral)
            # transform vector into matrix
            if as_shot_neutral.shape == (3,):
                as_shot_neutral = np.diag(1./as_shot_neutral)

            assert as_shot_neutral.shape == (3, 3)
            break
        repeat_count += 1
    
    white_balanced_image = np.dot(denoised_image, as_shot_neutral.T)
    metadata["wb_estimation"] = as_shot_neutral.tolist()

    # print(as_shot_neutral)
    return np.clip(white_balanced_image, 0.0, 1.0)


def xyz_transform(wb_image, color_matrix):
    if isinstance(color_matrix[0], Fraction):
        color_matrix = fractions2floats(color_matrix)
    xyz2cam = np.reshape(np.asarray(color_matrix), (3, 3))
    # normalize rows (needed?)
    xyz2cam = xyz2cam / np.sum(xyz2cam, axis=1, keepdims=True)
    # inverse
    cam2xyz = np.linalg.inv(xyz2cam)
    # for now, use one matrix  # TODO: interpolate btween both
    # simplified matrix multiplication
    xyz_image = cam2xyz[np.newaxis, np.newaxis, :, :] * wb_image[:, :, np.newaxis, :]
    xyz_image = np.sum(xyz_image, axis=-1)
    xyz_image = np.clip(xyz_image, 0.0, 1.0)
    return xyz_image


def xyz_to_srgb(xyz_image):
    # srgb2xyz = np.array([[0.4124564, 0.3575761, 0.1804375],
    #                      [0.2126729, 0.7151522, 0.0721750],
    #                      [0.0193339, 0.1191920, 0.9503041]])

    # xyz2srgb = np.linalg.inv(srgb2xyz)

    xyz2srgb = np.array([[3.2404542, -1.5371385, -0.4985314],
                         [-0.9692660, 1.8760108, 0.0415560],
                         [0.0556434, -0.2040259, 1.0572252]])

    # normalize rows (needed?)
    xyz2srgb = xyz2srgb / np.sum(xyz2srgb, axis=-1, keepdims=True)

    srgb_image = xyz2srgb[np.newaxis, np.newaxis,
                          :, :] * xyz_image[:, :, np.newaxis, :]
    srgb_image = np.sum(srgb_image, axis=-1)
    srgb_image = np.clip(srgb_image, 0.0, 1.0)
    return srgb_image


def apply_tmo_flash(Y, a):
    Y[Y == 0] = 1e-9
    return Y / (Y + a * np.exp(np.mean(np.log(Y))))


def apply_tmo_storm(Y, a, kernels):
    rows, cols = Y.shape
    Y[Y == 0] = 1e-9
    return sum([
        Y / (Y + a * np.exp(cv2.boxFilter(np.log(Y), -1, (int(min(rows // kernel, cols // kernel)),) * 2)))
        for kernel in kernels
    ]) / len(kernels)


def apply_tmo_nite(Y, CC, kernels):
    rows, cols = Y.shape
    Y[Y == 0] = 1e-9
    y_mu, y_std = max(Y.mean(), 0.001), Y.std()
    cc_std = CC.std()
    # tmo_offset = np.exp(y_mu * (cc_std / y_std) * 100)
    tmo_offset = 10. / np.sqrt(np.exp(np.log(y_mu) * (np.log(cc_std) / np.log(y_std))) * 100)
    # print(f"Y mean: {y_mu:.3f}, Y std: {y_std:.3f}, CC std: {cc_std:.3f}, Offset: {tmo_offset:.3f}")

    # tmo_scale = 8.5 + min(6.5, round(tmo_offset))
    tmo_scale = min(28., max(5., tmo_offset))
    # print(f"TMO scale: {tmo_scale}")
    return sum([
        Y / np.clip((Y + tmo_scale * np.exp(cv2.boxFilter(np.log(Y), -1, (int(min(rows // kernel, cols // kernel)),) * 2))), 0., 1.)
        for kernel in kernels
    ]) / len(kernels)


def perform_tone_mapping(source, metadata):
    ycc_source = rgb2ycc(source[:, :, ::-1])
    y_source = ycc_source[:, :, 0]
    cc_source = ycc_source[:, :, 1:]

    if metadata["tmo_type"].lower() == "flash":
        y_hat_source = apply_tmo_flash(y_source, metadata["tmo_scale"])
    elif metadata["tmo_type"].lower() == "storm":
        y_hat_source = apply_tmo_storm(y_source, metadata["tmo_scale"], metadata["tmo_kernels"])
    else:  # nite
        y_hat_source = apply_tmo_nite(y_source, cc_source, metadata["tmo_kernels"])

    ycc_nite = np.concatenate([
        np.expand_dims(y_hat_source, -1), 
        cc_source
    ], axis=-1)
    result = ycc2rgb(ycc_nite)[:, :, ::-1]
    if metadata["tmo_do_leap"]:
        target_mean_grayscale = 0.282  # 72 / 255
        result = np.clip(result, a_min=0., a_max=1.)
        grayscale = cv2.cvtColor(result * 255., cv2.COLOR_BGR2GRAY) / 255.
        result *= target_mean_grayscale / np.mean(grayscale)
    result = np.clip(result, a_min=0., a_max=1.)
    return result


def global_mean_contrast(input_im, beta=1.0):
    mu_ = input_im.mean(axis=(0, 1), keepdims=True)
    output_im = mu_ + beta * (input_im - mu_)
    output_im = np.where(0 > output_im, input_im, output_im)
    output_im = np.where(1 < output_im, input_im, output_im)
    return output_im


def s_curve_correction(input_im, alpha=0.5, lambd=0.5):
    ycc_ = rgb2ycc(input_im[:, :, ::-1])
    Y = ycc_[:, :, 0]
    Y_hat = alpha + np.where(
        Y >= alpha, 
        (1 - alpha) * np.power(((Y - alpha) / (1 - alpha)), lambd), 
        -alpha * np.power((1 - (Y / alpha)), lambd)
    )
    ycc_[:, :, 0] = Y_hat
    bgr_ = np.clip(ycc2rgb(ycc_)[:, :, ::-1], a_min=0., a_max=1.)
    return bgr_


def histogram_stretching(input_im):
    hsv = rgb2hsv(input_im[:, :, ::-1])
    V = hsv[:, :, 0]
    p0_01, p99 = np.percentile(V, (0.01, 99.99))
    if 0.7 > p99:
         _, p99 = np.percentile(V, (0.01, 99.5))

    V_hat = rescale_intensity(V, in_range=(p0_01, p99))
    hsv[:, :, 0] = V_hat
    bgr_ = np.clip(hsv2rgb(hsv), a_min=0., a_max=1.)[:, :, ::-1]
    return bgr_


def conditional_contrast_correction(input_im, threshold=0.5):
    ycc_ = rgb2ycc(input_im[:, :, ::-1])
    Y = ycc_[:, :, 0]
    y_avg = Y.mean()
    if y_avg > threshold:
        Y_hat = Y.copy()
        idx = Y_hat <= 0.0031308
        Y_hat[idx] *= 12.92
        Y_hat[idx == False] = (Y_hat[idx == False] ** (1.0 / 2.4)) * 1.055 - 0.055
    else:
        alpha = 0.5
        lambd = 1.2
        Y_hat = alpha + np.where(
            Y >= alpha, 
            (1 - alpha) * np.power(((Y - alpha) / (1 - alpha)), lambd), 
            -alpha * np.power((1 - (Y / alpha)), lambd)
        )
    ycc_[:, :, 0] = Y_hat
    bgr_ = np.clip(ycc2rgb(ycc_)[:, :, ::-1], a_min=0., a_max=1.)
    return bgr_


def memory_color_enhancement(data, color_space="srgb", illuminant="D65", clip_range=[0, 1], cie_version="1964"):
    target_hue = [30., -125., 100.]
    hue_preference = [20., -118., 130.]
    hue_sigma = [20., 10., 5.]
    is_both_side = [True, False, False]
    multiplier = [0.6, 0.6, 0.6]
    chroma_preference = [25., 14., 30.]
    chroma_sigma = [10., 10., 5.]

    # RGB to xyz
    data = rgb2xyz(data, color_space, clip_range)
    # xyz to lab
    data = xyz2lab(data, cie_version, illuminant)
    # lab to lch
    data = lab2lch(data)

    # hue squeezing
    # we are traversing through different color preferences
    height, width, _ = data.shape
    hue_correction = np.zeros((height, width), dtype=np.float32)
    for i in range(0, np.size(target_hue)):

        delta_hue = data[:, :, 2] - hue_preference[i]

        if is_both_side[i]:
            weight_temp = np.exp(-np.power(data[:, :, 2] - target_hue[i], 2) / (2 * hue_sigma[i] ** 2)) + \
                          np.exp(-np.power(data[:, :, 2] + target_hue[i], 2) / (2 * hue_sigma[i] ** 2))
        else:
            weight_temp = np.exp(-np.power(data[:, :, 2] - target_hue[i], 2) / (2 * hue_sigma[i] ** 2))

        weight_hue = multiplier[i] * weight_temp / np.max(weight_temp)

        weight_chroma = np.exp(-np.power(data[:, :, 1] - chroma_preference[i], 2) / (2 * chroma_sigma[i] ** 2))

        hue_correction = hue_correction + np.multiply(np.multiply(delta_hue, weight_hue), weight_chroma)

    # correct the hue
    data[:, :, 2] = data[:, :, 2] - hue_correction

    # lch to lab
    data = lch2lab(data)
    # lab to xyz
    data = lab2xyz(data, cie_version, illuminant)
    # xyz to rgb
    data = xyz2rgb(data, color_space, clip_range)

    data = outOfGamutClipping(data, range=clip_range[1])
    return data


def unsharp_masking(data, gaussian_kernel_size=[5, 5], gaussian_sigma=2.0, slope=1.5, tau_threshold=0.05, gamma_speed=4., clip_range=[0, 1]):
    # create gaussian kernel
    gaussian_kernel = gaussian(gaussian_kernel_size, gaussian_sigma)

    # convolve the image with the gaussian kernel
    # first input is the image
    # second input is the kernel
    # output shape will be the same as the first input
    # boundary will be padded by using symmetrical method while convolving
    if np.ndim(data) > 2:
        image_blur = np.empty(np.shape(data), dtype=np.float32)
        for i in range(0, np.shape(data)[2]):
            image_blur[:, :, i] = signal.convolve2d(data[:, :, i], gaussian_kernel, mode="same", boundary="symm")
    else:
        image_blur = signal.convolve2d(data, gaussian_kernel, mode="same", boundary="symm")

    # the high frequency component image
    image_high_pass = data - image_blur

    # soft coring (see in utility)
    # basically pass the high pass image via a slightly nonlinear function
    tau_threshold = tau_threshold * clip_range[1]

    # add the soft cored high pass image to the original and clip
    # within range and return
    def soft_coring(img_hp, slope, tau_threshold, gamma_speed):
        return slope * np.float32(img_hp) * (1. - np.exp(-((np.abs(img_hp / tau_threshold))**gamma_speed)))
    return np.clip(data + soft_coring(image_high_pass, slope, tau_threshold, gamma_speed), clip_range[0], clip_range[1])


def to_uint8(srgb):
    return (srgb * 255).astype(np.uint8)


def resize(img, width=None, height=None):
    if width is None or height is None:
        return img
    img_pil = Image.fromarray(img)
    out_size = (width, height)
    if img_pil.size == out_size:
        return img
    out_img = img_pil.resize(out_size, Image.Resampling.LANCZOS)
    out_img = np.array(out_img)
    return out_img


def fix_orientation(image, orientation):
    # 1 = Horizontal (normal)
    # 2 = Mirror horizontal
    # 3 = Rotate 180
    # 4 = Mirror vertical
    # 5 = Mirror horizontal and rotate 270 CW
    # 6 = Rotate 90 CW
    # 7 = Mirror horizontal and rotate 90 CW
    # 8 = Rotate 270 CW

    orientation_dict = {
        "Horizontal (normal)": 1,
        "Mirror horizontal": 2,
        "Rotate 180": 3,
        "Mirror vertical": 4,
        "Mirror horizontal and rotate 270 CW": 5,
        "Rotate 90 CW": 6,
        "Mirror horizontal and rotate 90 CW": 7,
        "Rotate 270 CW": 8
    }

    if type(orientation) is list:
        orientation = orientation[0]
    orientation = orientation_dict[orientation]
    if orientation == 1:
        pass
    elif orientation == 2:
        image = cv2.flip(image, 0)
    elif orientation == 3:
        image = cv2.rotate(image, cv2.ROTATE_180)
    elif orientation == 4:
        image = cv2.flip(image, 1)
    elif orientation == 5:
        image = cv2.flip(image, 0)
        image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)
    elif orientation == 6:
        image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
    elif orientation == 7:
        image = cv2.flip(image, 0)
        image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
    elif orientation == 8:
        image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)

    return image