import math
import os
from itertools import groupby
import numpy as np
from matplotlib import pyplot as plt, colors
from skimage.transform import resize

# todo different backend maybe diffrent files thata have partial
# todo depend on mode different min and max
from raven_utils.tools import is_num, il, image_type, lw, is_int, dict_from_list, if_images, download_file, \
    dir_decorator, list_dir

TRAIN = "train"
VAL = "val"
TEST = "test"
COLOR = "color"
SIZE = "size"
EXIST = "exist"
COR = "cor"
TARGET_SIZE = "target_size"


# https://github.com/qubvel/efficientnet/blob/master/efficientnet/preprocessing.py
MAP_INTERPOLATION_TO_ORDER = {
    "nearest": 0,
    "bilinear": 1,
    "biquadratic": 2,
    "bicubic": 3,
}


def center_crop_and_resize(image, image_size, crop_padding=32, interpolation="bicubic"):
    assert image.ndim in {2, 3}
    assert interpolation in MAP_INTERPOLATION_TO_ORDER.keys()

    h, w = image.shape[:2]

    padded_center_crop_size = int(
        (image_size / (image_size + crop_padding)) * min(h, w)
    )
    offset_height = ((h - padded_center_crop_size) + 1) // 2
    offset_width = ((w - padded_center_crop_size) + 1) // 2

    image_crop = image[
                 offset_height: padded_center_crop_size + offset_height,
                 offset_width: padded_center_crop_size + offset_width,
                 ]
    resized_image = resize(
        image_crop,
        (image_size, image_size),
        order=MAP_INTERPOLATION_TO_ORDER[interpolation],
        preserve_range=True,
    )

    return resized_image


def im_conv(image):
    return image.astype(np.float32) / 255.


def add_dim(x):
    return x[..., None]


def im_conv_black(image):
    return image[..., np.newaxis].astype(np.float32) / 255.


def im_resolve(image):
    return (image * 255).astype(np.uint8)


def minmax(data, axis=None, scale=1.0):
    return ((data - np.min(data, axis=axis)) / (np.max(data, axis=axis) - np.min(data, axis=axis))) * scale


def inverse(x):
    return 1.0 - x


def standarize(img, scaler=0.15, batch=True):
    indexes = tuple(range(1 if batch else 0, img.ndim))
    img = np.array(img)
    img -= np.mean(img, axis=indexes)[tuple([slice(None)] + [None] * len(indexes))]
    img /= np.std(img, axis=indexes)[tuple([slice(None)] + [None] * len(indexes))] + 1e-5
    img *= scaler
    return img


def image_standardization(img, scaler=0.15):
    indexes = tuple(range(1, 4))
    img -= np.mean(img, axis=indexes)
    img /= np.std(img, axis=indexes) + 1e-5
    img *= scaler
    return img


def clip_standardized(img, min=0, max=1):
    img += (max - min) / 2
    return np.clip(img, min, max)


def deprocess_image(img, clip=True):
    # Normalize array: center on 0., ensure variance is 0.15
    img = image_standardization(img)
    if clip:
        img = clip_standardized(img)
    return img


# def im_conv_tf(image,dtype=tf.float32):
#     return tf.image.convert_image_dtype(image,dtype)

# todo Grid on edges
# todo different channels
def draw_images(data, col=-1, row=-1, grid=True, border=None, mark=None, grid_args={}, border_kwargs={}, *args,
                **kwargs):
    figure, shape = create_grid(data, row=row, col=col, *args, **kwargs)
    if border is not None:
        add_border(figure, border, **border_kwargs)
    figure = create_image_from_grid(figure, *args, **kwargs)

    if grid:
        figure = add_grid(figure, shape[-3:], **grid_args)
    return figure


def draw_images2(data, col=-1, row=-1, grid=True, color=1, border=None, border_kwargs={}, *args,
                 **kwargs):
    figure, shape = create_grid(data, row=row, col=col, *args, **kwargs)
    if border is not None:
        add_border(figure, border, **border_kwargs)
    figure = fill_canvas(figure, color=color, grid=grid)
    return figure


