import argparse
import tensorflow as tf
import sys
import os
import numpy as np
import cv2
from random import random
import scipy.misc
from glob import glob
import matplotlib.pyplot as plt
import math


def make_image_square(image):

    if image.shape[0] > image.shape[1]:
        shift = int((image.shape[0] - image.shape[1]) / 2)
        print(shift)
        small_dim = image.shape[1]
        image = image[shift : shift + small_dim, :, :]

    else:
        shift = int((image.shape[1] - image.shape[0]) / 2)
        print(shift)
        small_dim = image.shape[0]
        image = image[:, shift : shift + small_dim]

    return image


def mix_with_background(path_background, frame, fg):

    images = [img for img in os.listdir(path_background) if img.endswith(".jpg")]

    # Assure that there is a binary image
    # fg[np.where(fg > 150)] = 255
    # fg[np.where(fg < 150)] = 0

    # fg[np.where(fg == 1)] = 255
    # fg[np.where(fg != 1)] = 0

    upper_limit = 0
    lower_limit = len(images) - 1
    num = np.uint8(random() * (upper_limit - lower_limit) + lower_limit)
    # print(images[num])
    bkg = cv2.imread(os.path.join(path_background, images[num]))
    # bkg = cv2.cvtColor(bkg, cv2.COLOR_BGR2RGB)
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    [index_r, index_c] = np.where((fg[:, :] >= 200))
    # [index_r, index_c] = np.where((fg[:,:] == 1))

    for i in range(1, len(index_r)):

        bkg[index_r[i], index_c[i], 0] = frame[index_r[i], index_c[i], 0]
        bkg[index_r[i], index_c[i], 1] = frame[index_r[i], index_c[i], 1]
        bkg[index_r[i], index_c[i], 2] = frame[index_r[i], index_c[i], 2]

    return bkg


def equalize(image):

    output = np.zeros((image.shape[0], image.shape[1]))

    img_yuv = cv2.cvtColor(image, cv2.COLOR_BGR2YUV)
    # This configuration achieves a very slight equalization
    clahe = cv2.createCLAHE(clipLimit=1.0, tileGridSize=(1, 1))
    img_yuv[:, :, 0] = clahe.apply(img_yuv[:, :, 0])
    output = cv2.cvtColor(img_yuv, cv2.COLOR_YUV2BGR)

    return output


def show_image(image, text=None):
    plt.imshow(image)
    plt.show()
    if text is not None:
        plt.title(text)


def show_two_images(image, image2, text=None, cmap=None, horizontal=True):

    plt.figure(1)

    if horizontal:
        plt.subplot(121)
    else:
        plt.subplot(211)

    if cmap:
        plt.imshow(image, cmap=cmap)
    else:
        plt.imshow(cv2.cvtColor(image.astype("float32"), cv2.COLOR_BGR2RGB))

    if text is not None:
        plt.title(text)

    if horizontal:
        plt.subplot(122)
    else:
        plt.subplot(212)

    if cmap:
        plt.imshow(image2, cmap=cmap)
    else:
        plt.imshow(cv2.cvtColor(image2.astype("float32"), cv2.COLOR_BGR2RGB))
    plt.show()


def show_three_images(image, image2, image3, text=None):

    plt.figure(1)
    plt.subplot(311)
    plt.imshow(image)

    if text is not None:
        plt.title(text)

    plt.subplot(312)

    # plt.imshow(image2,cmap='gray', vmin=0, vmax=1)
    plt.imshow(image2)
    plt.subplot(313)

    # plt.imshow(image3,cmap='gray', vmin=0, vmax=1)
    plt.imshow(image3)
    plt.show()


