import sys
import os
import glob
import numpy as np
import cv2
import imageio
from skimage.measure import label, regionprops
from skimage.color import label2rgb
import matplotlib.pyplot as plt
import pandas as pd
from pyntcloud import PyntCloud
from scipy.spatial import KDTree

DEBUG = False
NOISE_FLOOR = 10

class DepthImage:

    def __init__(self, file_prefix, types="XYZA"):
        self.depth_images = []
        self.depth_mask = None
        self.types = types
        self.z_index = self.types.find("Z")
        self.file_prefix = file_prefix
        self.load()

    def load(self):
        for t in self.types:
            file_name = self.file_prefix + '-DepthMap-' + t + '.tiff'
            self.depth_images.append(imageio.imread(file_name))

    def mask(self, image_mask):

        # rescale mask as needed
        if not image_mask.shape == self.depth_images[0].shape:
            image_mask = cv2.resize(image_mask, self.depth_images[0].shape)
        # apply mask to each channel
        for i in range(len(self.types)):
            self.depth_images[i] = self.depth_images[i] * image_mask
        # apply mask to depth mask
        self.depth_mask = self.depth_mask * image_mask

    def get_depth_mask(self, depth_range=[-500, 500]):
        if self.depth_mask is None:
            tmp_img = self.depth_images[self.z_index]
            tmp_img[np.isnan(tmp_img)] = depth_range[0] - 100
            self.depth_mask = np.logical_and(tmp_img > depth_range[0], tmp_img <= depth_range[1]).astype('uint8')

        # apply some filtering to the mask before returning
        # don't apply this to the internal mask as we may need it for finding valid depth pixels

        #kernel = np.ones((7, 7), np.uint8)
        #depth_mask = cv2.morphologyEx(self.depth_mask, cv2.MORPH_CLOSE, kernel)
        #depth_mask = img_fill(depth_mask.astype('uint8') * 255) / 255
        #cv2.imshow('depth_mask', cv2.resize(depth_mask, (1024, 1024))*255)
        #cv2.waitKey()
        return self.depth_mask

    def generate_point_cloud(self, depth_range=[-500, 500]):
        if self.depth_mask is None:
            self.get_depth_mask(depth_range)
        # label connected components in the depth image
        label_image = label(self.depth_mask)
        #cv2.imshow('label_image', label2rgb(label_image, image=self.depth_images[2]))
        #cv2.waitKey()
        self.point_cloud = []
        for region in regionprops(label_image):
            coords = region['coords']

            p = []
            p.append(np.mean(self.depth_images[0][coords[:, 0], coords[:, 1]]))
            p.append(np.mean(self.depth_images[1][coords[:, 0], coords[:, 1]]))
            p.append(np.mean(self.depth_images[2][coords[:, 0], coords[:, 1]]))
            p.append(len(coords[:, 0]))
            self.point_cloud.append(p)

        self.point_cloud = np.array(self.point_cloud)
        self.num_particles = self.point_cloud.shape[0]

        if self.point_cloud.shape[0] < 10:
            self.nnstat = 0
            self.nn_points = 0
        else:
            self.kdtree = KDTree(self.point_cloud)
            self.nndist, self.nn_points = self.kdtree.query(self.point_cloud, 2)
            self.nnstat = 2 * np.pi * self.num_particles * np.sum(self.nndist[:, 1] ** 2)

        #print(self.nndist[:, 1])
        #plt.hist(self.nndist[:, 1], bins=500)
        #plt.show()
        return self.point_cloud

def enlarge_bbox(box, scale, image):

    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]

    scaled_width = width * scale
    scaled_height = height * scale

    box[0] = shift_with_limit(box[0], -scaled_height / 2, 0)
    box[2] = shift_with_limit(box[2], scaled_height / 2, image.shape[0])
    box[1] = shift_with_limit(box[1], -scaled_width / 2, 0)
    box[3] = shift_with_limit(box[2], scaled_width / 2, image.shape[1])

    return box



