#!/usr/bin/python
import json
import math
import cv2
import numpy as np
import sys
import psycopg2
import psycopg2.extras
import allensdk.core.json_utilities as json
from neuron_morphology import swc
from allensdk.internal.core.lims_pipeline_module import PipelineModule

#from surrogate_strategy import prep_json

# TODO update run_python.sh to include this path
twok_dir = "/shared/bioapps/itk/itk_shared/jp2/build"
print("WARNING: adding directory to PYTHONPATH")
print("=>  %s" % twok_dir)
sys.path.append(twok_dir)
import jpeg_twok

########################################################################
# helper functions

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 = math.sqrt((x1-x0)*(x1-x0) + (y1-y0)*(y1-y0))
        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 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


color_table = []
color_table.append((255,  77,  77))
color_table.append((102, 102, 255))
color_table.append(( 25, 255,  25))
color_table.append((177, 166, 255))
color_table.append(( 46, 230, 230))
color_table.append((255,  77, 255))
color_table.append((128, 230,  46))
color_table.append((255, 166,  77))
color_table.append((179, 179, 179))
color_table.append(( 77, 255, 166))
color_table.append((229, 229,  46))
color_table.append((255,  51, 153))
color_table.append((166,  77, 255))
color_table.append((151, 166,  86))

color_table.append((153,   0,   0))
color_table.append((  0,  77, 153))
color_table.append((153, 153,   0))
color_table.append(( 77, 153,   0))
color_table.append((  0, 153, 153))
color_table.append(( 13, 128,  13))
color_table.append((153,  77,   0))
color_table.append((  0,   0, 153))
color_table.append((153,   0, 153))
color_table.append((  0, 179,  89))
color_table.append((102, 102, 102))
color_table.append(( 77,   0, 153))
color_table.append((153,   0,  77))
color_table.append(( 78,  89,  30))

def color_by_index(i):
    global color_table
    #
    return color_table[i % len(color_table)]

def draw_morphology(nrn, img, somax, somay, color_by_layer=False):
    global LINE_WIDTH, resolution
    #
    soma_col = (0, 0, 0)
    axon_col = (70, 130, 180)
    dend_col = (178, 34, 34)
    apical_col = (255, 127, 80)
    for c in nrn.compartment_list:
        x0 = int(c.node1.x / resolution + somax)
        x1 = int(c.node2.x / resolution + somax)
        y0 = int(c.node1.y / resolution + somay)
        y1 = int(c.node2.y / resolution + somay)
        if color_by_layer:
            color = color_table[c.node1.layer_num]
        else:
            color = soma_col
            if c.node2.t == 2:
                color = axon_col
            elif c.node2.t == 3:
                color = dend_col
            elif c.node2.t == 4:
                color = apical_col
        cv2.line(img, (x0,y0), (x1,y1), color, LINE_WIDTH)

def write_svg(svgname, jin, nrn):
    resolution = jin["resolution"]
    dx = jin["soma"]["position"][0] - nrn.soma_root().x / resolution
    dy = jin["soma"]["position"][1] - nrn.soma_root().y / resolution
    with open(svgname, "w") as f:
        #f.write('<?xml version="1.0" encoding="UTF-8" ?>\n')
        f.write('<svg xmlns="http://www.w3.org/2000/svg" version="1.1">\n')
        # soma
        soma = jin["soma"]["path"][0]
        coords = soma.split(',')
        f.write('  <polyline points="')
        for i in range(0, len(coords), 2):
            f.write('%f,%f ' % (float(coords[i]), float(coords[i+1])))
        f.write('" stroke="black" stoke-width="2" fill="none" />\n')
        for layer in jin["layers"]:
            f.write('  <polyline points="')
            path = layer["path"]
            coords = path.split(',')
            for i in range(0, len(coords), 2):
                f.write('%f,%f ' % (float(coords[i]), float(coords[i+1])))
            f.write('" stroke="black" stoke-width="2" fill="none" />\n')
        for c in nrn.compartment_list:
            try:
                color = color_table[c.node1.layer_num]
            except:
                color = (45, 67, 89)
            x0 = int(c.node1.x / resolution + dx)
            x1 = int(c.node2.x / resolution + dx)
            y0 = int(c.node1.y / resolution + dy)
            y1 = int(c.node2.y / resolution + dy)
            f.write('  <line x1="%f" y1="%f" x2="%f" y2="%f" style="stroke:rgb(%d,%d,%d)" />\n' % (x0, y0, x1, y1, color[0], color[1], color[2]))
        f.write('</svg>\n')