def draw_images3(data, col=-1, row=-1, grid=True, border=None, mark=None, grid_args={}, border_kwargs={}, *args,
                 **kwargs):
    if len(np.asarray(data).shape) < 4:
        return data

    return draw_images(data=data,
                       col=col,
                       row=row,
                       grid=grid,
                       border=border,
                       mark=mark,
                       grid_args=grid_args,
                       border_kwargs=border_kwargs,
                       *args,
                       **kwargs
                       )


def fill_canvas(a, color=1, grid=10):
    if not isinstance(a, np.ndarray):
        a = np.asarray(a)
    if grid is True:
        grid = 10
    shape = a.shape
    move_r = shape[2] + grid
    move_c = shape[3] + grid
    if color is None:
        color = 0
    if is_num(color):
        color = np.full((move_r * shape[0], move_c * shape[1]) + shape[-1:], fill_value=color, dtype=a.dtype)
    for r in range(shape[0]):
        for c in range(shape[1]):
            color[move_r * r:move_r * r + shape[2], move_c * c:move_c * c + shape[3]] = a[r, c]
    return color


def create_grid(
        data,
        row=-1,
        col=-1,
        mode="auto",
        # Method for creating grid if row and col not set. Even evenly for row and cows. Shape take shape from input data. Auto use shape for input data that have more then 4 dimention, otherwise even.
        swap=False,  # swaping row and columns
        *args, **kwargs):
    if il(data):
        for i, d in enumerate(data):
            if len(d.shape) > 4:
                data[i] = draw_images(d)
            elif len(d.shape) < 3:
                data[i] = d[..., np.newaxis]
        shapes = [d.shape for d in data]
        # padding to make all images same size
        if not all_equal(shapes):
            mh = np.max([sh[-3] for sh in shapes])
            mw = np.max([sh[-2] for sh in shapes])
            canvas = np.zeros((len(data), mh, mw, shapes[0][-1]))
            for j, d in enumerate(data):
                canvas[j, :d.shape[-3], :d.shape[-2]] = d
            data = canvas
    data = np.asarray(data)
    shape = data.shape
    # check in case of no channels
    if shape[-1] > 4:
        data = data[..., np.newaxis]
        shape = data.shape
    if swap and len(shape) > 4:
        data = data.swapaxes(-4, -5)
    h = shape[-3]
    w = shape[-2]
    c = shape[-1]
    if len(shape) > 4:
        im_no = shape[-5] * shape[-4]
    elif len(shape) == 4:
        im_no = shape[-4]
    else:
        return data, shape
    if row == -1 and col == -1:
        if mode == "auto":
            mode = "shape" if len(shape) > 4 else "even"
        if mode == "shape":
            row = shape[-5] if len(shape) > 4 else shape[-4]
        elif mode == "even":
            row = math.ceil(np.sqrt(im_no))
        else:
            row = shape[-5] if len(shape) > 4 else math.ceil(np.sqrt(im_no))
    # eh = -1 if row == -1 else row * h
    # ew = -1 if column == -1 else column * w
    # if number of images is less then available slots in grid add blank images
    if row == -1:
        eh = -1
        m = im_no % col
        if m != 0:
            data = np.concatenate([data, np.zeros((col - m,) + shape[-3:])], axis=-4)
    else:
        eh = row * h
    if col == -1:
        ew = -1
        m = im_no % row
        if m != 0:
            data = np.concatenate([data, np.zeros((row - m,) + shape[-3:])], axis=-4)
    else:
        ew = col * h
    figure = data.reshape(row, col, h, w, c)
    import time
    time.sleep(0.1)
    return figure, shape


def create_image_from_grid(grid, *args, **kwargs):
    shape = grid.shape
    figure = grid.swapaxes(1, 2).reshape(shape[0] * shape[2], shape[1] * shape[3], shape[4])
    return figure