def roi_detector(prefix, output_path, im_in, depth_img, min_area=300, pixel_scale=[78, 97], size_classes=[0.1, 0.15, 0.2, 0.25, 0.3, 0.5, 1, 2, 4, 8, 16], save_rois=False):

    # scale depth mask to image size
    depth_mask = depth_img.get_depth_mask()

    z_img = depth_img.depth_images[2]
    if not im_in.shape == depth_mask.shape:
        depth_mask = cv2.resize(depth_mask, im_in.shape)
        z_img = cv2.resize(z_img, im_in.shape)

    # fill in missing parts of depth mask from img_in
    # depth_mask[im_in > 150] = 1
    depth_mask = cv2.GaussianBlur(depth_mask, (15, 15), 11)
    depth_mask[depth_mask > 0] = 1
    depth_mask = img_fill(depth_mask.astype('uint8') * 255) / 255

    #cv2.imshow('edge', cv2.resize(depth_mask*255, (1024, 1024)))
    #cv2.waitKey()

    # label connected components in the depth image
    label_image = label(depth_mask)
    #cv2.imshow('label_image', label2rgb(label_image))
    #cv2.waitKey()
    rois = []
    particle_stats = np.zeros((1, len(size_classes)))
    depths = []

    im_out = cv2.cvtColor(im_in.copy(), cv2.COLOR_GRAY2RGB)

    counter = 0
    valid_roi_counter = 0
    esds = []
    all_regions = regionprops(label_image, im_in)
    for region in all_regions:

        coords = region['coords']

        box = enlarge_bbox(egion['bbox'], 2, im_in)

        depth_val = np.mean(z_img[[coords[:, 0], coords[:, 1]]])

        pixel_scale = 97 - (97-78)*depth_val/220 # from depth calibration
        #pixel_scale = 85

        # Put into size bin using ESD
        esd = 2*np.sqrt(region['area']/np.pi) * pixel_scale / 1000 # ESD in mm
        esds.append(esd)
        # print(esd)

        if region['area'] > min_area:
            roi = {}
            #roi['image'] = region['intensity_image']
            roi['image'] = im_in[box[0]:box[2], box[1]:box[3]]
            roi['file_prefix'] = prefix
            min_image = np.min(roi['image'][roi['image'] > 0])
            roi['image'] = 255*(roi['image'].astype('float') - min_image) / np.max(roi['image'])
            roi['image'][roi['image'] < 0] = 0
            roi['image'] = roi['image'].astype('uint8')

            im_out = cv2.rectangle(im_out,
                          (box[1], box[0]),
                          (box[3], box[2]),
                          (0, 255, 0),
                          2
            )

            valid_roi_counter += 1

            if save_rois:
                imageio.imwrite(
                    os.path.join(output_path, prefix + '-'
                        + str(counter) + '-'
                        + str(box[1]) + '-'
                        + str(box[0]) + '-'
                        + str(box[3]) + '-'
                        + str(box[2]) + '-'
                        + str(int(depth_val))
                        + '.png'
                    ),
                    roi['image']
                )

        counter += 1

    h, n = np.histogram(esds, bins=size_classes)

    print(h)

    return im_out, valid_roi_counter, h, len(all_regions)


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 preproc_image_pair(file_prefix, depth_range=[0, 255]):

    # load the total focus image
    tf_image = imageio.imread(file_prefix+'-TotalFocus.png')

    # Step 0 : generate depth mask from 3D image using one of the channels
    depth_img = DepthImage(file_prefix)
    mask = depth_img.get_depth_mask()

    # step 1 : Resize the mask to match the dims of the total focus image
    tf_width = tf_image.shape[1]
    tf_height = tf_image.shape[0]
    mask = cv2.resize(mask, (tf_width, tf_height))

    # step 3 : convert image to 8-bits
    tf_image = (tf_image/256).astype('uint8')

    # step 2 : mask out artifacts in the total focus image
    # the assumption here is that if there is no depth calculated for that pixel
    # it should be excluded from the image.
    masked_tf = mask * tf_image

    masked_tf = (masked_tf).astype('uint8')

    # cv2.imshow('edge', cv2.resize(mask, (1024, 1024)))
    # cv2.waitKey()


    # step 4: mask out the background artifacts around valid depths
    masked_tf[masked_tf < NOISE_FLOOR] = NOISE_FLOOR
    tf_image[tf_image < NOISE_FLOOR] = NOISE_FLOOR

    #lum_mask = masked_tf.copy()
    #lum_mask[lum_mask > 0] = 1

    # These last two step may not be needed

    # step 5 : segment the masked total focus image and fill holes in the segmentation
    edge_image = auto_canny(masked_tf, sigma=0.33)
    kernel = np.ones((9, 9), np.uint8)
    edge_image = cv2.morphologyEx(edge_image, cv2.MORPH_CLOSE, kernel)
    edge_filled = img_fill(edge_image).astype('uint8')
    edge_filled[edge_filled > 0] = 1

    # step 6: mask the depth images with the segmented total focus image
    depth_img.mask(edge_filled)
    #cv2.imshow('depth', cv2.resize(depth_img.depth_images[3], (1024, 1024)))
    #cv2.waitKey()

    return tf_image, masked_tf, depth_img