########################################################################
########################################################################
#
# global values, shared between functions
resolution = None   # microns-per-pixel, from input.json
LINE_WIDTH = 1  # default pen width
DOWNSAMPLE_STEPS = 1    # default image pyramid level
# 
def main(jin):
    global resolution, LINE_WIDTH, DOWNSAMPLE_STEPS
    jout = {}
    spec_id = jin["specimen_id"]

    ####################################################################
    if "line_width" in jin:
        LINE_WIDTH = int(jin["line_width"])

    # according to Staci, accuracy of 20x trace is approximately to the
    #   level of a cell soma (8-10um), which is approx 22 pixels
    # downsampling by 2 won't significantly alter the accuracy of
    #   the annotations and is well within the stated margin of error. this
    #   lends itself to more efficient processing, and better thumbnails
    if "downsample_steps" in jin:
        DOWNSAMPLE_STEPS = int(jin["downsample_steps"])

    ############################################
    # derived constants
    RADS = []
    RADS.append(29)
    RADS.append(15)
    RADS.append(7)
    RADS.append(3)
    DOWNSAMPLE = 1
    for i in range(DOWNSAMPLE_STEPS):
        DOWNSAMPLE *= 2
    GAUS_RAD = RADS[DOWNSAMPLE_STEPS]

    ####################################################################

    # calculate soma position and store in jin structure
    soma_res = jin["soma"]["path"]
    soma_path = soma_res[0].split(',')
    soma_x = np.array(soma_path[0::2], dtype=float)
    soma_y = np.array(soma_path[1::2], dtype=float)
    soma_path = []
    for i in range(len(soma_x)):
        soma_path.append([soma_x[i],soma_y[i]])
    soma_path = np.array(soma_path, np.int32)
    soma_pos = calculate_centroid(soma_x, soma_y)
    jin["soma"]["position"] = soma_pos

    resolution = jin["resolution"]
    swc_name = jin["storage_directory"] + jin["swc_file"]

    ###########################
    # before doing the heavy work, load external objects to make sure they're
    #   available

    # read morphology
    morph = swc.read_swc(swc_name)

    # read 20x
    fname_20x = jin["20x"]["img_path"]
    image_20x = jpeg_twok.read(fname_20x, reduction_factor=DOWNSAMPLE_STEPS)
    abs_width = image_20x.shape[1]
    abs_height = image_20x.shape[0]
    print("20x image size %dx%d at pyramid level %d" % (abs_width, abs_height, DOWNSAMPLE_STEPS))

    # get soma position in 20x pixel space, at present downsample level
    # resolution converts soma coords in microns to pixels
    #   (resolution is microns / pixel)
    resolution *= DOWNSAMPLE    # divide pixels by X means mult res by same
    print("Image resolution (microns/pixel): %f" % resolution)
    dx = jin["soma"]["position"][0] / DOWNSAMPLE - morph.soma_root().x / resolution
    dy = jin["soma"]["position"][1] / DOWNSAMPLE - morph.soma_root().y / resolution
    dx = int(dx)
    dy = int(dy)

    ##############################
    # no point in processing the entire image, as only a small part 
    #   is relevant
    # select min/max values for x,y of all polygons. restrict analysis 
    #   to there
    min_x = 1e10
    min_y = 1e10
    max_x = 0
    max_y = 0
    layers = jin["layers"]
    for layer in layers:
        path_array = np.array(layer["path"].split(','))
        x = np.array(path_array[0::2], dtype=float)
        y = np.array(path_array[1::2], dtype=float)
        min_x = min(min_x, x.min())
        min_y = min(min_y, y.min())
        max_x = max(max_x, x.max())
        max_y = max(max_y, y.max())

    min_x /= DOWNSAMPLE
    min_y /= DOWNSAMPLE
    max_x /= DOWNSAMPLE
    max_y /= DOWNSAMPLE

    # add a border around polygons to provide context
    BORDER = 200
    BORDER /= DOWNSAMPLE
    TOP = min_y - BORDER
    LEFT = min_x - BORDER
    RIGHT = max_x + BORDER
    BOTTOM = max_y + BORDER

    # make sure border doesn't extend beyond image limits
    TOP = max(TOP, 0)
    LEFT = max(LEFT, 0)
    RIGHT = min(RIGHT, abs_width-1)
    BOTTOM = min(BOTTOM, abs_height-1)
    WIDTH = int(RIGHT - LEFT)
    HEIGHT = int(BOTTOM - TOP)

    # adjust soma location for top and left of visible image area
    dx -= LEFT
    dy -= TOP
    #print "inset soma position", dx, dy


    # make frame for each polygon and blur. blur radious should be
    #   approx the size of largest gap or overlap between polygons. this
    #   is for estimating which polygon each point is a best fit in
    layers = jin["layers"]
    for layer in layers:
        path_array = np.array(layer["path"].split(','))
        x = np.array(path_array[0::2], dtype=float)
        x /= DOWNSAMPLE
        x -= LEFT
        y = np.array(path_array[1::2], dtype=float)
        y /= DOWNSAMPLE
        y -= TOP
        #print layer["label"]
        #print x.min(), y.min()
        #print x.max(), y.max()
        xy = []
        for i in range(len(x)):
            xy.append([x[i],y[i]])
        raw_frame = np.zeros((HEIGHT, WIDTH))
        path = np.array(xy)
        cv2.fillPoly(raw_frame, np.int32([path]), 255)
        frame = cv2.blur(raw_frame, (GAUS_RAD, GAUS_RAD))
        layer["frame"] = frame      # blurred polygon
        layer["raw_frame"] = raw_frame  # raw polygon

    # collapse all polys into single array, with value at each position 
    #   corresponding to the index of the polygon that the pixel falls 
    #   into, or -1 if there's no match
    master = np.zeros((HEIGHT, WIDTH, 3))
    master_idx = np.zeros((HEIGHT, WIDTH), dtype=int)
    master_idx -= 1
    for y in range(HEIGHT):
        for x in range(WIDTH):
            peak = 0
            idx = -1
            for i in range(len(layers)):
                frame = layers[i]["frame"]
                val = frame[y][x]
                if val > peak:
                    peak = val
                    idx = i
            if idx >= 0:
                master[y][x] = color_by_index(idx)
                master_idx[y][x] = idx

    #################################################
    # draw standard morphology on colored layers
    draw_morphology(morph, master, int(dx), int(dy))
    outfile = "layer_%d.png" % spec_id
    print "saving " + outfile
    cv2.imwrite(outfile, master)
    jout["morph_layers"] = outfile

    #################################################
    # draw standard morphology on 20x image
    img = image_20x[TOP:BOTTOM,LEFT:RIGHT]
    print img.shape
    draw_morphology(morph, img, int(dx), int(dy))
    outfile = "blockface_%d.png" % spec_id
    print "saving " + outfile
    cv2.imwrite(outfile, img)
    jout["morph_20x"] = outfile

    #################################################
    # associate SWC nodes with morphology layers
    jout["reconstruction_id"] = jin["reconstruction_id"]
    reconstruction = {}
    errs = 0
    for n in morph.node_list:
        x = dx + int(n.x / resolution)
        y = dy + int(n.y / resolution)
        desc = n.to_dict()
        idx = -1
        try:
            idx = master_idx[y][x]
        except:
            errs += 1
        if idx >= 0:
            desc["label"] = layers[idx]["label"]
            n.layer_num = idx
        else:
            desc["label"] = "unknown"
            n.layer_num = -1
        reconstruction[n.n] = desc
    if errs > 0:
        raise Exception("Unable to map %d nodes to a cortical layer" % errs)
    jout["reconstruction"] = reconstruction

    #################################################
    # draw layer & morphology SVG
    outfile = "outline_%d.svg" % spec_id
    print "saving " + outfile
    jout["outline_svg"] = outfile
    write_svg(outfile, jin, morph)
    
    #################################################
    # draw layer-colored morphology on 20x image
    img = image_20x[TOP:BOTTOM,LEFT:RIGHT]
    draw_morphology(morph, img, int(dx), int(dy), True)
    outfile = "layered_blockface_%d.png" % spec_id
    print "saving " + outfile
    cv2.imwrite(outfile, img)
    jout["colored_morph_20x"] = outfile

    #################################################
    # draw layer-colored morphology on empty polygons
    img = np.zeros((HEIGHT, WIDTH, 3))
    layers = jin["layers"]
    for layer in layers:
        path_array = np.array(layer["path"].split(','))
        x = np.array(path_array[0::2], dtype=float)
        x /= DOWNSAMPLE
        x -= LEFT
        y = np.array(path_array[1::2], dtype=float)
        y /= DOWNSAMPLE
        y -= TOP
        for i in range(1,len(x)):
            cv2.line(img, (int(x[i-1]),int(y[i-1])), (int(x[i]),int(y[i])), (255, 255, 255), 1)
        cv2.line(img, (int(x[-1]),int(y[-1])), (int(x[0]),int(y[0])), (255, 255, 255), 1)
    draw_morphology(morph, img, int(dx), int(dy), True)
    outfile = "outline_%d.png" % spec_id
    print "saving " + outfile
    cv2.imwrite(outfile, img)
    jout["colored_morph_poly"] = outfile

    return jout


# mouse
#ims_id = "489909914"
#spec_id = 488679042

#ims_id = "491762612"
#spec_id = 490387590

# human
#ims_id = "487992082"
#spec_id = 488386504

#ims_id = "488759189"
#spec_id = 488418027

#spec_id = 528015670

if __name__ == "__main__":
    module = PipelineModule()
    jin = module.input_data()   # loads input.json
    # "get" input json
    #jin = prep_json(spec_id)
    #json.write("in_%d.json" % spec_id, jin)
    jout = main(jin)
    module.write_output_data(jout)  # writes output.json
    #json.write("out_%d.json" % spec_id, jout)