# todo For more dimention
def add_grid(figure, shape, left=True, right=True, color=1.00, hor=True, ver=True):
    revers = False
    if len(figure.shape) == 2:
        figure = figure[..., None]
        reverse = True
    if hor:
        figure[shape[0]::shape[0], :, :] = color
    if ver:
        figure[:, shape[1]::shape[1], :] = color
    if left:
        figure[0, :, :] = color
        figure[0, :] = color
    if right:
        figure[figure.shape[0] - 1, :, :] = color
        figure[:, ::figure.shape[1] - 1, :] = color
    if revers:
        figure = figure[..., 0]
    return figure


def add_border(figure, cor, exist=True, size=4, color="auto"):
    if color == "auto":
        info = image_type(figure)
        if info.mode == "RGB":
            color = np.array(colors.to_rgb("red")) * info.range_[1]
        else:
            color = info.range_[1]
        if info.reverse:
            color = info.range_[1] - color

    shape = figure.shape
    for c in lw(cor):
        if not isinstance(c, dict):
            c = dict_from_list([COR, EXIST, SIZE, COLOR], c)
        c = {
            **{
                EXIST: exist if il(exist) else [exist] * 4,
                SIZE: size,
                COLOR: color
            },
            **c
        }

        if is_int(lw(c[EXIST])[0]):
            c[EXIST] = K.mask(4, c[EXIST])
        if is_int(c[COR]):
            c[COR] = c[COR] // shape[1], c[COR] % shape[1]
            isinstance(lw(c[EXIST])[0], np.integer)

        cr = c[COR]
        if c[EXIST][0]:
            figure[cr[0], cr[1], 0:c[SIZE], :] = c[COLOR]
        if c[EXIST][1]:
            figure[cr[0], cr[1], :, shape[3] - c[SIZE]:shape[3]] = c[COLOR]
        if c[EXIST][2]:
            figure[cr[0], cr[1], shape[2] - c[SIZE]:shape[2], :] = c[COLOR]
        if c[EXIST][3]:
            figure[cr[0], cr[1], :, 0:c[SIZE]] = c[COLOR]
    return figure


# todo hgih tmp fix
def get_digitilizer(min=0, max=1, steps=10, mode="numpy"):
    bins = np.linspace(min, max, steps)
    bins = bins + (bins[1] - bins[0]) / 2
    # if mode == "tf":
    #     from tensorflow.ops import math_ops
    #     func = math_ops._bucketize
    #     bins = list(bins)
    # else:
    func = np.digitize

    def digitilizer(x):
        return func(x, bins)

    return digitilizer


def save_image_pil(path, data, *args, **kwargs):
    from PIL import Image
    if not isinstance(data, np.ndarray):
        data = np.array(data)
    if if_images(data):
        data = draw_images2(data, *args, **kwargs)
    if np.max(data) <= 1:
        data = im_resolve(data)
    if np.shape(data)[-1] == 1:
        data = data[..., 0]
    img = Image.fromarray(np.array(data, dtype="uint8"))
    # test if black or wihte
    if data.shape[-1] not in [3, 4]:
        img = img.convert('L')
    # todo maybe just check type for np.float
    if path is None:
        img.show()
    else:
        img.save(path)


def save_image_pil2(path, data, *args, **kwargs):
    from PIL import Image
    if np.max(data) <= 1:
        data = im_resolve(data)
    if len(np.shape(data)) == 3 and np.shape(data)[-1] == 1:
        data = data[..., 0]
    img = Image.fromarray(data)
    if len(data.shape) == 2 or (len(data.shape == 3) and data.shape[2] == 1):
        img = img.convert('L')
    # todo maybe just check type for np.float
    if path is None:
        img.show()
    else:
        img.save(path)


#
# https://towardsdatascience.com/how-to-create-a-gif-from-matplotlib-plots-in-python-6bec6c0c952c
# duration, fps, loop
def save_gif(path, data,sort_fn="first", *args, **kwargs):
    data = get_paths_(data, sort_fn=sort_fn)
    import imageio
    with imageio.get_writer(path, mode='I', *args, **kwargs) as writer:
        for filename in data:
            image = imageio.imread(filename)
            writer.append_data(image)


