#%% modules
import cv2, copy, os, sys, warnings
import numpy as np
import pandas as pd

from math import inf

from phenopype.settings import colours
from phenopype.core.preprocessing import invert_image
from phenopype.utils_lowlevel import _create_mask_bool, _image_viewer, _auto_line_width

#%% functions


def blur(obj_input, kernel_size=5, method="averaging", sigma_color=75, sigma_space=75):
    """
    Apply a blurring algorithm to the image.

    Parameters
    ----------
    obj_input : array or container
        input object
    kernel_size: int, optional
        size of the blurring kernel (has to be odd - even numbers will be ceiled)
    method: {averaging, gaussian, median, bilateral} str, optional
        blurring algorithm
    sigma_colour: int, optional
        for 'bilateral'
    sigma_space: int, optional
        for 'bilateral'

    Returns
    -------
    image : array or container
        blurred image

    """
    ## kwargs
    if kernel_size % 2 == 0:
        kernel_size = kernel_size + 1

    ## load image
    if obj_input.__class__.__name__ == "ndarray":
        image = obj_input
    elif obj_input.__class__.__name__ == "container":
        image = copy.deepcopy(obj_input.image)

    ## method
    if method == "averaging":
        image = cv2.blur(image, (kernel_size, kernel_size))
    elif method == "gaussian":
        image = cv2.GaussianBlur(image, (kernel_size, kernel_size), 0)
    elif method == "median":
        image = cv2.medianBlur(image, kernel_size)
    elif method == "bilateral":
        image = cv2.bilateralFilter(image, kernel_size, sigma_color, sigma_space)
    else:
        image = image

    ## return
    if obj_input.__class__.__name__ == "container":
        obj_input.image = image
    else:
        return image


def draw(
    obj_input,
    overwrite=False,
    tool="line",
    line_colour="black",
    line_width="auto",
    **kwargs
):
    """
    Draw lines, rectangles or polygons onto a colour or binary image. Can be 
    used to connect. disconnect or erase contours. 

    Parameters
    ----------
    obj_input : array or container
        input object
    overwrite : bool, optional
        if a container is supplied, or when working from the pype, should any 
        exsting drawings be overwritten
    tool : {line, polygon, rectangle} str, optional
        type of tool to use for drawing
    line_colour : {"black", "white", "green", "red", "blue"} str, optional
        line or filling (for rectangle and polygon) colour
    line_width : int, optional
        line width

    Returns
    -------
    image : array or container
        image with drawings

    """
    ## kwargs
    flag_overwrite = overwrite
    test_params = kwargs.get("test_params", {})

    ## load image
    df_draw, df_image_data = None, None
    if obj_input.__class__.__name__ == "ndarray":
        image = obj_input
        if df_image_data.__class__.__name__ == "NoneType":
            df_image_data = pd.DataFrame({"filename": "unknown"}, index=[0])
    elif obj_input.__class__.__name__ == "container":
        image = copy.deepcopy(obj_input.image)
        df_image_data = obj_input.df_image_data
        if hasattr(obj_input, "df_draw"):
            df_draw = obj_input.df_draw
    else:
        print("wrong input format.")
        return

    ## more kwargs
    if line_width == "auto":
        line_width = _auto_line_width(image)
    if tool in ["rect", "rectangle", "poly", "polygon"]:
        line_width = -1

    while True:
        ## check if exists
        if not df_draw.__class__.__name__ == "NoneType" and flag_overwrite == False:
            print("- polylines already drawn (overwrite=False)")
            break
        elif not df_draw.__class__.__name__ == "NoneType" and flag_overwrite == True:
            print("- draw polylines (overwriting)")
            pass
        ## future option: edit drawings
        # elif not df_draw.__class__.__name__ == "NoneType" and flag_edit == True:
        #     print("- draw polylines (editing)")
        #     pass
        elif df_draw.__class__.__name__ == "NoneType":
            print("- draw polylines")
            pass

        ## method
        out = _image_viewer(
            image,
            tool=tool,
            draw=True,
            line_width=line_width,
            line_colour=line_colour,
            previous=test_params,
        )

        ## abort
        if not out.done:
            if obj_input.__class__.__name__ == "ndarray":
                print("terminated polyline creation")
                return
            elif obj_input.__class__.__name__ == "container":
                print("- terminated polyline creation")
                return True

        ## create df
        df_draw = pd.DataFrame({"tool": tool}, index=[0])
        df_draw["line_width"] = line_width
        df_draw["colour"] = line_colour
        df_draw["coords"] = str(out.point_list)

        break

    ## draw
    for idx, row in df_draw.iterrows():
        coord_list = eval(row["coords"])
        for coords in coord_list:
            if row["tool"] in ["line", "lines"]:
                cv2.polylines(
                    image,
                    np.array([coords]),
                    False,
                    colours[row["colour"]],
                    row["line_width"],
                )
            elif row["tool"] in ["rect", "rectangle", "poly", "polygon"]:
                cv2.fillPoly(image, np.array([coords]), colours[row["colour"]])

    ## return
    if obj_input.__class__.__name__ == "ndarray":
        df_draw = pd.concat(
            [
                pd.concat([df_image_data] * len(df_draw)).reset_index(drop=True),
                df_draw.reset_index(drop=True),
            ],
            axis=1,
        )
        return image
    elif obj_input.__class__.__name__ == "container":
        obj_input.df_draw = df_draw
        obj_input.image = image


