########################################################################
# library code
import math
import argparse
import sys
import numpy as np
from scipy.spatial.distance import euclidean
import neuron_morphology.swc as swc
import skimage.draw

#import allensdk.core.json_utilities as json


def calculate_centroid(x, y):
    ''' Calculates the center of a polygon, using weighted averages
    of vertex locations
    '''
    assert len(x) == len(y), "Vertex arrays are of incorrect shape"
    tot_len = 0.0
    tot_x = 0.0
    tot_y = 0.0
    for i in range(len(x)):
        x0 = x[i-1]
        y0 = y[i-1]
        x1 = x[i]
        y1 = y[i]
        seg_len = euclidean((x0, y0), (x1, y1))
        tot_len += seg_len
        tot_x += seg_len * x0
        tot_x += seg_len * x1
        tot_y += seg_len * y0
        tot_y += seg_len * y1
    tot_x /= 2.0 * tot_len
    tot_y /= 2.0 * tot_len
    return tot_x, tot_y


def construct_affine(theta):
    tr_rot = [np.cos(theta), np.sin(theta), 0,
              -np.sin(theta), np.cos(theta), 0,
              0, 0, 1,
              0, 0, 0
             ]
    return tr_rot


#def get_pia_wm_rotation_transform(soma_coords, wm_coords, pia_coords, resolution):
#    # get soma position using weighted average of vertices
#    sx, sy = convert_coords_str(soma_coords)
#    soma_x, soma_y = calculate_centroid(sx, sy)
#
#    #pia_proj = project_to_polyline(pia_coords, avg_soma_position)
#    #wm_proj = project_to_polyline(wm_coords, avg_soma_position)
#    px, py, wx, wy = calculate_shortest(soma_x, soma_y, pia_coords, wm_coords)
#    theta = vector_angle((0, 1), np.asarray([px,py]) - np.asarray([wx,wy]))
#    print theta
#
#    depth = euclidean((soma_x, soma_y), (px, py))
#    height = euclidean((wx, wy), (px, py))
#    return theta, resolution*depth, resolution*height, [px, py, wx, wy]


def convert_coords_str(coord_str):
    vals = coord_str.split(',')
    x = np.array(vals[0::2], dtype=float)
    y = np.array(vals[1::2], dtype=float)
    return x, y

def calculate_shortest(soma_x, soma_y, pia, wm):
    """ Calculates shortest distance through a point on the polygon wm
    through the soma coordinates (soma_x, soma_y) and
    through a point on the polygon pia.

    Returns the x,y points in pia and wm that define the endpoints of this
    shortest line.
    """
    pia_xs, pia_ys = convert_coords_str(pia)
    wm_xs, wm_ys = convert_coords_str(wm)
    #################
    # calculate canvas size and make canvas
    width = max(pia_xs)
    width = int(max(width, max(wm_xs)))
    height = max(pia_ys)
    height = int(max(height, max(wm_ys)))

    canvas = np.zeros((height+10, width+10, 3), dtype=np.uint8)
    
    for i in range(1,len(pia_xs)):
        lr, lc = skimage.draw.line(int(pia_ys[i-1]), int(pia_xs[i-1]), int(pia_ys[i]), int(pia_xs[i]))
        canvas[lr,lc,2] = 255

    for i in range(1,len(wm_xs)):
        lr, lc = skimage.draw.line(int(wm_ys[i-1]), int(wm_xs[i-1]), int(wm_ys[i]), int(wm_xs[i]))
        canvas[lr,lc,0] = 255

    # get points in white matter trace
    wp_y, wp_x = np.nonzero(canvas[:,:,0])

    ##################
    # draw an extended line from each wm pix through the soma
    # (there are usually less WM pix than pia pix, so this should be 
    #   faster than iterating through pia pix)
    # make array of blue (pia) channel only to
    pia = canvas[:,:,2]
    # draw line from each wm pix through the soma and into infinity
    #   (line terminates if/when it intersects with pia trace)
    min_dist = None
    min_coord = None    # stores [ pia_x, pia_y, wm_x, wm_y ]
    for i in range(len(wp_x)):
        x0 = wp_x[i]
        y0 = wp_y[i]
        x1 = soma_x
        y1 = soma_y
        #############################################
        # adapted from Bresenham's line algorithm, from rosetacode
        dx = abs(x1 - x0)
        dy = abs(y1 - y0)
        x, y = x0, y0
        sx = -1 if x0 > x1 else 1
        sy = -1 if y0 > y1 else 1
        if dx > dy:
            err = dx / 2.0
            while x >= 0 and x < width and y >= 0 and y < height:
                if pia[y,x] > 0:
                    dist = euclidean((x0, y0), (x, y))
                    if min_dist is None or min_dist > dist:
                        min_dist = dist
                        min_coord = [x, y, x0, y0]
                    break
                err -= dy
                if err < 0:
                    y += sy
                    err += dx
                x += sx
        else:
            err = dy / 2.0
            while x >= 0 and x < width and y >= 0 and y < height:
                if pia[y,x] > 0:
                    dist = euclidean((x0, y0), (x, y))
                    if min_dist is None or min_dist > dist:
                        min_dist = dist
                        min_coord = [x, y, x0, y0]
                    break
                err -= dx
                if err < 0:
                    x += sx
                    err += dy
                y += sy

    if min_dist is None:
        print("Unable to connect pia to WM through soma")
        px = None
        py = None
        wx = None
        wy = None
    else:
        px = min_coord[0]
        py = min_coord[1]
        wx = min_coord[2]
        wy = min_coord[3]
    return px, py, wx, wy


