import os
import cv2
import numpy as np
import argparse 
from skimage.filters import gaussian
from scipy.ndimage.interpolation import map_coordinates
from tqdm import tqdm
from PIL import Image


def single2uint(img):
    return np.uint8((img.clip(0, 1)*255.).round())


def uint2single(img):
    return np.float32(img/255.)


def Laplacian_edge_detector(img):
    # input: [0, 1]
    # return: [0, 1] (H, W, 3)
    img = np.clip(img*255, 0, 255).astype(np.uint8) # (H, W, 3)
    img = cv2.GaussianBlur(img, (3, 3), 0)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = cv2.Laplacian(img, cv2.CV_16S) # (H, W)
    img = cv2.convertScaleAbs(img)
    img = img.astype(np.float32) / 255.
    img = np.expand_dims(img, 2).repeat(3, axis=2) # (H, W, 3)
    return img


def Laplacian_edge_detector_uint8(img):
    # input: [0, 255]
    # return: [0, 255] (H, W, 3)
    img = cv2.GaussianBlur(img, (3, 3), 0)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = cv2.Laplacian(img, cv2.CV_16S) # (H, W)
    img = cv2.convertScaleAbs(img)
    img = np.expand_dims(img, 2).repeat(3, axis=2) # (H, W, 3)
    return img


def Canny_edge_detector(img):
    # input: [0, 1]
    # return: [0, 1] (H, W, 3)
    img = np.clip(img*255, 0, 255).astype(np.uint8) # (H, W, 3)
    img = cv2.GaussianBlur(img, (3, 3), 0)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = cv2.Canny(img, 50, 200) # (H, W)
    img = cv2.convertScaleAbs(img)
    img = img.astype(np.float32) / 255.
    img = np.expand_dims(img, 2).repeat(3, axis=2) # (H, W, 3)
    return img


def Canny_edge_detector_uint8(img):
    # input: [0, 255]
    # return: [0, 255] (H, W, 3)
    img = cv2.GaussianBlur(img, (3, 3), 0)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = cv2.Canny(img, 50, 200) # (H, W)
    img = cv2.convertScaleAbs(img)
    img = np.expand_dims(img, 2).repeat(3, axis=2) # (H, W, 3)
    return img


def Sobel_edge_detector(img):
    # input: [0, 1]
    # return: [0, 1] (H, W, 3)
    img = np.clip(img*255, 0, 255).astype(np.uint8) # (H, W, 3)
    img = cv2.GaussianBlur(img, (3, 3), 0)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = cv2.Sobel(img, cv2.CV_16S, 1, 1) # (H, W)
    img = cv2.convertScaleAbs(img)
    img = img.astype(np.float32) / 255.
    img = np.expand_dims(img, 2).repeat(3, axis=2) # (H, W, 3)
    return img


def erosion(img, kernel_size=5):
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    img = cv2.erode(img, kernel, iterations=1)
    return img


def dilatation(img, kernel_size=5):
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    img = cv2.dilate(img, kernel, iterations=1)
    return img


def opening(img):
    return dilatation(erosion(img))


def closing(img):
    return erosion(dilatation(img))


def morphological_gradient(img):
    return dilatation(img) - erosion(img)


def top_hat(img):
    return img - opening(img)


def black_hat(img):
    return closing(img) - img


def adjust_contrast(image, clip_limit=2.0, tile_grid_size=(8, 8)):
    
    image = single2uint(image)
    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)

    l, a, b = cv2.split(lab)

    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
    l_eq = clahe.apply(l)

    lab_eq = cv2.merge((l_eq, a, b))
    result = cv2.cvtColor(lab_eq, cv2.COLOR_LAB2BGR)

    result = uint2single(result)
    return result


def embossing(img):
    kernel = np.array([[0, -1, -1],
                       [1, 0, -1],
                       [1, 1, 0]])
    return cv2.filter2D(img, -1, kernel)


def hough_transform_line_detection(img):
    img = single2uint(img)
    dst = cv2.Canny(img, 50, 200, apertureSize=3)
    cdst = cv2.cvtColor(dst, cv2.COLOR_GRAY2BGR)
    lines = cv2.HoughLinesP(dst, 1, np.pi / 180, 230, None, 0, 0)
    if lines is not None:
        for i in range(0, len(lines)):
            rho = lines[i][0][0]
            theta = lines[i][0][1]
            a = np.cos(theta)
            b = np.sin(theta)

            x0 = a * rho
            y0 = b * rho
            pt1 = (int(x0 + 1000*(-b)), int(y0 + 1000*(a)))

            pt2 = (int(x0 - 1000*(-b)), int(y0 - 1000*(a)))
            cv2.line(img, pt1, pt2, (0, 0, 255), 3, cv2.LINE_AA)
    
    return uint2single(img)