def find_contours(
    obj_input,
    df_image_data=None,
    approximation="simple",
    retrieval="ext",
    offset_coords=(0, 0),
    min_nodes=3,
    max_nodes=inf,
    min_area=0,
    max_area=inf,
    min_diameter=0,
    max_diameter=inf,
    subset=None,
):
    """
    Find objects in binarized image. The input image needs to be a binarized image
    or a container on which a binarizing segmentation algorithm (e.g. threshold or 
    watershed) has been performed. A flavor of different approximation algorithms 
    and retrieval intstructions can be applied. The contours can also be filtered
    for minimum area, diameter or nodes (= number of corners). 

    Parameters
    ----------
    obj_input : array or container
        input object (binary)
    df_image_data : DataFrame, optional
        an existing DataFrame containing image metadata 
    approximation : {"none", "simple", "L1", "KCOS"] str, optional
        contour approximation algorithm:
            - none: no approximation, all contour coordinates are returned
            - simple: only the minimum coordinates required are returned
            - L1: Teh-Chin chain approximation algorithm 
            - KCOS: Teh-Chin chain approximation algorithm 
    retrieval : {"ext", "list", "tree", "ccomp", "flood"} str, optional
        contour retrieval procedure:
            - ext: retrieves only the extreme outer contours
            - list: retrieves all of the contours without establishing any 
              hierarchical relationships
            - tree: retrieves all of the contours and reconstructs a full 
              hierarchy of nested contours
            - ccomp: retrieves all of the contours and organizes them into a 
              two-level hierarchy (parent and child)
            - flood: flood-fill algorithm
    offset_coords : tuple, optional
        offset by which every contour point is shifted.
    subset: {"parent", "child"} str, optional
        retain only a subset of the parent-child order structure
    min_nodes : int, optional
        minimum number of coordinates
    max_nodes : int, optional
        maximum number of coordinates 
    min_area : int, optional
        minimum contour area in pixels
    max_area : int, optional
        maximum contour area in pixels
    min_diameter : int, optional
        minimum diameter of boundary circle
    max_diameter : int, optional
        maximum diameter of boundary circle

    Returns
    -------
    df_contours : DataFrame or container
        contains contours

    """

    ## kwargs
    retr_alg = {
        "ext": cv2.RETR_EXTERNAL,  ## only external
        "list": cv2.RETR_LIST,  ## all contours
        "tree": cv2.RETR_TREE,  ## fully hierarchy
        "ccomp": cv2.RETR_CCOMP,  ## outer perimeter and holes
        "flood": cv2.RETR_FLOODFILL,
    }  ## not sure what this does
    approx_alg = {
        "none": cv2.CHAIN_APPROX_NONE,  ## no approximation of the contours, all points
        "simple": cv2.CHAIN_APPROX_SIMPLE,  ## minimal corners
        "L1": cv2.CHAIN_APPROX_TC89_L1,  ## algorithm 1
        "KCOS": cv2.CHAIN_APPROX_TC89_KCOS,
    }  ## algorithm 2
    flag_subset = subset

    ## load image
    if obj_input.__class__.__name__ == "ndarray":
        image = obj_input
        if df_image_data.__class__.__name__ == "NoneType":
            df_image_data = pd.DataFrame({"filename": "unknown"}, index=[0])
    elif obj_input.__class__.__name__ == "container":
        image = obj_input.image
        df_image_data = obj_input.df_image_data
    else:
        print("wrong input format.")
        return

    ## check
    if len(image.shape) > 2:
        print("Multi-channel array supplied - need binary array.")

    ## method
    image, contour_list, hierarchy = cv2.findContours(
        image=image,
        mode=retr_alg[retrieval],
        method=approx_alg[approximation],
        offset=offset_coords,
    )

    ## filtering
    if not contour_list.__class__.__name__ == "NoneType" and len(contour_list) > 0:
        contour_dict = {}
        idx = 0
        for contour, hier in zip(contour_list, hierarchy[0]):

            ## number of contour nodes
            if len(contour) > min_nodes and len(contour) < max_nodes:

                ## measure basic shape features
                center, radius = cv2.minEnclosingCircle(contour)
                center = int(center[0]), int(center[1])
                diameter = int(radius * 2)
                area = int(cv2.contourArea(contour))

                ## contour hierarchy
                if hier[3] == -1:
                    cont_order = "parent"
                else:
                    cont_order = "child"

                ## retain only a subset
                if not flag_subset.__class__.__name__ == "NoneType":
                    if flag_subset == cont_order:
                        pass
                    else:
                        continue

                ## feature filter
                if all(
                    [
                        diameter > min_diameter and diameter < max_diameter,
                        area > min_area and area < max_area,
                    ]
                ):
                    idx += 1
                    contour_label = str(idx)
                    contour_dict[contour_label] = {
                        "contour": contour_label,
                        "center": center,
                        "diameter": diameter,
                        "area": area,
                        "order": cont_order,
                        "idx_child": hier[2],
                        "idx_parent": hier[3],
                        "coords": contour,
                    }
    else:
        print("No contours found.")

    ## output
    df_contours = pd.DataFrame(contour_dict).T
    df_contours.reset_index(drop=True, inplace=True)
    df_contours = pd.concat(
        [
            pd.concat([df_image_data] * len(df_contours)).reset_index(drop=True),
            df_contours.reset_index(drop=True),
        ],
        axis=1,
    )

    if obj_input.__class__.__name__ == "ndarray":
        return df_contours
    elif obj_input.__class__.__name__ == "container":
        obj_input.df_contours = df_contours


