# -*- coding: utf-8 -*-
"""
cvtools - image processing tools for plankton images

"""
import time
import json
import os
import sys
import glob
import datetime
import pickle
import random, string
from math import pi
import cv2
from skimage import morphology, measure, exposure, restoration
from skimage import transform
from skimage.feature import register_translation
from skimage.filters import threshold_otsu, scharr, gaussian
import numpy as np
from scipy import ndimage, spatial
import xmlsettings

def img_fill(im_in):  # n = binary image threshold

    # Copy the thresholded image.
    im_floodfill = im_in.copy()

    # Mask used to flood filling.
    # Notice the size needs to be 2 pixels than the image.
    h, w = im_in.shape[:2]
    mask = np.zeros((h + 2, w + 2), np.uint8)

    # Floodfill from point (0, 0)
    cv2.floodFill(im_floodfill, mask, (0, 0), 255)

    # Invert floodfilled image
    im_floodfill_inv = cv2.bitwise_not(im_floodfill)

    # Combine the two images to get the foreground.
    fill_image = im_in | im_floodfill_inv

    return fill_image


def auto_canny(image, sigma=0.66):
    # compute the median of the single channel pixel intensities
    v = np.median(image)

    # apply automatic Canny edge detection using the computed median
    lower = int(max(0, (1.0 - sigma) * v))
    upper = int(min(255, (1.0 + sigma) * v))
    edged = cv2.Canny(image, lower, upper)

    # return the edged image
    return edged

def enlarge_bbox(box, scale, image):

    new_box = [0, 0, 0, 0]

    def shift_with_limit(val, shift, limit):

        val = int(val + shift)
        if val*np.sign(shift) > limit:
            val = limit
        return val

    width = box[3] - box[1]
    height = box[2] - box[0]

    # divide scale by 2 such that scaling by 2 will make the box dims grow by 2
    extra_width = width * scale / 2
    extra_height = height * scale / 2

    new_box[0] = shift_with_limit(box[0], -extra_height / 2, 0)
    new_box[2] = shift_with_limit(box[2], extra_height / 2, image.shape[0])
    new_box[1] = shift_with_limit(box[1], -extra_width / 2, 0)
    new_box[3] = shift_with_limit(box[3], extra_width / 2, image.shape[1])

    return tuple(new_box)

def make_gaussian(size, fwhm = 3, center=None):
    """ Make a square gaussian kernel.
    size is the length of a side of the square
    fwhm is full-width-half-maximum, which
    can be thought of as an effective radius.
    """

    x = np.arange(0, size, 1, float)
    y = x[:,np.newaxis]

    if center is None:
        x0 = y0 = size // 2
    else:
        x0 = center[0]
        y0 = center[1]

    output = np.exp(-4*np.log(2) * ((x-x0)**2 + (y-y0)**2) / fwhm**2)
    output = output/np.sum(output)

    return output

# import raw image
def import_image(abs_path,filename,raw=True,bayer_pattern=cv2.COLOR_BAYER_RG2RGB):

    # Load and convert image as needed
    img_c = cv2.imread(os.path.join(abs_path,filename),cv2.IMREAD_UNCHANGED)
    if raw:
        img_c = cv2.cvtColor(img_c,bayer_pattern)

    # force image to be 3-channel if it is not
    if len(img_c.shape) < 3:
        img_c = cv2.cvtColor(img_c, cv2.COLOR_GRAY2RGB)

    return img_c

# convert image to 8 bit with or without autoscaling
def convert_to_8bit(img,auto_scale=True):

    # Convert to 8 bit and autoscale
    if auto_scale:

        result = np.float32(img)-np.min(img)
        result[result<0.0] = 0.0
        if np.max(img) != 0:
            result = result/np.max(img)

        img_8bit = np.uint8(255*result)
    else:
        img_8bit = np.unit8(img)

    return img_8bit


