import sys

import neuron_morphology.swc as swc
import neuron_morphology.features.feature_extractor as feature_extractor
from allensdk.internal.core.lims_pipeline_module import PipelineModule


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

def main(jin):
    try:
        swc_file = jin["swc_file"]
        xform = jin["pia_transform"]
        depth = jin["relative_soma_depth"]
    except:
        print("** Unable to find requisite fields in input json");
        raise


    ####################################################################
    # calculate features

    try:
        nrn = swc.read_swc(swc_file)
    except:
        print("** Error reading swc file")
        raise

    try:
        aff = []
        for i in range(12):
            aff.append(xform["tvr_%02d" % i])
        nrn.apply_affine(aff)
    except:
        print("** Error applying affine transform")
        raise

    #try:
    #    # save a copy of affine-corrected file
    #    tmp_swc_file = swc_file[:-4] + "_pia.swc"
    #    nrn.write(tmp_swc_file)
    #except:
    #    # treat this as a soft error and print a warning
    #    print("Note: unable to write copy of affine corrected pia file")

    try:
        features = feature_extractor.MorphologyFeatures(nrn, depth)
        data = {}
        data["axon"] = features.axon
        data["cloud"] = features.axon_cloud
        data["dendrite"] = features.dendrite
        data["basal_dendrite"] = features.basal_dendrite
        data["apical_dendrite"] = features.apical_dendrite
        data["all_neurites"] = features.all_neurites
    except:
        print("** Error calculating morphology features")
        raise

    # make output of new module backwards compatible with previous module
    md = {}
    feat = {}
    feat["number_of_stems"] = data["dendrite"]["num_stems"]
    feat["max_euclidean_distance"] = data["dendrite"]["max_euclidean_distance"]
    feat["max_path_distance"] = data["dendrite"]["max_path_distance"]
    feat["overall_depth"] = data["dendrite"]["depth"]
    feat["total_volume"] = data["dendrite"]["total_volume"]
    feat["average_parent_daughter_ratio"] = data["dendrite"]["mean_parent_daughter_ratio"]
    feat["average_diameter"] = data["dendrite"]["average_diameter"]
    feat["total_length"] = data["dendrite"]["total_length"]
    feat["nodes_over_branches"] = data["dendrite"]["neurites_over_branches"]
    feat["overall_width"] = data["dendrite"]["width"]
    feat["number_of_nodes"] = data["dendrite"]["num_nodes"]
    feat["average_bifurcation_angle_local"] = data["dendrite"]["bifurcation_angle_local"]
    feat["number_of_bifurcations"] = data["dendrite"]["num_bifurcations"]
    feat["average_fragmentation"] = data["dendrite"]["mean_fragmentation"]
    feat["number_of_tips"] = data["dendrite"]["num_tips"]
    feat["average_contraction"] = data["dendrite"]["contraction"]
    feat["average_bifuraction_angle_remote"] = data["dendrite"]["bifurcation_angle_remote"]
    feat["number_of_branches"] = data["dendrite"]["num_branches"]
    feat["total_surface"] = data["dendrite"]["total_surface"]
    feat["max_branch_order"] = data["dendrite"]["max_branch_order"]
    feat["soma_surface"] = data["dendrite"]["soma_surface"]
    feat["overall_height"] = data["dendrite"]["height"]


    md["features"] = feat
    data["morphology_data"] = md

    return data


if __name__=='__main__':
    module = PipelineModule()
    jin = module.input_data()
    jout = main(jin)
    module.write_output_data(jout)