def morphology(obj_input, kernel_size=5, shape="rect", operation="close", iterations=1):
    """
    Performs advanced morphological transformations using erosion and dilation 
    as basic operations. Provides different kernel shapes and a suite of operation
    types (read more about morphological operations here:
    https://docs.opencv.org/master/db/df6/tutorial_erosion_dilatation.html)

    Parameters
    ----------
    obj_input : array or container
        input object
    kernel_size: int, optional
        size of the morphology kernel (has to be odd - even numbers will be 
        ceiled)
    shape : {"rect", "cross", "ellipse"} str, optional
        shape of the kernel
    operation : {erode, dilate, open, close, gradient, tophat, blackhat, hitmiss} str, optional
        the morphology operation to be performed:
            - erode: remove pixels from the border 
            - dilate: add pixels to the border
            - open: erosion followed by dilation
            - close: dilation followed by erosion
            - gradient: difference between dilation and erosion of an input image
            - tophat: difference between input image and opening of input image
            - blackhat: difference between input image and closing of input image
            - hitmiss: find patterns in binary images (read more here:
              https://docs.opencv.org/master/db/d06/tutorial_hitOrMiss.html)
    iterations : int, optional
        number of times to run morphology operation

    Returns
    -------
    image : array or container
        processed image

    """
    ## kwargs
    if kernel_size % 2 == 0:
        kernel_size = kernel_size + 1
    shape_list = {
        "cross": cv2.MORPH_CROSS,
        "rect": cv2.MORPH_RECT,
        "ellipse": cv2.MORPH_ELLIPSE,
    }
    operation_list = {
        "erode": cv2.MORPH_ERODE,
        "dilate": cv2.MORPH_DILATE,
        "open": cv2.MORPH_OPEN,
        "close": cv2.MORPH_CLOSE,
        "gradient": cv2.MORPH_GRADIENT,
        "tophat ": cv2.MORPH_TOPHAT,
        "blackhat": cv2.MORPH_BLACKHAT,
        "hitmiss": cv2.MORPH_HITMISS,
    }
    kernel = cv2.getStructuringElement(shape_list[shape], (kernel_size, kernel_size))

    ## load image
    if obj_input.__class__.__name__ == "ndarray":
        image = obj_input
    elif obj_input.__class__.__name__ == "container":
        image = obj_input.image

    ## method
    image = cv2.morphologyEx(
        image, op=operation_list[operation], kernel=kernel, iterations=iterations
    )

    ## return
    if obj_input.__class__.__name__ == "container":
        obj_input.image = image
    else:
        return image


