"""Extract small image containing single seed from larger image."""

import argparse

import numpy as np

from skimage.filters import threshold_otsu
from skimage.transform import hough_circle

from jicbioimage.core.image import Image

from jicbioimage.transform import (
    find_edges_sobel
)

from scipy.ndimage.morphology import (
    binary_erosion, 
    binary_dilation,
    distance_transform_cdt,

)

def find_n_best_hough_circles(radii, hough_res, n):
    """Given the radii and hough accumulators for those radii, find the n
    accumulators with the best circles, returning the centers and radii of
    each of those circles in the form (x1, y1, r1), (x2, y2, r2), ...."""

    n_radii = len(radii)

    max_by_radii = [(np.max(hough_res[r,:,:]), r) for r in range(n_radii)]
    max_by_radii.sort(reverse=True)

    best_scores = max_by_radii[:2]

    def flatten_where_result(where_result):
        return [e[0] for e in where_result]

    circles = []
    for score, index in best_scores:
        x, y = flatten_where_result(np.where(hough_res[index,:,:]==score))
        r = radii[index]
        circles.append((x, y, r))

    return circles

def find_inner_circle_parameters(plane_array, rmin=200, rmax=250):
    """Given a single planar image (e.g. a section from CT data), find the
    locations of the inner of two circles in the image."""

    xdim, ydim = plane_array.shape

    edges = find_edges_sobel(plane_array)

    hough_radii = np.arange(rmin, rmax, 3)
    hough_res = hough_circle(edges, hough_radii)

    # Find the two clearest circles
    c1, c2 = find_n_best_hough_circles(hough_radii, hough_res, 2)

    # Work out which is the inner circle
    r1 = c1[2]
    r2 = c2[2]
    if r1 > r2:
        inner_circle_radius = r2
        cx, cy, r = c2
    else:
        inner_circle_radius = r1
        cx, cy, r = c1

    return cx, cy, r

def strip_outside_circle(input_array, center, radius):
    """Return array generated by setting all values in input_array outside 
    circle of given center and radius to zero."""

    cx, cy = center
    r = radius
    xdim, ydim = input_array.shape

    y, x = np.ogrid[-cx:xdim-cx,-cy:ydim-cy]
    # Small adjustment for aliasing
    r = r - 2
    mask = x*x + y*y >= r*r

    output_array = np.copy(input_array)
    output_array[mask] = 0

    return output_array

def find_single_seed(image_filename, output_filename):

    image = Image.from_file(image_filename)

    w, h = 500, 500
    tube_section = image[1024-w:1024+w,1024-h:1024+h]

    threshold = threshold_otsu(tube_section)

    thresholded = tube_section > threshold

    x, y, r = find_inner_circle_parameters(thresholded, 400, 500)

    # FIXME - think routine is finding outer circle

    stripped = strip_outside_circle(thresholded, (x, y), 300)

    eroded = binary_erosion(stripped, structure=np.ones((10, 10)))

    float_coords = map(np.mean, np.where(eroded > 0))
    ix, iy = map(int, float_coords)

    w, h = 100, 100
    selected = tube_section[ix-w:ix+w,iy-h:iy+h]

    with open(output_filename, 'wb') as f:
        f.write(selected.view(Image).png())

def main():
    parser = argparse.ArgumentParser(__doc__)
    parser.add_argument('image_filename', help="Path to larger image.")
    parser.add_argument('output_filename', help="Image to output.")

    args = parser.parse_args()

    find_single_seed(args.image_filename, args.output_filename)

if __name__ == "__main__":
    main()