def intensity_features(img, obj_mask):
    res = {}

    # assume that obj_mask contains one connected component
    prop = measure.regionprops(obj_mask.astype(np.uint8), img)[0]
    res["mean_intensity"] = prop.mean_intensity

    intensities = prop.intensity_image[prop.image]
    res["median_intensity"] = np.median(intensities)
    res["std_intensity"] = np.std(intensities)
    res["perc_25_intensity"] = np.percentile(intensities, 25)
    res["perc_75_intensity"] = np.percentile(intensities, 75)

    centroid = np.array(prop.centroid)
    weighted_centroid = np.array(prop.weighted_centroid)
    displacement = weighted_centroid - centroid
    displacement_image = np.linalg.norm(displacement / img.shape)
    if prop.minor_axis_length > 0:
        displacement_minors = np.linalg.norm(displacement) / prop.minor_axis_length
    else:
        displacement_minors = 0
    res['mass_displace_in_images'] = displacement_image
    res['mass_displace_in_minors'] = displacement_minors

    res["moment_hu_1"] = prop.weighted_moments_hu[0]
    res["moment_hu_2"] = prop.weighted_moments_hu[1]
    res["moment_hu_3"] = prop.weighted_moments_hu[2]
    res["moment_hu_4"] = prop.weighted_moments_hu[3]
    res["moment_hu_5"] = prop.weighted_moments_hu[4]
    res["moment_hu_6"] = prop.weighted_moments_hu[5]
    res["moment_hu_7"] = prop.weighted_moments_hu[6]

    return res

