# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4
import errno
import json
import matplotlib.image
import numpy
import os
import scipy.ndimage
import scipy.spatial
import skimage.filter
import skimage.feature
import skimage.io
import skimage.measure
import skimage.morphology
import sys

random_numbers = numpy.random.rand(2**12, 3)
random_numbers[0] = numpy.zeros(3)
rg_numbers = numpy.copy(random_numbers)
rg_numbers[:,2] = 0
r_cmap = matplotlib.colors.ListedColormap(random_numbers)
rg_cmap = matplotlib.colors.ListedColormap(rg_numbers)

def open_lsm(filename):
    image = skimage.io.MultiImage(filename)
    return [image[i] for i in xrange(0, len(image), 2)]

def segment_cells(image , min_distance=4, threshold=None):
    #smoothed = skimage.filter.gaussian_filter(stack, 1)
    #smoothed = skimage.filter.rank.enhance_contrast(smoothed, skimage.morphology.disk(5))
    threshold = threshold if threshold is not None else skimage.filter.threshold_otsu(image)
    binary = image > threshold
    distance = scipy.ndimage.distance_transform_edt(binary)
    local_maxima = skimage.feature.peak_local_max(distance,
            min_distance = min_distance, indices = False)
    markers, n_markers = scipy.ndimage.label(local_maxima)
    segmented = skimage.morphology.watershed(-distance, markers, mask = binary)
    print "Threshold: {0} Count: {1}".format(threshold, numpy.max(segmented))
    return segmented

def convex_hull(coordinates):
    if len(coordinates) < 4:
        return coordinates
    if len(numpy.unique(coordinates[..., 0])) < 2:
        return coordinates
    return coordinates[scipy.spatial.ConvexHull(coordinates, qhull_options="QJ").vertices]

def main(argv):
    filename = argv[0]
    min_distance = int(argv[1]) if len(argv) > 1 else 4
    threshold = int(argv[2]) if len(argv) > 2 else None
    dir_name = "out{0}".format(os.getpid())
    try:
        os.mkdir(dir_name)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(dir_name):
            pass
        else: raise
    if filename.endswith("lsm"):
        stack = open_lsm(filename)
    else:
        stack = skimage.io.MultiImage(filename)
    print len(stack)
    stack_file = open("{0}/stack.txt".format(dir_name), "w")
    for i, image in enumerate(stack):
        contours, segmented, contour_map = contour_stack(image, i, min_distance, threshold)
        with open("{0}/{1}.json".format(dir_name, i), "w") as f:
            json.dump(contours, f)
        skimage.io.imsave("{0}/{1}.png".format(dir_name, i), image)
        comp = rg_cmap(segmented)
        comp[...,2] = skimage.img_as_float(image[...,1])
        matplotlib.image.imsave("{0}/{1}_seg.png".format(dir_name, i), segmented, cmap=r_cmap)
        matplotlib.image.imsave("{0}/{1}_con.png".format(dir_name, i), contour_map, cmap="gray")
        matplotlib.image.imsave("{0}/{1}_comp.png".format(dir_name, i), comp)
        stack_file.write("{0}/{1}_comp.png\n".format(dir_name,i))
    with open("{0}/feed.json".format(dir_name), "w") as feed_f:
        json.dump({'layer_ids': range(len(stack))}, feed_f)

def contour_stack(image, layer, min_distance=4, threshold=None):
    segmented = segment_cells(image[...,0], min_distance, threshold)
    shape = numpy.shape(segmented)
    contour_map = numpy.zeros(numpy.shape(segmented), dtype = bool)
    seg_props_red = skimage.measure.regionprops(segmented, intensity_image=image[...,0])
    seg_props_green = skimage.measure.regionprops(segmented, intensity_image=image[...,1])
    seg_props_blue = skimage.measure.regionprops(segmented, intensity_image=image[...,2])
    contours = []
    for i in xrange(len(seg_props_red)):
        contour = convex_hull(seg_props_blue[i].coords)
        contour_map[zip(*contour)] = True
        # Swap coordinates first to make (A, B) -> (B, A)
        contour = numpy.array((contour[...,1], contour[...,0])).T
        properties = {
            'centroid': list(seg_props_blue[i].centroid)[::-1],
            'color_avg': {
                'b': seg_props_blue[i].mean_intensity,
                'g': seg_props_green[i].mean_intensity,
                'r': seg_props_red[i].mean_intensity,
            },
            'color_std': {
                'b': numpy.std(image[...,2][zip(*seg_props_blue[i].coords)]),
                'g': numpy.std(image[...,1][zip(*seg_props_blue[i].coords)]),
                'r': numpy.std(image[...,0][zip(*seg_props_blue[i].coords)])
            },
            'id': seg_props_blue[i].label,
            'pixel_count': int(seg_props_blue[i].area),
            'points': contour.tolist()
        }
        contours.append(properties)
    return {'contours': contours, "id": layer, "imageFilename":"{0}.png".format(layer)}, segmented, contour_map

if __name__ == "__main__":
    main(sys.argv[1:])