# https://stackoverflow.com/questions/44947505/how-to-make-a-movie-out-of-images-in-python
# https://docs.opencv.org/4.x/dd/d9e/classcv_1_1VideoWriter.html#ad59c61d8881ba2b2da22cff5487465b5
# https://github.com/ContinuumIO/anaconda-issues/issues/223
# https://softron.zendesk.com/hc/en-us/articles/207695697-List-of-FourCC-codes-for-video-codecs
def save_video(path, data, *args, fps=1, fourcc="MJPG",sort_fn="first", **kwargs):
    data = get_paths_(data, sort_fn=sort_fn)

    import cv2

    frame = cv2.imread(data[0])
    height, width, layers = frame.shape

    if isinstance(fourcc, str):
        fourcc = cv2.VideoWriter_fourcc(*fourcc)
    video = cv2.VideoWriter(path, fourcc, fps, (width, height))

    for image in data:
        video.write(cv2.imread(image))

    cv2.destroyAllWindows()
    video.release()


def get_paths_(data, sort_fn="first"):
    if isinstance(data, str):
        data = list_dir(data, sort_fn=sort_fn)
    return data


def add_text(ax, text, pos=None, star_x=25, start_y=25, row_size=30):
    text = lw(text)
    if pos is None:
        pos = [(star_x, start_y + i * row_size) for i in range(len(text))]
    pos = lw(pos)
    for i, t in enumerate(text):
        # draw.text(pos[i], t,fill="black")
        ax.text(pos[i][0], pos[i][1], t)


@dir_decorator
def save_image(path, data, description=None, size=None, dpi=None, *args, **kwargs):
    if not isinstance(data, np.ndarray):
        data = np.array(data)
    fig = plt.figure()
    if size is not None:
        fig.set_size_inches(size)
    # if title is not None:
    #     plt.title(title)
    ax = plt.Axes(fig, [0., 0., 1., 1.])
    ax.set_axis_off()
    fig.add_axes(ax)
    if len(data.shape) == 2 or data.shape[-1] == 1:
        # plt.set_cmap('Greys_r')
        plt.set_cmap('Greys')
    else:
        plt.set_cmap('hot')
    if if_images(data):
        data = draw_images(data, *args, **kwargs)
    ax.imshow(data, aspect='equal')
    if description is not None:
        add_text(ax, description, star_x=data.shape[1])

    params = {
        "fname": path,
        "bbox_inches": 'tight',
        "pad_inches": 0
    }
    if dpi is not None:
        params['dpi'] = dpi
    if path is None:
        plt.show()
    else:
        plt.savefig(**params)
    plt.clf()


SAMPLE_IMAGES = {
    "panda": "~/all/dataset/examples/panda.jpg",
    "cat": "~/all/dataset/examples/cat.jpg",
}


def load_image(path=None, dtype="uint8", *args, **kwargs):
    from tensorflow.keras.preprocessing import image
    if path is None:
        path = os.path.expanduser(SAMPLE_IMAGES["panda"])
        if not os.path.isfile(path):
            os.makedirs(os.path.dirname(path), exist_ok=True)
            download_file("https://upload.wikimedia.org/wikipedia/commons/f/fe/Giant_Panda_in_Beijing_Zoo_1.JPG", path)
        if not TARGET_SIZE in kwargs:
            kwargs[TARGET_SIZE] = (224, 224)
    elif path == "cat":
        path = os.path.expanduser(SAMPLE_IMAGES["cat"])
        kwargs[TARGET_SIZE] = (224, 224)

    img = image.load_img(path, *args, **kwargs)
    img = image.img_to_array(img, dtype=dtype)
    return img


def to_rgb(img, add_dim=True, check_dim=True):
    shape = np.shape(img)
    if add_dim and shape[-1] > 4:
        img = img[..., None]
    shape = np.shape(img)
    if check_dim and shape[-1] != 1:
        return img
    return np.tile(img, reps=(1, 1, 3))

# https://stackoverflow.com/questions/3844801/check-if-all-elements-in-a-list-are-identical
def all_equal(iterable, key=None):
    g = groupby(iterable, key=key)
    return next(g, True) and not next(g, False)