def hough_circle_detection(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    circles = cv2.HoughCircles(gray, cv2.HOUGH_GRADIENT, 1, 100, param1=100, param2=30, minRadius=50, maxRadius=200)
    circles = np.uint16(np.around(circles))
    for i in circles[0, :]:
        cv2.circle(img, (i[0], i[1]), i[2], (0, 0, 255), 2)
    return img


def disk(radius, alias_blur=0.1, dtype=np.float32):
    if radius <= 8:
        L = np.arange(-8, 8 + 1)
        ksize = (3, 3)
    else:
        L = np.arange(-radius, radius + 1)
        ksize = (5, 5)
    X, Y = np.meshgrid(L, L)
    aliased_disk = np.array((X ** 2 + Y ** 2) <= radius ** 2, dtype=dtype)
    aliased_disk /= np.sum(aliased_disk)

    # supersample disk to antialias
    return cv2.GaussianBlur(aliased_disk, ksize=ksize, sigmaX=alias_blur)


def defocus_blur(image, level=(1, 0.1)):
    c = level
    kernel = disk(radius=c[0], alias_blur=c[1])

    channels = []
    for d in range(3):
        channels.append(cv2.filter2D(image[:, :, d], -1, kernel))
    channels = np.array(channels).transpose((1, 2, 0))  # 3x64x64 -> 64x64x3

    return np.clip(channels, 0, 1)


def masks_CFA_Bayer(shape):
    pattern = "RGGB"
    channels = dict((channel, np.zeros(shape)) for channel in "RGB")
    for channel, (y, x) in zip(pattern, [(0, 0), (0, 1), (1, 0), (1, 1)]):
        channels[channel][y::2, x::2] = 1
    return tuple(channels[c].astype(bool) for c in "RGB")


def cfa4_to_rgb(CFA4):
    RGB = np.zeros((CFA4.shape[0]*2, CFA4.shape[1]*2, 3), dtype=np.uint8)
    RGB[0::2, 0::2, 0] = CFA4[:, :, 0]  # R
    RGB[0::2, 1::2, 1] = CFA4[:, :, 1]  # G on R row
    RGB[1::2, 0::2, 1] = CFA4[:, :, 2]  # G on B row
    RGB[1::2, 1::2, 2] = CFA4[:, :, 3]  # B

    return RGB


def mosaic_CFA_Bayer(RGB):
    RGB = single2uint(RGB)
    R_m, G_m, B_m = masks_CFA_Bayer(RGB.shape[0:2])
    mask = np.concatenate(
        (R_m[..., np.newaxis], G_m[..., np.newaxis], B_m[..., np.newaxis]), axis=-1
    )
    mosaic = np.multiply(mask, RGB)  # mask*RGB
    CFA = mosaic.sum(2).astype(np.uint8)

    CFA4 = np.zeros((RGB.shape[0] // 2, RGB.shape[1] // 2, 4), dtype=np.uint8)
    CFA4[:, :, 0] = CFA[0::2, 0::2]
    CFA4[:, :, 1] = CFA[0::2, 1::2]
    CFA4[:, :, 2] = CFA[1::2, 0::2]
    CFA4[:, :, 3] = CFA[1::2, 1::2]

    rgb = cfa4_to_rgb(CFA4)
    rgb = uint2single(rgb)
    return rgb


def simulate_barrel_distortion(image, k1=0.02, k2=0.01):
    height, width = image.shape[:2]
    mapx, mapy = np.meshgrid(np.arange(width), np.arange(height))
    mapx = 2 * mapx / (width - 1) - 1
    mapy = 2 * mapy / (height - 1) - 1
    r = np.sqrt(mapx**2 + mapy**2)
    mapx = mapx * (1 + k1 * r**2 + k2 * r**4)
    mapy = mapy * (1 + k1 * r**2 + k2 * r**4)
    mapx = (mapx + 1) * (width - 1) / 2
    mapy = (mapy + 1) * (height - 1) / 2
    distorted_image = cv2.remap(image, mapx.astype(np.float32), mapy.astype(np.float32), cv2.INTER_LINEAR)
    return distorted_image


def simulate_pincushion_distortion(image, k1=-0.02, k2=-0.01):
    height, width = image.shape[:2]
    mapx, mapy = np.meshgrid(np.arange(width), np.arange(height))
    mapx = 2 * mapx / (width - 1) - 1
    mapy = 2 * mapy / (height - 1) - 1
    r = np.sqrt(mapx**2 + mapy**2)
    mapx = mapx * (1 + k1 * r**2 + k2 * r**4)
    mapy = mapy * (1 + k1 * r**2 + k2 * r**4)
    mapx = (mapx + 1) * (width - 1) / 2
    mapy = (mapy + 1) * (height - 1) / 2
    distorted_image = cv2.remap(image, mapx.astype(np.float32), mapy.astype(np.float32), cv2.INTER_LINEAR)
    return distorted_image


def rgb2gray(rgb):
    return np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140])


def spatter(x, severity=1):
    c = [(0.65, 0.3, 4, 0.69, 0.6, 0),
         (0.65, 0.3, 3, 0.68, 0.6, 0),
         (0.65, 0.3, 2, 0.68, 0.5, 0),
         (0.65, 0.3, 1, 0.65, 1.5, 1),
         (0.67, 0.4, 1, 0.65, 1.5, 1)][severity - 1]
    x_PIL = x
    x = np.array(x, dtype=np.float32) / 255.

    liquid_layer = np.random.normal(size=x.shape[:2], loc=c[0], scale=c[1])

    liquid_layer = gaussian(liquid_layer, sigma=c[2])
    liquid_layer[liquid_layer < c[3]] = 0
    if c[5] == 0:
        liquid_layer = (liquid_layer * 255).astype(np.uint8)
        dist = 255 - cv2.Canny(liquid_layer, 50, 150)
        dist = cv2.distanceTransform(dist, cv2.DIST_L2, 5)
        _, dist = cv2.threshold(dist, 20, 20, cv2.THRESH_TRUNC)
        dist = cv2.blur(dist, (3, 3)).astype(np.uint8)
        dist = cv2.equalizeHist(dist)
        ker = np.array([[-2, -1, 0], [-1, 1, 1], [0, 1, 2]])
        dist = cv2.filter2D(dist, cv2.CV_8U, ker)
        dist = cv2.blur(dist, (3, 3)).astype(np.float32)

        m = cv2.cvtColor(liquid_layer * dist, cv2.COLOR_GRAY2BGRA)
        m /= np.max(m, axis=(0, 1))
        m *= c[4]
        # water is pale turqouise
        color = np.concatenate((175 / 255. * np.ones_like(m[..., :1]),
                                238 / 255. * np.ones_like(m[..., :1]),
                                238 / 255. * np.ones_like(m[..., :1])), axis=2)

        color = cv2.cvtColor(color, cv2.COLOR_BGR2BGRA)

        if len(x.shape) < 3 or x.shape[2] < 3:
            add_spatter_color = cv2.cvtColor(np.clip(m * color, 0, 1),
                                             cv2.COLOR_BGRA2BGR)
            add_spatter_gray = rgb2gray(add_spatter_color)

            return (np.clip(x + add_spatter_gray, 0, 1) * 255).astype(np.uint8)

        else:

            x = cv2.cvtColor(x, cv2.COLOR_BGR2BGRA)

            return (cv2.cvtColor(np.clip(x + m * color, 0, 1),
                                cv2.COLOR_BGRA2BGR) * 255).astype(np.uint8)
    else:
        m = np.where(liquid_layer > c[3], 1, 0)
        m = gaussian(m.astype(np.float32), sigma=c[4])
        m[m < 0.8] = 0

        x_rgb = np.array(x_PIL)

        # mud brown
        color = np.concatenate((63 / 255. * np.ones_like(x_rgb[..., :1]),
                                42 / 255. * np.ones_like(x_rgb[..., :1]),
                                20 / 255. * np.ones_like(x_rgb[..., :1])),
                               axis=2)
        color *= m[..., np.newaxis]
        if len(x.shape) < 3 or x.shape[2] < 3:
            x *= (1 - m)
            return (np.clip(x + rgb2gray(color), 0, 1) * 255).astype(np.uint8)

        else:
            x *= (1 - m[..., np.newaxis])
            return (np.clip(x + color, 0, 1) * 255).astype(np.uint8)


# mod of https://gist.github.com/erniejunior/601cdf56d2b424757de5
def elastic_transform(image, severity=3):
    image = np.array(image, dtype=np.float32) / 255.
    shape = image.shape
    shape_size = shape[:2]

    sigma = np.array(shape_size) * 0.01
    alpha = [250 * 0.05, 250 * 0.065, 250 * 0.085, 250 * 0.1, 250 * 0.12][
        severity - 1]
    max_dx = shape[0] * 0.005
    max_dy = shape[0] * 0.005

    dx = (gaussian(np.random.uniform(-max_dx, max_dx, size=shape[:2]),
                   sigma, mode='reflect', truncate=3) * alpha).astype(
        np.float32)
    dy = (gaussian(np.random.uniform(-max_dy, max_dy, size=shape[:2]),
                   sigma, mode='reflect', truncate=3) * alpha).astype(
        np.float32)

    if len(image.shape) < 3 or image.shape[2] < 3:
        x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]))
        indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx, (-1, 1))
    else:
        dx, dy = dx[..., np.newaxis], dy[..., np.newaxis]
        x, y, z = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]),
                              np.arange(shape[2]))
        indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx,
                                                          (-1, 1)), np.reshape(
            z, (-1, 1))
    return np.clip(
        map_coordinates(image, indices, order=1, mode='reflect').reshape(
            shape), 0, 1) * 255