def show_four_images(image, image2, image3, image4, text=None):

    plt.figure(1)
    plt.subplot(221)
    plt.imshow(image)

    if text is not None:
        plt.title(text)

    plt.subplot(222)
    plt.imshow(image2)
    # plt.imshow(image2, cmap='gray', vmin=0, vmax=1)

    plt.subplot(223)
    plt.imshow(image3)
    # plt.imshow(image3, cmap='gray', vmin=0, vmax=1)

    plt.subplot(224)
    plt.imshow(image4)
    # plt.imshow(image4, cmap='gray', vmin=0, vmax=1)
    plt.show()


def show_six_images(image, image2, image3, image4, image5, image6, text=None):

    plt.figure(1)
    plt.subplot(231)
    plt.imshow(image)

    if text is not None:
        plt.title(text)

    plt.subplot(232)

    plt.imshow(image2, cmap="gray", vmin=0, vmax=1)

    plt.subplot(233)

    plt.imshow(image3, cmap="gray", vmin=0, vmax=1)

    plt.subplot(234)

    plt.imshow(image4, cmap="gray", vmin=0, vmax=1)

    plt.subplot(235)

    plt.imshow(image5, cmap="gray", vmin=0, vmax=1)

    plt.subplot(236)

    plt.imshow(image6, cmap="gray", vmin=0, vmax=1)

    plt.show()


def show_image_per_channel(image, text):

    plt.figure(1)
    plt.subplot(311)
    plt.imshow(image[:, :, 0])
    plt.subplot(312)
    plt.imshow(image[:, :, 1])
    plt.subplot(313)
    plt.imshow(image[:, :, 2])

    plt.show()
    if text is not None:
        plt.title(text)


def pixels_to_labels(image):

    print(image.shape)
    output = np.zeros((image.shape[0], image.shape[1]))

    output[np.where(image > 200)] = 255
    output[np.where(image < 200)] = 0


def jpg_to_png(path):

    image_paths = glob(os.path.join(path, "*.jpg"))

    for i in range(0, len(image_paths)):
        print(i)
        name = image_paths[i]
        idx = name.rfind("/")
        image = scipy.misc.imread(image_paths[i])
        scipy.misc.imsave(
            os.path.join(path, name[idx + 1 : -4] + ".png"), image.astype(np.uint8)
        )  # Really important to convert to uint8


def jpeg_to_jpg(path):

    image_paths = glob(os.path.join(path, "*.jpeg"))

    for i in range(0, len(image_paths)):
        print(i)
        name = image_paths[i]
        idx = name.rfind("/")
        image = scipy.misc.imread(image_paths[i])
        scipy.misc.imsave(
            os.path.join(path, name[idx + 1 : -5] + ".jpg"), image
        )  # Really important to convert to uint8


def overlap_image_with_label(image: np.ndarray, mask: np.ndarray) -> np.ndarray:
    """
    This function overlaps the mask with the image.
    In other words, it only plots the image where the
    segmentation mask is >= 1
    Args:
    - image (numpy array): RGB image of shape (1, 480, 640, 3)
    - label (numpy array): segmentation mask of shape (480, 640)
    Returns:
    - overlap (numpy array): overlapped image
    """
    binary_mask = (mask > 0).astype(np.uint8)
    overlapped = image.copy().squeeze(0)
    overlapped[binary_mask == 0] = 0

    return overlapped


def show_x_images(images, titles=None, cmap=None, horizontal=False):
    num_images = len(images)

    if num_images == 1:
        plt.imshow(images[0])
        if titles:
            plt.title(titles[0])
        plt.show()
        return

    if horizontal:
        cols = num_images
        rows = 1
    else:
        cols = math.ceil(math.sqrt(num_images))
        rows = math.ceil(num_images / cols)

    plt.figure(figsize=(15, 5))
    for i, image in enumerate(images):
        plt.subplot(rows, cols, i + 1)
        if cmap:
            plt.imshow(image, cmap=cmap)
        else:
            plt.imshow(cv2.cvtColor(image.astype("float32"), cv2.COLOR_BGR2RGB))

        if titles and i < len(titles):
            plt.title(titles[i])

    plt.show()
    return