from typing import List, Generator, Optional, Tuple
from fileutils import FILE, dslice, search, DIRIN, DIROUT, DIRCACHE
import window
from math import pi
import cv2
import pandas as pd
import numpy as np
from scipy import ndimage as ndi
from skimage.filters import threshold_otsu, threshold_li
from skimage import feature, measure, restoration, segmentation
from pythreshold.global_th.entropy.kapur import kapur_threshold

PIXELS_TO_UM = 3.2  # 1 pixel equals 3.2 um

def rball(img, *argv, **args):
    return img - restoration.rolling_ball(img, *argv, **args)


def watershed_new(img, blur: int, min_thresh: int, dist: int, min_size: int, min_roundness: float, min_mean_brightness: int, ignore: List[Tuple[int, int]]):
    MAX_VAL = 255 if img.dtype == "uint8" else 65535
    BRIGHT_MUL = 65535/255 if img.dtype == "uint8" else 1

    c_img = cv2.cvtColor(img,cv2.COLOR_GRAY2RGB)

    if blur > 0:
        b_img = cv2.GaussianBlur(img, (blur, blur), 0)
        thresh_val = kapur_threshold(b_img)
        thresh_val = thresh_val if thresh_val >= min_thresh else min_thresh
        thresh = cv2.threshold(cv2.bitwise_not(b_img), MAX_VAL-thresh_val, MAX_VAL, cv2.THRESH_BINARY_INV)[1]
    else:
        thresh_val = kapur_threshold(img)
        thresh_val = thresh_val if thresh_val >= min_thresh else min_thresh
        thresh = cv2.threshold(cv2.bitwise_not(img), MAX_VAL-thresh_val, MAX_VAL, cv2.THRESH_BINARY_INV)[1]
    thresh = segmentation.clear_border(thresh)
    thresh = ndi.binary_fill_holes(thresh)

    distance = ndi.distance_transform_edt(thresh)
    if dist < 1:
        localMax = distance > distance.min()
    else:
        localMax = np.zeros_like(img, dtype=np.bool)
        localMax[tuple(feature.peak_local_max(distance, exclude_border=False, min_distance=dist, labels=thresh).T)] = True


    markers = ndi.label(localMax, structure=np.ones((3, 3)))[0]
    labels = segmentation.watershed(-img, markers, mask=thresh)

    grain_num = 1
    meta = {
        "outlines": np.full_like(img, MAX_VAL, dtype=img.dtype),
        "data": [[ "Grain","Area","Mean","Min","Max","Circ", "Eccentricity", "IntDen","RawIntDen","AR","Round","Solidity" ]]
    }

    for prop, label in list(zip(measure.regionprops(labels, intensity_image=img), np.unique(labels)[1:]))[:50]:
        bbox = prop["bbox"]

        if prop["area"]*PIXELS_TO_UM**2 < min_size \
           or 1-prop["eccentricity"] < min_roundness \
           or prop["mean_intensity"]*BRIGHT_MUL < min_mean_brightness*257 \
           or grain_num > 50 \
           or any([bbox[1] <= x <= bbox[3] and bbox[0] <= y <= bbox[2]
                   for (x, y) in ignore]):
            continue

        # Create mask of image size and mark labeled area
        mask = np.zeros(img.shape, dtype=img.dtype)
        mask[labels == label] = MAX_VAL
        # cnts = measure.find_contours(mask.copy())
        # print(cnts)
        # c_img[cnts] = (65535,65535)
        # measure.drawContours()

        # c_img = segmentation.mark_boundaries(c_img, mask, (0,255,255))
        # Create and draw contour around labeled area
        if cv2.__version__.startswith("3."):
            cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL,
                                    cv2.CHAIN_APPROX_SIMPLE)[1]
        else:
            cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL,
                                    cv2.CHAIN_APPROX_SIMPLE)[0]

        c_img = cv2.drawContours(c_img, cnts, -1, (0,MAX_VAL,MAX_VAL), 1)
        meta["outlines"] = cv2.drawContours(meta["outlines"], cnts, -1, (0,0), 1)
        c = max(cnts, key=cv2.contourArea)
        (x,y), r = cv2.minEnclosingCircle(c)
        x = int(float(x) + float(r)/2)
        y = int(float(y) + float(r)/2)
        cv2.putText(c_img, str(grain_num), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
        cv2.putText(meta["outlines"], str(grain_num), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0), 1)
        meta["data"].append([
            grain_num,
            round(prop["area"]*PIXELS_TO_UM**2, 4),
            round(prop["mean_intensity"]*BRIGHT_MUL, 4),
            round(prop["min_intensity"]*BRIGHT_MUL, 4),
            round(prop["max_intensity"]*BRIGHT_MUL, 4),
            (1 if (4 * pi * prop.area) / (prop.perimeter_crofton * prop.perimeter_crofton) >=1 else round((4 * pi * prop.area) / (prop.perimeter_crofton * prop.perimeter_crofton), 4)),
            round(1-prop["eccentricity"], 4), # 4*pi*(prop["area"])/prop["perimeter"]**2,
            round((prop["area"]*PIXELS_TO_UM**2)*prop["mean_intensity"]*BRIGHT_MUL, 4),
            round(prop["intensity_image"].sum()*BRIGHT_MUL, 4),
            (None if prop["minor_axis_length"] == 0 else round(prop["major_axis_length"]/prop["minor_axis_length"], 4)),
            (None if prop["major_axis_length"] == 0 else round(4*prop["area"]*PIXELS_TO_UM**2/(pi*prop["major_axis_length"]**2)/10, 4)),
            round(prop["solidity"], 4)
        ])
        grain_num += 1
    return c_img, meta

def main():
    """
    tif metadata documentation:
    threshold-low  -> lower bound bit range (0)
    threshold-high -> upper bound bit range (65535)
    https://docs.opencv.org/3.0-beta/modules/core/doc/operations_on_arrays.html#cv2.normalize
    """
    files = [FILE(p) for p in search([DIRIN]) if p.endswith(".tif")]
    if files != []:
        files.sort()
        normalize = ["normalized", (cv2.normalize, [], {"dst": None, "alpha": 0, "beta": 65535, "norm_type": cv2.NORM_MINMAX}, True)]
        to8bit = ["8bit", (cv2.convertScaleAbs, [] , {"alpha": 255.0/65535.0}, True)]
        rollingball = ["rball", (rball, [], {"radius": 50.0}, True)]

        wat = ["wshed", (watershed_new, [], {"blur": 0, "dist": 0, "min_thresh": 20, "min_size": 250, "min_roundness": 0.2, "min_mean_brightness": 50, "ignore": []}, False)]
        [f.opqueue.__init__([normalize, to8bit, rollingball, wat]) for f in files]
        [f.apply(dslice(f.opqueue,None,-1)) for f in files]

        window.init(files)
        window.reset()
        return True
    print(f"ERR: No files in {DIRIN}, please add some.")
    return False

if __name__ == '__main__':
    if main():
        window.run()