def make_dir(base_dir, new_dir):
    new_path = os.path.join(base_dir, new_dir)
    if not os.path.exists(new_path):
        print(new_path)
        os.makedirs(new_path)
    return new_path

if __name__=="__main__":

    if len(sys.argv) < 4:
        print('Usage: RaytrixParticleProc [input dir] [pattern] [output dir]')
        exit()

    focus_images = sorted(glob.glob(os.path.join(sys.argv[1], sys.argv[2] + '-TotalFocus.png')))

    output_dir = sys.argv[3]

    if len(focus_images) > 0:
        # create output dirs
        roi_dir = make_dir(output_dir, 'rois')
        stats_dir = make_dir(output_dir, 'stats')
        images_dir = make_dir(output_dir, 'images')
        clouds_dir = make_dir(output_dir, 'clouds')



    stats_file = os.path.join(stats_dir, 'basic_stats.csv')
    total_rois = 0

    for img in focus_images:

        sub_fmt = "{:07d}".format(int(total_rois / 5000)*5000)
        roi_sub_dir = os.path.join(roi_dir, sub_fmt)
        if not os.path.exists(roi_sub_dir):
            make_dir(roi_dir, sub_fmt)

        prefix = img[0:-15]  # remove the '-TotalFocus.png'
        print(prefix)
        tf_image, masked_tf, depth_img = preproc_image_pair(prefix)
        cloud = depth_img.generate_point_cloud()

        if DEBUG:
            cv2.imshow('masked_tf', cv2.resize(masked_tf, (1024, 1024)))
            cv2.imshow('masked depth', depth_img.depth_images[3])
            cv2.waitKey()
        else:
            #imageio.imwrite(prefix+'-masked-tf.jpg', masked_tf)
            im_out, n_rois, h, n_particles = roi_detector(os.path.basename(prefix), roi_sub_dir, tf_image, depth_img, save_rois=True)
            total_rois += n_rois

            with open(stats_file,"a+") as f:
                output = os.path.basename(prefix) + ","
                output += str(n_particles) + ","
                for i in range(len(h)):
                    output += str(h[i]) + ","
                output += str(depth_img.num_particles) + ","
                output += str(depth_img.nnstat)
                f.write(output + "\n")

            imageio.imwrite(os.path.join(images_dir, os.path.basename(prefix) + '-masked-tf.png'), im_out)

            if np.prod(cloud.shape) != 0:
                cloud = PyntCloud(pd.DataFrame(
                    # same arguments that you are passing to visualize_pcl
                    data=cloud, columns=["x", "y", "z", "size"])
                )
                cloud.to_file(os.path.join(clouds_dir, os.path.basename(prefix) + ".ply"))