# extract simple features and create a binary representation of the image
def quick_features(img, save_to_disk=False, abs_path='', file_prefix='', cfg = []):
    """
    :param img: 8-bit array
    """
    # Pull out some settings from cfg if available
    if cfg:
        min_obj_area = cfg.get('MinObjectArea', 100)
        objs_per_roi = cfg.get('ObjectsPerROI', 1)
        deconv = cfg.get("Deconvolve").lower() == 'true'
        edge_thresh = cfg.get('EdgeThreshold', 2.5)
        use_jpeg = cfg.get("UseJpeg").lower() == 'true'
        raw_color = cfg.get("SaveRawColor").lower() == 'true'
        compute_edges = cfg.get("ComputeEdges", False)
    else:
        min_obj_area = 100
        objs_per_roi = 1
        deconv = False
        use_jpeg = False
        raw_color = True
        edge_thresh = 2.5
        compute_edges = False

    # Define an empty dictionary to hold all features
    features = {}

    features['rawcolor'] = np.copy(img)
    # compute features from gray image
    gray = np.uint8(np.mean(img, 2))
    bw_img = gray > cfg.get("LumThreshold", 3)

    # clear bw_img
    bw_img_all = bw_img.copy()

    # Compute morphological descriptors
    # use foreground as label image
    #label_img = morphology.label(bw_img, neighbors=8, background=0)
    label_img = bw_img.astype('uint8')
    props = measure.regionprops(label_img, gray)

    bw_img = 0*bw_img

    props = sorted(props, reverse=True, key=lambda k: k.area)

    if len(props) > 0:

        # Init mask with the largest area object in the roi
        bw_img = (label_img) == props[0].label
        bw_img_all = bw_img.copy()

        base_area = props[0].area

        # use only the features from the object with the largest area
        avg_count = 0

        if len(props) > objs_per_roi:
            n_objs = objs_per_roi
        else:
            n_objs = len(props)

        for f in range(0 ,n_objs):

            if props[f].area > min_obj_area:
                bw_img_all = bw_img_all + (label_img== props[f].label)
                avg_count = avg_count + 1

            if f >= objs_per_roi:
                break

        # Take the largest object area as the roi area
        # no average
        avg_area = props[0].area
        avg_maj = props[0].major_axis_length
        avg_min = props[0].minor_axis_length
        avg_or = props[0].orientation
        avg_eccentricity = props[0].eccentricity
        avg_solidity = props[0].solidity

        # Calculate intensity features only for largest
        features_intensity = intensity_features(gray, bw_img)
        features['intensity_gray'] = features_intensity

        features_intensity = intensity_features(img[::, ::, 0], bw_img)
        features['intensity_red'] = features_intensity

        features_intensity = intensity_features(img[::, ::, 1], bw_img)
        features['intensity_green'] = features_intensity

        features_intensity = intensity_features(img[::, ::, 2], bw_img)
        features['intensity_blue'] = features_intensity

        # Check for clipped image
        if np.max(bw_img_all) == 0:
            bw = bw_img_all
        else:
            bw = bw_img_all/np.max(bw_img_all)

        clip_frac = float(np.sum(bw[:,1]) +
                np.sum(bw[:,-2]) +
                np.sum(bw[1,:]) +
                np.sum(bw[-2,:]))/(2*bw.shape[0]+2*bw.shape[1])
        features['clipped_fraction'] = clip_frac

        # Save simple features of the object
        features['area'] = avg_area
        features['minor_axis_length'] = avg_min
        features['major_axis_length'] = avg_maj
        if avg_maj == 0:
            features['aspect_ratio'] = 1
        else:
            features['aspect_ratio'] = avg_min/avg_maj
        features['orientation'] = avg_or
        features['eccentricity'] = avg_eccentricity
        features['solidity'] = avg_solidity
        features['estimated_volume'] = 4.0 / 3 * pi * avg_maj * avg_min * avg_min
        #
        #

        # print "Foreground Objects: " + str(avg_count)

    else:

        features['clipped_fraction'] = 0.0

        # Save simple features of the object
        features['area'] = 0.0
        features['minor_axis_length'] = 0.0
        features['major_axis_length'] = 0.0
        features['aspect_ratio'] = 1
        features['orientation'] = 0.0
        features['eccentricity'] = 0
        features['solidity'] = 0
        features['estimated_volume'] = 0
        features['intensity_gray'] = 0
        features['intensity_red'] = 0
        features['intensity_green'] = 0
        features['intensity_blue'] = 0

    # Masked background with Gaussian smoothing, image sharpening, and
    # reduction of chromatic aberration

    # mask the raw image with smoothed foreground mask
    blurd_bw_img = gaussian(bw_img_all, 3)
    img[:, :, 0] = img[:, :, 0]*blurd_bw_img
    img[:, :, 1] = img[:, :, 1]*blurd_bw_img
    img[:, :, 2] = img[:, :, 2]*blurd_bw_img

    # Make a guess of the PSF for sharpening
    psf = make_gaussian(5, 3, center=None)

    # sharpen each color channel and then reconbine

    if np.max(img) == 0:
        img = np.float32(img)
    else:
        img = np.float32(img)/np.max(img)

    if deconv:

        img[img == 0] = 0.0001
        img[:, :, 0] = restoration.richardson_lucy(img[:, :, 0], psf, 7)
        img[:, :, 1] = restoration.richardson_lucy(img[:, :, 1], psf, 7)
        img[:, :, 2] = restoration.richardson_lucy(img[:, :, 2], psf, 7)

    # Estimate color channel shifts and try to align.
    # this works for most images but some still retain and offset.
    # need to figure out why...
    # r_shift, r_error, r_diffphase = register_translation(img[:,:,1], img[:,:,2],1)
    # b_shift, b_error, b_diffphase = register_translation(img[:,:,1], img[:,:,0],1)

    # # this swap of values is needed for some reason
    # if r_shift[0] < 0 and r_shift[1] < 0:
       # r_shift = -r_shift

    # if b_shift[0] < 0 and b_shift[1] < 0:
       # b_shift = -b_shift

    # r_tform = transform.SimilarityTransform(scale=1,rotation=0,translation=r_shift)
    # img[:,:,2] = transform.warp(img[:,:,2],r_tform)

    # b_tform = transform.SimilarityTransform(scale=1,rotation=0,translation=b_shift)
    # img[:,:,0] = transform.warp(img[:,:,0],b_tform)

    # Rescale image to uint8 0-255
    img[img < 0] = 0

    if np.max(img) == 0:
        img = np.uint8(255*img)
    else:
        img = np.uint8(255*img/np.max(img))

    features['image'] = img
    features['binary'] = 255*bw_img_all

    # Save the binary image and also color image if requested
    if save_to_disk:

        # convert and save images
        # Raw color (no background removal)
        if use_jpeg:
            if raw_color:
                cv2.imwrite(os.path.join(abs_path, file_prefix+"_rawcolor.jpeg"), features['rawcolor'])
            # Save the processed image and binary mask
            cv2.imwrite(os.path.join(abs_path, file_prefix+".jpeg"), features['image'])
        else:
            if raw_color:
                cv2.imwrite(os.path.join(abs_path, file_prefix+"_rawcolor.png"), features['rawcolor'])
            # Save the processed image and binary mask
            cv2.imwrite(os.path.join(abs_path, file_prefix+".png"), features['image'])

        # Binary should also be saved png
        cv2.imwrite(os.path.join(abs_path, file_prefix+"_binary.png"), features['binary'])

    return features