def frost(x, severity=2):
    c = [(1, 0.4),
         (0.8, 0.6),
         (0.7, 0.7),
         (0.65, 0.7),
         (0.6, 0.75)][severity - 1]

    idx = np.random.randint(5)
    filename = [os.path.join("degradation_toolkit/frost", 'frost1.png'),
                os.path.join("degradation_toolkit/frost", 'frost2.png'),
                os.path.join("degradation_toolkit/frost", 'frost3.png'),
                os.path.join("degradation_toolkit/frost", 'frost4.jpg'),
                os.path.join("degradation_toolkit/frost", 'frost5.jpg'),
                os.path.join("degradation_toolkit/frost", 'frost6.jpg')][idx]
    frost = Image.open(filename)
    frost = frost.convert("RGB")
    frost = np.array(frost)
    # frost = cv2.imread(filename)
    frost = uint2single(frost)
    frost_shape = frost.shape
    x_shape = np.array(x).shape

    # resize the frost image so it fits to the image dimensions
    scaling_factor = 1
    if frost_shape[0] >= x_shape[0] and frost_shape[1] >= x_shape[1]:
        scaling_factor = 1
    elif frost_shape[0] < x_shape[0] and frost_shape[1] >= x_shape[1]:
        scaling_factor = x_shape[0] / frost_shape[0]
    elif frost_shape[0] >= x_shape[0] and frost_shape[1] < x_shape[1]:
        scaling_factor = x_shape[1] / frost_shape[1]
    elif frost_shape[0] < x_shape[0] and frost_shape[1] < x_shape[
        1]:  # If both dims are too small, pick the bigger scaling factor
        scaling_factor_0 = x_shape[0] / frost_shape[0]
        scaling_factor_1 = x_shape[1] / frost_shape[1]
        scaling_factor = np.maximum(scaling_factor_0, scaling_factor_1)

    scaling_factor *= 1.1
    new_shape = (int(np.ceil(frost_shape[1] * scaling_factor)),
                 int(np.ceil(frost_shape[0] * scaling_factor)))
    frost_rescaled = cv2.resize(frost, dsize=new_shape,
                                interpolation=cv2.INTER_CUBIC)

    # randomly crop
    x_start, y_start = np.random.randint(0, frost_rescaled.shape[0] - x_shape[
        0]), np.random.randint(0, frost_rescaled.shape[1] - x_shape[1])

    if len(x_shape) < 3 or x_shape[2] < 3:
        frost_rescaled = frost_rescaled[x_start:x_start + x_shape[0],
                         y_start:y_start + x_shape[1]]
        frost_rescaled = rgb2gray(frost_rescaled)
    else:
        frost_rescaled = frost_rescaled[x_start:x_start + x_shape[0],
                         y_start:y_start + x_shape[1]][..., [2, 1, 0]]
    return c[0] * np.array(x) + c[1] * frost_rescaled