#!/usr/bin/python
import os, sys
import allensdk.internal.core.swc as swc


def resave_swc(orig_swc, new_file):
    """ Reads SWC file into AllenSDK Morphology object and resaves
    it. This can fix some problems in an SWC file that may disrupt
    other software tools reading the file (e.g., NEURON)

    Parameters
    ----------
        orig_swc: string
        Name of SWC file to read

        new_file: string
        Name of output SWC file
    """
    try:
        morphology = swc.read_swc(orig_swc)
    except:
        print("Failed to read SWC file '%s'" % orig_swc)
        raise
    try:
        morphology.save(new_file)
    except:
        print("Failed to save SWC file '%s'" % new_file)


class TestNode(object):
    def __init__(self, n, t, x, y, z, r, pn):
        # these values correspond to columns in an SWC file
        self.n = n
        self.t = t
        self.x = x
        self.y = y
        self.z = z
        self.r = r
        self.pn = pn
        self.children = []  # IDs of child nodes

    def __str__(self):
        """ create string with node information in succinct, 
        single-line form """
        return "%d %d %.4f %.4f %.4f %.4f %d %s" % (self.n, self.t, self.x, self.y, self.z, self.r, self.pn, str(self.children))



def validate_swc(swc_file):
    """ 
    Tests SWC files for compatibility with AllenSDK

    To be compatible with NEURON, SWC files must have the following properties:
        1) a single root node with parent ID '-1'
        2) sequentially increasing ID numbers
        3) immediate children of the soma cannot branch

    To be compatible with feature analysis, SWC files can only have node
    types in the range 1-4:
        1 = soma
        2 = axon
        3 = [basal] dendrite
        4 = apical dendrite
    """
    success = True  # be optimistic

    # see if SWC file is readable by internal tools
    print("Validating " + swc_file)
    try:
        morphology = swc.read_swc(swc_file)
    except:
        print("Fatal error reading SWC file")
        return False

    for node in morphology.node_list:
        if node.t < 1 or node.t > 4:
            print("Expecting type between 1 and 4, but found %d" % node.t)
            print("File has unrecognized node type(s)")
            print("----------------------------------")
            success = False
            break

    # make sure all dendrite nodes are in tree 0
    # this is because modeling requires a full dendrite morphology
    for node in morphology.node_list:
        if (node.t == 3 or node.t == 4) and node.tree_id != 0:
            print("Dendrite node(s) exist in disconnected trees")
            print("This breaks an SDK modeling requirement")
            print("----------------------------------")
            success = False
            break

    # if we've made it here, file is OK for using Morphology class, and 
    #   should be valid with internal processing. It may also be able
    #   to be convertable for NEURON use by resaving it

    nodes = []
    node_table = [] # lookup table by node num
    line_num = 1
    try:
        with open(swc_file, "r") as f:
            for line in f:
                # remove comments
                if line.lstrip().startswith('#'):
                    continue
                # read values. expected SWC format is:
                #   ID, type, x, y, z, rad, parent
                # x, y, z and rad are floats. the others are ints
                toks = line.split()
                vals = TestNode(
                        n =  int(toks[0]),
                        t =  int(toks[1]),
                        x =  float(toks[2]),
                        y =  float(toks[3]),
                        z =  float(toks[4]),
                        r =  float(toks[5]),
                        pn = int(toks[6].rstrip()),
                    )
                # store this node
                while len(nodes) <= vals.n:
                    nodes.append(None)
                nodes[vals.n] = vals
                #nodes.append(vals)
                #
                if vals.n < 0:
                    print("Negative node ID not allowed")
                    print("Node: " + str(vals))
                    return False
                while vals.n >= len(node_table):
                    node_table.append(None)
                node_table[vals.n] = vals
                # increment line number (used for error reporting only)
                line_num += 1
    except:
        err = "File not recognized as valid SWC file.\n"
        err += "Problem parsing line %d\n" % line_num
        if line is not None:
            err += "Content: '%s'\n" % line
        raise IOError(err)

    try:
        for node in nodes:
            if node is None:
                continue
            par = None
            if node.pn >= 0:
                par = node_table[node.pn]
                par.children.append(node.n)
    except:
        print("Error reading SWC file -- fail to link child to parent")
        print("Node:    %s" % str(node))
        print("----------------------------------")
        success = False

    # verify presence and number of soma and root nodes
    num_soma_nodes = sum([ int(c is not None and c.t == 1) for c in nodes ])
    if num_soma_nodes == 0:
        print("SWC must have at least one soma node.  Found: %d" % num_soma_nodes)
        print("----------------------------------")
        success = False
    elif num_soma_nodes > 1:
        print("Warning: File has multiple soma nodes. This can interfere with feature analysis in some external software (e.g., vaa3d)")
        print("----------------------------------")

    num_root_nodes = sum([ int(c is not None and c.pn == -1) for c in nodes ])
    # case of no root nodes covered by rule below that ID of child must
    #   be greater than that of parent
    if num_root_nodes > 1:
        print("Warning: File has multiple root nodes. This can interfere with feature analysis in some external software (e.g., vaa3d)")
        print("----------------------------------")

    # get a list of all of the ids, make sure they are unique while we're at it
    all_ids = set()
    for node in nodes:
        if node is None:
            continue
        iid = int(node.n)
        if iid in all_ids:
            print("Node ID %s is not unique." % node.n)
            print("----------------------------------")
            success = False
            break
        pid = int(node.pn)
        if iid < pid:
            print("Node (%d) has a smaller ID that its parent (%d)" % (iid, pid))
            print("----------------------------------")
            success = False
            break
        all_ids.add(iid)
        
    # make sure that first root node is soma
    for n in nodes:
        if n is not None:
            root = n
            break
    #root = nodes[0]
    if root.t != 1:
        # see if soma has a root
        if sum([int(c is not None and c.t == 2 and c.pn == -1) for c in nodes]) == 0:
            print("No soma root found in file")
            print("----------------------------------")
            success = False
        print("First root node is not soma")
        print("This should be fixable by calling resave_swc() on the file if there is a soma root in the file")
        print("----------------------------------")
        success = False

    # verify that children of the root have max one child
    for root_child_id in root.children:
        root_child = nodes[root_child_id]
        num_grand_children = len(root_child.children)
        if num_grand_children > 1:
            print("Child of root (%s) has more than one child (%d)" % ( root_child_id, num_grand_children ))
            print("----------------------------------")
            success = False

    # sort the ids and make sure there are no gaps
    sorted_ids = sorted(all_ids)
    for i in xrange(1, len(sorted_ids)):
        if sorted_ids[i] - sorted_ids[i-1] != 1:
            print("Node IDs are not sequential")
            print("This can be fixed by calling resave_swc() on the file")
            print("----------------------------------")
            success = False
    return success
    
    
    
def main():
    argc = len(sys.argv)
    if argc < 1:
        print("usage: python %s <swc_file> [<swc_file ...]")
        print("")
        print("Validate an SWC file for use with NEURON")
        sys.exit(1)
    try:
        for i in range(1, argc):
            if validate_swc(sys.argv[i]) == True:
                print "    PASS"
            else:
                print "    FAIL"
                exit(1)
    except Exception as e:
        print "    FAIL"
        print str(e)
        exit(1)
if __name__ == "__main__": main()