def threshold(
    obj_input,
    df_masks=None,
    method="otsu",
    constant=1,
    blocksize=99,
    value=127,
    channel="gray",
    invert=False,
):
    """
    Applies a fixed-level threshold to create a binary image from a grayscale 
    image or a multichannel image (gray, red, green or blue channel can be selected).
    Three types of thresholding algorithms are supported. Binary mask coordinates
    can be supplied to include or exclude areas.

    Parameters
    ----------
    obj_input : array or container
        input object
    df_masks : DataFrame, optional
        contains mask coordinates
    method : {"otsu", "adaptive", "binary"} str, optional
        type of thresholding algorithm:
            - otsu: use Otsu algorithm to choose the optimal threshold value
            - adaptive: dynamic threshold values across image (uses arguments
              "blocksize" and "constant")
            - binary: fixed threshold value (uses argument "value")
    blocksize: int, optional
        Size of a pixel neighborhood that is used to calculate a threshold 
        value for the pixel (has to be odd - even numbers will be ceiled; for
        "adaptive" method)
    constant : int, optional
        value to subtract from binarization output (for "adaptive" method)
    value : {between 0 and 255} int, optional
        thesholding value (for "binary" method)
    channel {"gray", "red", "green", "blue"}: str, optional
        which channel of the image to use for thresholding 
    invert : bool, optional
        invert all pixel values PRIOR to applying threshold values

    Returns
    -------
    image : array or container
        binary image

    """

    ##kwargs
    if blocksize % 2 == 0:
        blocksize = blocksize + 1
    if value > 255:
        value = 255
    elif value < 0:
        value = 0

    ## load image
    if obj_input.__class__.__name__ == "ndarray":
        image = copy.deepcopy(obj_input)
    elif obj_input.__class__.__name__ == "container":
        image = copy.deepcopy(obj_input.image)
        if hasattr(obj_input, "df_masks"):
            df_masks = copy.deepcopy(obj_input.df_masks)

    ## channel
    if len(image.shape) == 3:
        if channel == "gray" or channel == "grayscale":
            image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        elif channel == "g" or channel == "green":
            image = image[:, :, 0]
        elif channel == "r" or channel == "red":
            image = image[:, :, 1]
        elif channel == "blue" or channel == "b":
            image = image[:, :, 2]

    if invert:
        image = invert_image(image)

    ## method
    if method == "otsu":
        ret, image = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    elif method == "adaptive":
        image = cv2.adaptiveThreshold(
            image,
            255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY_INV,
            blocksize,
            constant,
        )
    elif method == "binary":
        ret, image = cv2.threshold(image, value, 255, cv2.THRESH_BINARY_INV)

    ## apply masks
    if not df_masks.__class__.__name__ == "NoneType":
        ## include == True
        mask_bool, include_idx = np.zeros(image.shape, dtype=bool), 0
        for index, row in df_masks.iterrows():
            if row["include"] == True:
                if not row["mask"] == "":
                    coords = eval(row["coords"])
                    mask_bool = np.logical_or(mask_bool, _create_mask_bool(image, coords))
                    print('- include mask "' + row["mask"] + '" pixels')
                    include_idx += 1
        if include_idx > 0:
            image[mask_bool == False] = 0
        ## include == False
        for index, row in df_masks.iterrows():
            if row["include"] == False:
                if not row["mask"] == "":
                    coords = eval(row["coords"])
                    image[_create_mask_bool(image, coords)] = 0
                    print('- exclude mask "' + row["mask"] + '" pixels')

    ## return
    if obj_input.__class__.__name__ == "ndarray":
        return image
    elif obj_input.__class__.__name__ == "container":
        obj_input.image = image
        obj_input.image_bin = image