def dist_proj_point_lineseg(p, q1, q2):
    # based on c code from http://stackoverflow.com/questions/849211/shortest-distance-between-a-point-and-a-line-segment
    l2 = euclidean(q1, q2) ** 2
    if l2 == 0:
        return euclidean(p, q1) # q1 == q2 case
    t = max(0, min(1, np.dot(p - q1, q2 - q1) / l2))
    proj = q1 + t * (q2 - q1)
    return euclidean(p, proj), proj


def project_to_polyline(boundary, soma):
    x, y = convert_coords_str(boundary)
    points = zip(x, y)
    dists_projs = [dist_proj_point_lineseg(soma, np.array(q1), np.array(q2))
                 for q1, q2 in zip(points[:-1], points[1:])]
    min_idx = np.argmin(np.array([d[0] for d in dists_projs]))
    return dists_projs[min_idx][1]


def vector_angle(v1, v2):
    return np.arctan2(v2[1], v2[0]) - np.arctan2(v1[1], v1[0])


########################################################################
# pipeline code
import allensdk.internal.core.swc as swc
from allensdk.internal.core.lims_pipeline_module import PipelineModule


def main(jin):
    # per IT-14567, blockface analysis is no longer required
    #########################################################################
    ## analyze blockface image
    #try:
    #    soma = jin["blockface"]["Soma"]["path"]
    #    pia = jin["blockface"]["Pia"]["path"]
    #    wm = jin["blockface"]["White Matter"]["path"]
    #    res = float(jin["blockface"]["Pia"]["resolution"])
    #except:
    #    print("** Error -- missing requisite blockface field(s) in input json")
    #    raise
    #
    ## get soma position using weighted average of vertices
    #try:
    #    sx, sy = convert_coords_str(soma)
    #    soma_x, soma_y = calculate_centroid(sx, sy)
    #except:
    #    print("** Error -- unable to calculate soma information (blockface)")
    #    raise
    #
    ## calculate shortest path
    #try:
    #    px, py, wx, wy = calculate_shortest(soma_x, soma_y, pia, wm)
    ## calculate theta and affine
    #    theta = vector_angle((0, 1), np.asarray([px,py]) - np.asarray([wx,wy]))
    ## calculate soma depth and cortical thickness
    #    depth = res * euclidean((soma_x, soma_y), (px, py))
    #    blk_thickness = res * euclidean((wx, wy), (px, py))
    #except:
    #    print("** Error calculating shortest path (blockface)")
    #    raise
    #
    #blockface = {}
    #blockface["pia_intersect"] = [ px, py ]
    #blockface["wm_intersect"] = [ wx, wy ]
    #blockface["soma_center"] = [ soma_x, soma_y ]
    #blockface["soma_depth_um"] = depth
    #try:
    #    blockface["soma_depth_relative"] = depth / blk_thickness
    #except:
    #    blockface["soma_depth_relative"] = -1.0    # NaN is not friendly to ruby
    #blockface["cort_thickness_um"] = blk_thickness
    #blockface["theta"] = theta
    
    ########################################################################
    # analyze primary (20x) image
    try:
        soma = jin["primary"]["Soma"]["path"]
        pia = jin["primary"]["Pia"]["path"]
        wm = jin["primary"]["White Matter"]["path"]
        res = float(jin["primary"]["Pia"]["resolution"])
    except:
        print("** Error -- missing requisite primary (20x) field(s) in input json")
        raise
    
    # get soma position using weighted average of vertices
    try:
        sx, sy = convert_coords_str(soma)
        soma_x, soma_y = calculate_centroid(sx, sy)
    except:
        print("** Error -- unable to calculate soma information (primary)")
        raise
    try:
        # calculate shortest path
        px, py, wx, wy = calculate_shortest(soma_x, soma_y, pia, wm)
        # calculate theta and affine
        theta = vector_angle((0, 1), np.asarray([px,py]) - np.asarray([wx,wy]))
        tr_rot = construct_affine(theta)
        inv_tr_rot = construct_affine(-theta)
        # calculate soma depth and cortical thickness
        depth = res * euclidean((soma_x, soma_y), (px, py))
        raw_thickness = res * euclidean((wx, wy), (px, py))
    except:
        print("** Error calculating shortest path (primary)")
        raise
    
    
    primary = {}
    primary["pia_intersect"] = [ px, py ]
    primary["wm_intersect"] = [ wx, wy ]
    primary["soma_center"] = [ soma_x, soma_y ]
    primary["soma_depth_um"] = depth
    try:
        primary["soma_depth_relative"] = depth / raw_thickness
    except:
        primary["soma_depth_relative"] = -1.0   # NaN is not friendly to ruby
    primary["cort_thickness_um"] = raw_thickness
    primary["theta"] = theta
    
    try:
        scale = raw_thickness / blk_thickness
    except:
        scale = -1.0    # NaN is not ruby-friendly
    
    soma_coords_avail = False
    if "swc_file" in jin:
        # if SWC file available, extract soma position from it
        try:
            nrn = swc.read_swc(jin["swc_file"])
            root = nrn.soma_root()
            soma_x = root.x
            soma_y = root.y
            soma_z = root.z
            soma_coords_avail = True
        except:
            # treat this as a fatal error -- if SWC was specified then
            #   it should be used
            print("**** Error reading SWC file '%s'" % jin["swc_file"])
            raise

    if not soma_coords_avail:
        # hope that the 63x data is available. As of May 2017, this seems
        #   to no longer be supplied to the module, but just in case...
        # IT-14567 continue if 63x data not available
        try:
            info = jin["soma_63x"]
            sx, sy = convert_coords_str(info["path_63x"])
            soma_x, soma_y = calculate_centroid(sx, sy)
            soma_x *= float(info["resolution"])
            soma_y *= float(info["resolution"])
            soma_z = float(info["idx"]) * float(info["thickness"])
            soma_coords_avail = True
        except:
            print("** Error reading soma 63x info from input json")
            print("** Translation component of affine matrix is invalid **")

    if not soma_coords_avail:
        raise Exception("** Error: Unable to construct translation component of affine")
    
    try:
        # apply affine rotation to soma position
        translate_x = soma_x*tr_rot[0] + soma_y*tr_rot[1] + soma_z*tr_rot[2]
        translate_y = soma_x*tr_rot[3] + soma_y*tr_rot[4] + soma_z*tr_rot[5]
        translate_z = soma_x*tr_rot[6] + soma_y*tr_rot[7] + soma_z*tr_rot[8]
        # apply translation vector to transform
        tr_rot[ 9] = -translate_x
        tr_rot[10] = -translate_y - depth
        tr_rot[11] = -translate_z
    except:
        print("** Error calculating affine tranform (math fault?)")
        raise
    soma_x = -translate_x
    soma_y = -translate_y - depth
    soma_z = -translate_z
    
    # apply affine rotation to soma position
    translate_x = soma_x*inv_tr_rot[0] + soma_y*inv_tr_rot[1] + soma_z*inv_tr_rot[2]
    translate_y = soma_x*inv_tr_rot[3] + soma_y*inv_tr_rot[4] + soma_z*inv_tr_rot[5]
    translate_z = soma_x*inv_tr_rot[6] + soma_y*inv_tr_rot[7] + soma_z*inv_tr_rot[8]
    
    inv_tr_rot[ 9] = -translate_x
    inv_tr_rot[10] = -translate_y
    inv_tr_rot[11] = -translate_z
    
    try:
        # upright transform. based on rotation in 20x image
        upright = {}
        for i in range(12):
            upright["tvr_%02d" % i] = tr_rot[i]
            upright["trv_%02d" % i] = inv_tr_rot[i]
        jout = {}
        jout["primary"] = primary
        #jout["blockface"] = blockface  # per IT-14567, disable blockface
        jout["upright"] = upright
        alignment = {}
        alignment["scale"] = scale
        alignment["rotate_x"] = theta
        alignment["rotate_y"] = theta
        alignment["rotate_z"] = 0.0
        alignment["scale_x"] = 1.0
        alignment["scale_y"] = 1.0
        alignment["scale_z"] = 1.0
        alignment["skew_x"] = 0.0
        alignment["skew_y"] = 0.0
        alignment["skew_z"] = 0.0
        jout["alignment"] = alignment
    except:
        print("** Internal error **")
        raise
    
    return jout

    #
    # test transform -- bar.swc should match the source file
    #print("source swc: " + jin["swc_file"])
    #print tr_rot
    #morph2 = swc.read_swc(jin["swc_file"])
    #morph2.apply_affine(tr_rot)
    #morph2.save("foo.swc")
    #morph3 = swc.read_swc("foo.swc")
    #morph3.apply_affine(inv_tr_rot)
    #morph3.save("bar.swc")

if __name__ == "__main__":
    # read module input. PipelineModule object automatically parses the 
    #   command line to pull out input.json and output.json file names
    module = PipelineModule()
    jin = module.input_data()   # loads input.json
    jout = main(jin)
    module.write_output_data(jout)  # writes output.json

