import math
import os
import shutil

import cv2
import numpy as np
import tifffile


def draw_gradient(x, y, channel, radius, map):
    """
    Draw gradients to create a localization map.

    Parameters
    ----------
    x : int
    y : int
    channel : int
    radius : int
    map : numpy.ndarray
        empty localization map
    Returns
    -------
    map : numpy.ndarray
        localization map with drawn gradients
    """

    x_min = max(x - radius, 0)
    x_max = min(x + radius, map.shape[0])
    y_min = max(y - radius, 0)
    y_max = min(y + radius, map.shape[1])

    for i in range(x_min, x_max):
        for j in range(y_min, y_max):
            distance = math.sqrt((i - x) ** 2 + (j - y) ** 2)
            if distance < radius:
                gray = int(-(255 / radius) * distance + 255)
                map[i, j, channel] = gray

    return map


def redraw_gradients(original_map, target_shape, radius):
    """
    Re-draw gradients to create new localization map.

    Parameters
    ----------
    original_map : numpy.ndarray
    target_shape : tuple
    radius : int

    Returns
    -------
    map_new : numpy.ndarray
        localozation map with re-drawn gradients
    """

    original_resolution = original_map.shape
    map_new = np.zeros(
        (target_shape[0], target_shape[1], original_map.shape[2]), dtype=np.uint8
    )
    x_ratio = target_shape[0] / original_resolution[0]
    y_ratio = target_shape[1] / original_resolution[1]
    points = np.argwhere(original_map == 255)

    for point in points:
        x = int(point[0] * x_ratio)
        y = int(point[1] * y_ratio)
        z = int(point[2])
        draw_gradient(x, y, z, radius, map_new)

    return map_new


def click_event(event, x, y, flags, params):
    """
    OpenCV click event function.

    .. note:: Imlemented for maximum of three classes.           (left
    click, middle click, right click)
    """

    if event == cv2.EVENT_LBUTTONDOWN:
        print(f"Class 0: [{x},{y}]")

        cv2.circle(img, (x, y), 5, color=255, thickness=2)
        cv2.imshow(f"{image}", img)

        draw_gradient(y, x, 0, 5, map)

    if event == cv2.EVENT_MBUTTONDOWN:
        print(f"Class 1: [{x},{y}]")

        cv2.circle(img, (x, y), 10, color=127, thickness=2)
        cv2.imshow(f"{image}", img)

        draw_gradient(y, x, 1, 10, map)

    if event == cv2.EVENT_RBUTTONDOWN:
        print(f"Class 2: [{x},{y}]")

        cv2.circle(img, (x, y), 15, color=0, thickness=2)
        cv2.imshow(f"{image}", img)

        draw_gradient(y, x, 2, 15, map)


if __name__ == "__main__":

    mode = "label"
    n_classes = 3
    radius = 10

    if mode == "label":
        load_path = "/Users/jan/Pictures/DP/vaid/vaid_512/to_label"
        save_path = "/Users/jan/Pictures/DP/vaid/vaid_512"

        for image in sorted(os.listdir(load_path)):
            if image[-4:] == ".png":
                print(f"Image: {image[:-4]}")
                img = cv2.imread(f"{load_path}/{image}", cv2.IMREAD_UNCHANGED)
                height, width = img.shape[0], img.shape[1]
                map = np.zeros_like(img, shape=(height, width, n_classes))

                cv2.imshow(f"{image}", img)
                cv2.setMouseCallback(f"{image}", click_event)
                cv2.waitKey(0)
                cv2.destroyAllWindows()

                shutil.move(f"{load_path}/{image}", f"{save_path}/images/{image}")
                tifffile.imwrite(
                    f"{save_path}/maps/{image[:-4]}_map.tiff",
                    map,
                    planarconfig="CONTIG",
                )

    elif mode == "relabel":
        load_path = ""
        save_path = ""

        for map in sorted(os.listdir(load_path)):
            if "_map.tiff" in map:
                print(f"map: {map[:-5]}")

                original_map = tifffile.imread(f"{load_path}/{map}")
                target_shape = (512, 512, n_classes)

                synt = redraw_gradients(original_map, target_shape, radius)

                tifffile.imwrite(f"{save_path}/{map}", synt, planarconfig="CONTIG")