def watershed(
    obj_input,
    image_thresh=None,
    iterations=3,
    kernel_size=3,
    distance_cutoff=0.5,
    distance_mask=0,
    distance_type="l1",
):
    """
    Performs non-parametric marker-based segmentation - useful if many detected 
    contours are touching or overlapping with each other. Input image should be 
    a binary image that serves as the true background. Iteratively, edges are 
    eroded, the difference serves as markers.

    Parameters
    ----------
    obj_input : array or container
        input object
    image_thresh: array, optional
        thresholded image n
    kernel_size: int, optional
        size of the diff-kernel (has to be odd - even numbers will be ceiled)
    iterations : int, optional
        number of times to apply diff-operation
    distance_cutoff : {between 0 and 1} float, optional
        watershed distance transform cutoff (larger values discard more pixels)
    distance_mask : {0, 3, 5} int, optional
        size of distance mask - not all sizes work with all distance types (will
        be coerced to 0)
    distance_type : {"user", "l1", "l2", "C", "l12", "fair", "welsch", "huber"} str, optional
        distance transformation type

    Returns
    -------
    image : array or container
        binary image

    """

    ##kwargs
    distance_type_list = {
        "user": cv2.DIST_USER,
        "l1": cv2.DIST_L1,
        "l2": cv2.DIST_L2,
        "C": cv2.DIST_C,
        "l12": cv2.DIST_L12,
        "fair": cv2.DIST_FAIR,
        "welsch": cv2.DIST_WELSCH,
        "huber": cv2.DIST_HUBER,
    }
    if kernel_size % 2 == 0:
        kernel_size = kernel_size + 1

    ## load image
    if obj_input.__class__.__name__ == "ndarray":
        image = copy.deepcopy(obj_input)
        thresh = copy.deepcopy(image_thresh)
    elif obj_input.__class__.__name__ == "container":
        thresh = copy.deepcopy(obj_input.image)
        image = copy.deepcopy(obj_input.image_copy)

    if thresh.__class__.__name__ == "NoneType":
        print("No thresholded version of image provided for watershed - aborting.")
        return
    if len(thresh.shape) == 3:
        thresh = cv2.cvtColor(thresh, cv2.COLOR_BGR2GRAY)

    ## sure background
    ## note: sure_bg is set as the thresholded input image
    sure_bg = copy.deepcopy(thresh)

    ## sure foreground
    if distance_type in ["user", "l12", "fair", "welsch", "huber"]:
        distance_mask = 0
    opened = morphology(
        thresh,
        operation="open",
        shape="ellipse",
        kernel_size=kernel_size,
        iterations=iterations,
    )
    dist_transform = cv2.distanceTransform(
        opened, distance_type_list[distance_type], distance_mask
    )

    dist_transform = cv2.normalize(
        dist_transform, dist_transform, 0, 1.0, cv2.NORM_MINMAX
    )
    ret, sure_fg = cv2.threshold(
        dist_transform, distance_cutoff * dist_transform.max(), 1, 0
    )

    ## finding unknown region
    sure_fg = sure_fg.astype("uint8")
    sure_fg[sure_fg == 1] = 255
    unknown = cv2.subtract(sure_bg, sure_fg)

    ## marker labelling
    ret, markers = cv2.connectedComponents(sure_fg)
    markers = markers + 1
    markers[unknown == 255] = 0

    ## watershed
    markers = cv2.watershed(image, markers)

    ## convert to contours
    watershed_mask = np.zeros(image.shape[:2], np.uint8)
    watershed_mask[markers == -1] = 255
    watershed_mask[0 : watershed_mask.shape[0], 0] = 0
    watershed_mask[0 : watershed_mask.shape[0], watershed_mask.shape[1] - 1] = 0
    watershed_mask[0, 0 : watershed_mask.shape[1]] = 0
    watershed_mask[watershed_mask.shape[0] - 1, 0 : watershed_mask.shape[1]] = 0

    ## return
    if obj_input.__class__.__name__ == "ndarray":
        return watershed_mask
    elif obj_input.__class__.__name__ == "container":
        obj_input.image = watershed_mask
