# Allen Institute Software License - This software license is the 2-clause BSD
# license plus a third clause that prohibits redistribution for commercial
# purposes without further permission.
#
# Copyright 2015-2016. Allen Institute. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Redistributions for commercial purposes are not permitted without the
# Allen Institute's written permission.
# For purposes of this license, commercial purposes is the incorporation of the
# Allen Institute's software into anything for which you will charge fees or
# other compensation. Contact terms@alleninstitute.org for commercial licensing
# opportunities.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
import csv
import copy
import math
import six

# Morphology nodes have the following fields. SWC fields are numeric.
NODE_ID = 'id'
NODE_TYPE = 'type'
NODE_X = 'x'
NODE_Y = 'y'
NODE_Z = 'z'
NODE_R = 'radius'
NODE_PN = 'parent'
SWC_COLUMNS = [NODE_ID, NODE_TYPE, NODE_X, NODE_Y, NODE_Z, NODE_R, NODE_PN]

NODE_TREE_ID = 'tree_id'
NODE_CHILDREN = 'children'

# shorthand for dictionary entries, to shorten sometimes long code lines
_N = NODE_ID
_TYP = NODE_TYPE
_X = NODE_X
_Y = NODE_Y
_Z = NODE_Z
_R = NODE_R
_P = NODE_PN
_C = NODE_CHILDREN
_TID = NODE_TREE_ID


########################################################################
def read_swc(file_name, columns="NOT_USED", numeric_columns="NOT_USED"):
    """
    Read in an SWC file and return a Morphology object.

    Parameters
    ----------
    file_name: string
        SWC file name.

    Returns
    -------
    Morphology
        A Morphology instance.
    """
    compartments = []
    line_num = 1
    try:
        with open(file_name, "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 = Compartment({
                    NODE_ID: int(toks[0]),
                    NODE_TYPE: int(toks[1]),
                    NODE_X: float(toks[2]),
                    NODE_Y: float(toks[3]),
                    NODE_Z: float(toks[4]),
                    NODE_R: float(toks[5]),
                    NODE_PN: int(toks[6].rstrip())
                })
                # store this compartment
                compartments.append(vals)
                # increment line number (used for error reporting only)
                line_num += 1
    except ValueError:
        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)

    return Morphology(compartment_list=compartments)


########################################################################
########################################################################
class Compartment(dict):
    """
    A dictionary class storing information about a single morphology node
    """

    def __init__(self, *args, **kwargs):
        super(Compartment, self).__init__(*args, **kwargs)
        if (NODE_ID not in self or
                NODE_TYPE not in self or
                NODE_X not in self or
                NODE_Y not in self or
                NODE_Z not in self or
                NODE_R not in self or
                NODE_PN not in self):
            raise ValueError(
                "Compartment was not initialized with requisite fields")
        # Each unconnected graph has its own ID. This is the ID
        #   of graph that the node resides in
        self[NODE_TREE_ID] = -1

        # IDs of child nodes
        self[NODE_CHILDREN] = []

    def print_node(self):
        """ print out compartment information with field names """
        print("%d %d %.4f %.4f %.4f %.4f %d %s %d" % (self[_N], self[_TYP], self[
              _X], self[_Y], self[_Z], self[_R], self[_P], str(self[_C]), self[_TID]))


class Morphology(object):
    """
    Keep track of the list of compartments in a morphology and provide
    a few helper methods (soma, tree information, pruning, etc).
    """

    SOMA = 1
    AXON = 2
    DENDRITE = 3
    BASAL_DENDRITE = 3
    APICAL_DENDRITE = 4

    NODE_TYPES = [SOMA, AXON, DENDRITE, BASAL_DENDRITE, APICAL_DENDRITE]

    def __init__(self, compartment_list=None, compartment_index=None):
        """
        Try to initialize from a list of compartments first, then from
        a dictionary indexed by compartment id if that fails, and finally just
        leave everything empty.

        Parameters
        ----------
        compartment_list: list
            list of compartment dictionaries

        compartment_index: dict
            dictionary of compartments indexed by id
        """
        self._compartment_list = []
        self._compartment_index = {}

        ##############################################
        # define tree list here for clarity, even though it's reset below
        #   when nodes are assigned
        self._tree_list = []

        ##############################################
        # construct the compartment list and index
        # first try to do so using the compartment list, then try using
        #   the compartment index and if that fails then complain
        if compartment_list:
            self.compartment_list = compartment_list
        elif compartment_index:
            self.compartment_index = compartment_index
        ##############################################
        # verify morphology is consistent with morphology rules (e.g.,
        #   no dendrite branching from an axon)
        num_errors = self._check_consistency()
        if num_errors > 0:
            raise ValueError("Morphology appears to be inconsistent")
        ##############################################
        # root node (this must be part of the soma)
        self._soma = None
        for i in range(len(self.compartment_list)):
            seg = self.compartment_list[i]
            if seg[NODE_TYPE] == Morphology.SOMA and seg[NODE_PN] < 0:
                if self._soma is not None:
                    raise ValueError("Multiple somas detected in SWC file")
                self._soma = seg

    ####################################################################
    ####################################################################
    # class properties, and helper functions for them

    @property
    def compartment_list(self):
        """ Return the compartment list.  This is a property to ensure that the
        compartment list and compartment index are in sync. """
        return self._compartment_list

    @compartment_list.setter
    def compartment_list(self, compartment_list):
        """ Update the compartment list.  Update the compartment index. """
        self._set_compartments(compartment_list)

    @property
    def compartment_index(self):
        """ Return the compartment index.  This is a property to ensure that the
        compartment list and compartment index are in sync. """
        return self._compartment_index

    @compartment_index.setter
    def compartment_index(self, compartment_index):
        """ Update the compartment index.  Update the compartment list. """
        self._set_compartments(compartment_index.values())

    @property
    def num_trees(self):
        """ Return the number of trees in the morphology. A tree is
        defined as everything following from a single root compartment. """
        return len(self._tree_list)

    # TODO add filter for number of nodes of a particular type
    @property
    def num_nodes(self):
        """ Return the number of compartments in the morphology. """
        return len(self.compartment_list)

    # internal function
    def _set_compartments(self, compartment_list):
        """
        take a list of SWC-like objects and turn those into morphology
        nodes need to be able to initialize from a list supplied by an SWC
        file while also being able to initialize from the compartment list
        of an existing Morphology object. As nodes in a morphology object
        contain reference to nodes in that object, make a shallow copy
        of input nodes and overwrite known references (ie, the
        'children' array)
        """
        self._compartment_list = []
        for obj in compartment_list:
            seg = copy.copy(obj)
            seg[NODE_TREE_ID] = -1
            seg[NODE_CHILDREN] = []
            self._compartment_list.append(seg)
        # list data now set. remove holes in sequence and re-index
        self._reconstruct()

    @property
    def soma(self):
        """ Returns root node of soma, if present"""
        return self._soma

    @property
    def root(self):
        """ [deprecated] Returns root node of soma, if present. Use 'soma' instead of 'root'"""
        return self._soma

    ####################################################################
    ####################################################################
    # tree and node access

    def tree(self, n):
        """
        Returns a list of all Morphology Nodes within the specified
        tree. A tree is defined as a fully connected graph of nodes.
        Each tree has exactly one root.

        Parameters
        ----------
        n: integer
            ID of desired tree

        Returns
        -------
        A list of all morphology objects in the specified tree, or None
        if the tree doesn't exist
        """
        if n < 0 or n >= len(self._tree_list):
            return None
        return self._tree_list[n]

    def node(self, n):
        """
        Returns the morphology node having the specified ID.

        Parameters
        ----------
        n: integer
            ID of desired node

        Returns
        -------
        A morphology object having the specified ID, or None if such a
        node doesn't exist
        """
        # undocumented feature -- if a node is supplied instead of a
        #   node ID, the node is returned and no error is triggered
        return self._resolve_node_type(n)

    def parent_of(self, seg):
        """ Returns parent of the specified node.

        Parameters
        ----------
        seg: integer or Morphology Object
            The ID of the child node, or the child node itself

        Returns
        -------
        A morphology object, or None if no parent exists or if the
        specified node ID doesn't exist
        """
        # if ID passed in, make sure it's converted to a compartment
        # don't trap for exception here -- if supplied segment is
        #   incorrect, make sure the user knows about it
        seg = self._resolve_node_type(seg)
        # return parent of specified node
        if seg is not None and seg[NODE_PN] >= 0:
            return self._compartment_list[seg[NODE_PN]]
        return None

    def children_of(self, seg):
        """ Returns a list of the children of the specified node

        Parameters
        ----------
        seg: integer or Morphology Object
            The ID of the parent node, or the parent node itself

        Returns
        -------
        A list of the child morphology objects. If the ID of the parent
        node is invalid, None is returned.
        """
        seg = self._resolve_node_type(seg)
        return [self._compartment_list[c] for c in seg[NODE_CHILDREN]]

    ###################################################################
    ###################################################################
    # Information querying and data manipulation

    # internal function. takes an integer and returns the node having
    #   that ID. IF a node is passed in instead, it is returned
    def _resolve_node_type(self, seg):
        # if compartment passed then we don't need to convert anything
        # if compartment not passed, try converting value to int
        #   and using that as an index
        if not isinstance(seg, Compartment):
            try:
                seg = int(seg)
                if seg < 0 or seg >= len(self._compartment_list):
                    return None
                seg = self._compartment_list[seg]
            except ValueError:
                raise TypeError(
                    "Object not recognized as morphology node or index")
        return seg

    def change_parent(self, child, parent):
        """ Change the parent of a node. The child node is adjusted to
        point to the new parent, the child is taken off of the previous
        parent's child list, and it is added to the new parent's child list.

        Parameters
        ----------
        child: integer or Morphology Object
            The ID of the child node, or the child node itself

        parent: integer or Morphology Object
            The ID of the parent node, or the parent node itself

        Returns
        -------
        Nothing
        """
        child_seg = self._resolve_node_type(child)
        parent_seg = self._resolve_node_type(parent)
        # if child has former parent, remove it from parent's child list
        if child_seg[NODE_PN] >= 0:
            old_par = self.node(child_seg[NODE_PN])
            old_par[NODE_CHILDREN].remove(child_seg[NODE_ID])
        parent_seg[NODE_CHILDREN].append(child_seg[NODE_ID])
        child_seg[NODE_PN] = parent_seg[NODE_ID]

    # returns a list of nodes located within dist of x,y,z
    def find(self, x, y, z, dist, node_type=None):
        """ Returns a list of Morphology Objects located within 'dist'
        of coordinate (x,y,z). If node_type is specified, the search
        will be constrained to return only nodes of that type.

        Parameters
        ----------
        x, y, z: float
            The x,y,z coordinates from which to search around

        dist: float
            The search radius

        node_type: enum (optional)
            One of the following constants: SOMA, AXON, DENDRITE,
            BASAL_DENDRITE or APICAL_DENDRITE

        Returns
        -------
        A list of all Morphology Objects matching the search criteria
        """
        found = []
        for seg in self.compartment_list:
            dx = seg[NODE_X] - x
            dy = seg[NODE_Y] - y
            dz = seg[NODE_Z] - z
            if math.sqrt(dx * dx + dy * dy + dz * dz) <= dist:
                if node_type is None or seg[NODE_TYPE] == node_type:
                    found.append(seg)
        return found

    def compartment_list_by_type(self, compartment_type):
        """ Return an list of all compartments having the specified
        compartment type.

        Parameters
        ----------
        compartment_type: int
            Desired compartment type

        Returns
        -------
        A list of of Morphology Objects
        """
        return [x for x in self._compartment_list if x[NODE_TYPE] == compartment_type]

    def compartment_index_by_type(self, compartment_type):
        """ Return an dictionary of compartments indexed by id that all have
        a particular compartment type.

        Parameters
        ----------
        compartment_type: int
            Desired compartment type

        Returns
        -------
        A dictionary of Morphology Objects, indexed by ID
        """
        return {c[NODE_ID]: c for c in self._compartment_list if c[NODE_TYPE] == compartment_type}

    def save(self, file_name):
        """ Write this morphology out to an SWC file

        Parameters
        ----------
        file_name: string
            desired name of your SWC file
        """
        f = open(file_name, "w")
        f.write("#n,type,x,y,z,radius,parent\n")
        for seg in self.compartment_list:
            f.write("%d %d " % (seg[NODE_ID], seg[NODE_TYPE]))
            f.write("%0.4f " % seg[NODE_X])
            f.write("%0.4f " % seg[NODE_Y])
            f.write("%0.4f " % seg[NODE_Z])
            f.write("%0.4f " % seg[NODE_R])
            f.write("%d\n" % seg[NODE_PN])
        f.close()

    # keep for backward compatibility, but don't publish in docs
    def write(self, file_name):
        self.save(file_name)

    def sparsify(self, modulo, compress_ids=False):
        """ Return a new Morphology object that has a given number of non-leaf,
        non-root nodes removed.  IDs can be reassigned so as to be continuous.

        Parameters
        ----------
        modulo: int
           keep 1 out of every modulo nodes.

        compress_ids: boolean
           Reassign ids so that ids are continuous (no missing id numbers).

        Returns
        -------
        Morphology
            A new morphology instance
        """
        compartments = self.compartment_index
        root = self.root
        keep = {}
        # figure out which compartments to toss
        ct = 0
        for i, c in six.iteritems(compartments):
            pid = c[NODE_PN]
            cid = c[NODE_ID]
            ctype = c[NODE_TYPE]
            # keep the root, soma, junctions, and the first child of the root
            # (for visualization)
            if pid < 0 or len(c[NODE_CHILDREN]) != 1 or pid == root[NODE_ID] or ctype == Morphology.SOMA:
                keep[cid] = True
            else:
                keep[cid] = (ct % modulo) == 0
            ct += 1

        # hook children up to their new parents
        for i, c in six.iteritems(compartments):
            comp_id = c[NODE_ID]
            if keep[comp_id] is False:
                parent_id = c[NODE_PN]
                while keep[parent_id] is False:
                    parent_id = compartments[parent_id][NODE_PN]
                for child_id in c[NODE_CHILDREN]:
                    compartments[child_id][NODE_PN] = parent_id

        # filter out the orphans
        sparsified_compartments = {k: v for k,
                                   v in six.iteritems(compartments) if keep[k]}
        if compress_ids:
            ids = sorted(sparsified_compartments.keys(), key=lambda x: int(x))
            id_hash = {fid: str(i + 1) for i, fid in enumerate(ids)}
            id_hash[-1] = -1
            # build the final compartment index
            out_compartments = {}
            for cid, compartment in six.iteritems(sparsified_compartments):
                compartment[NODE_ID] = id_hash[cid]
                compartment[NODE_PN] = id_hash[compartment[NODE_PN]]
                out_compartments[compartment[NODE_ID]] = compartment
            return Morphology(compartment_index=out_compartments)
        else:
            return Morphology(compartment_index=sparsified_compartments)

    ####################################################################
    ####################################################################
    def _reconstruct(self):
        """
        internal function that restructures data and establishes
        appropriate internal linking. data is re-order, removing 'holes'
        in sequence so that each object ID corresponds to its position
        in compartment list. trees are (re)calculated
        parent-child indices are recalculated as is compartment table
        construct a map between new and old IDs
        """
        remap = {}
        # everything defaults to root. this way if a parent was deleted
        #   the child will become a new root
        for i in range(len(self.compartment_list)):
            remap[i] = -1
        # map old old node numbers to new ones. reset n to the new ID
        #   and put node in new list
        new_id = 0
        tmp_list = []
        for seg in self.compartment_list:
            if seg is not None:
                remap[seg[NODE_ID]] = new_id
                seg[NODE_ID] = new_id
                tmp_list.append(seg)
                new_id += 1
        # use map to reset parent values. copy objs to new list
        for seg in tmp_list:
            if seg[NODE_PN] >= 0:
                seg[NODE_PN] = remap[seg[NODE_PN]]
        # replace compartment list with newly created node list
        self._compartment_list = tmp_list
        # reconstruct parent/child relationship links
        # forget old relations
        for seg in self.compartment_list:
            seg[NODE_CHILDREN] = []
        # add each object to its parents child list
        for seg in self.compartment_list:
            par_num = seg[NODE_PN]
            if par_num >= 0:
                self.compartment_list[par_num][
                    NODE_CHILDREN].append(seg[NODE_ID])
        # update tree lists
        self._separate_trees()
        ############################
        # Rebuild internal index and links between parents and children
        self._compartment_index = {
            c[NODE_ID]: c for c in self.compartment_list}
        # compartment list is complete and sequential so don't need index
        #   to resolve relationships
        # for each node, reset children array
        # for each node, add self to parent's child list
        for seg in self._compartment_list:
            seg[NODE_CHILDREN] = []
        for seg in self._compartment_list:
            if seg[NODE_PN] >= 0:
                self._compartment_list[seg[NODE_PN]][
                    NODE_CHILDREN].append(seg[NODE_ID])
        # verify that each node ID is the same as its position in the
        #   compartment list
        for i in range(len(self.compartment_list)):
            if i != self.node(i)[NODE_ID]:
                raise RuntimeError(
                    "Internal error detected -- compartment list not properly formed")

    def append(self, node_list):
        """ Add additional nodes to this Morphology. Those nodes must
        originate from another morphology object.

        Parameters
        ----------
        node_list: list of Morphology nodes
        """
        # construct a map between new and old IDs of added nodes
        remap = {}
        for i in range(len(node_list)):
            remap[i] = -1
        # map old old node numbers to new ones. reset n to the new ID
        # append new nodes to existing node list
        old_count = len(self.compartment_list)
        new_id = old_count
        for seg in node_list:
            if seg is not None:
                remap[seg[NODE_ID]] = new_id
                seg[NODE_ID] = new_id
                self._compartment_list.append(seg)
                new_id += 1
        # use map to reset parent values. copy objs to new list
        for i in range(old_count, len(self.compartment_list)):
            seg = self.compartment_list[i]
            if seg[NODE_PN] >= 0:
                seg[NODE_PN] = remap[seg[NODE_PN]]
        self._reconstruct()

    def stumpify_axon(self, count=10):
        """ Remove all axon compartments except the first 'count'
        nodes, as counted from the connected axon root.

        Parameters
        ----------
        count: Integer
            The length of the axon 'stump', in number of compartments
        """
        # find connected axon root
        axon_root = None
        for seg in self.compartment_list:
            if seg[NODE_TYPE] == Morphology.AXON:
                par_id = seg[NODE_PN]
                if par_id >= 0:
                    par = self.compartment_list[par_id]
                    if par[NODE_TYPE] != Morphology.AXON:
                        axon_root = seg
                        break
        if axon_root is None:
            return
        # flag the first 'count' nodes from the axon root
        ax = axon_root
        for i in range(count):
            # ignore bifurcations -- go 'count' deep on one line only
            ax["flag"] = i
            children = ax[NODE_CHILDREN]
            if len(children) > 0:
                ax = children[0]
        # strip out all axons that aren't flagged
        for i in range(len(self.compartment_list)):
            seg = self.compartment_list[i]
            if seg[NODE_TYPE] == Morphology.AXON:
                if "flag" not in seg:
                    self.compartment_list[i] = None
        self._reconstruct()

    # strip out everything but the soma and the specified SWC type
    def strip_all_other_types(self, node_type, keep_soma=True):
        """ Strips everything from the morphology except for the
        specified type.
        Parent and child relationships are updated accordingly, creating
        new roots when necessary.

        Parameters
        ----------
        node_type: enum
            The compartment type to keep in the morphology.
            Use one of the following constants: SOMA, AXON, DENDRITE,
            BASAL_DENDRITE, or APICAL_DENDRITE

        keep_soma: Boolean (optional)
            True (default) if soma nodes should remain in the
            morpyhology, and False if the soma should also be stripped
        """
        flagged_for_removal = {}
        # scan nodes and see which ones should be removed. keep a record
        #   of them
        for seg in self.compartment_list:
            if seg[NODE_TYPE] == node_type:
                remove = False
            elif seg[NODE_TYPE] == 1 and keep_soma:
                remove = False
            else:
                remove = True
            if remove:
                flagged_for_removal[seg[NODE_ID]] = True
        # remove selected nodes andreset parent links
        for i in range(len(self.compartment_list)):
            seg = self.compartment_list[i]
            if seg[NODE_ID] in flagged_for_removal:
                # eliminate node
                self.compartment_list[i] = None
            elif seg[NODE_PN] in flagged_for_removal:
                # parent was eliminated. make this a new root
                seg[NODE_PN] = -1
        self._reconstruct()

    # strip out the specified SWC type
    def strip_type(self, node_type):
        """ Strips all compartments of the specified type from the
        morphology.
        Parent and child relationships are updated accordingly, creating
        new roots when necessary.

        Parameters
        ----------
        node_type: enum
            The compartment type to strip from the morphology.
            Use one of the following constants: SOMA, AXON, DENDRITE,
            BASAL_DENDRITE, or APICAL_DENDRITE
        """
        flagged_for_removal = {}
        for seg in self.compartment_list:
            if seg[NODE_TYPE] == node_type:
                remove = True
            else:
                remove = False
            if remove:
                flagged_for_removal[seg[NODE_ID]] = True
        for i in range(len(self.compartment_list)):
            seg = self.compartment_list[i]
            if seg[NODE_ID] in flagged_for_removal:
                # eliminate node
                self.compartment_list[i] = None
            elif seg[NODE_PN] in flagged_for_removal:
                # parent was eliminated. make this a new root
                seg[NODE_PN] = -1
        self._reconstruct()

    # strip out the specified SWC type
    def convert_type(self, old_type, new_type):
        """ Converts all compartments from one type to another.
        Nodes of the original type are not affected so this
        procedure can also be used as a merge procedure.

        Parameters
        ----------
        old_type: enum
            The compartment type to be changed.
            Use one of the following constants: SOMA, AXON, DENDRITE,
            BASAL_DENDRITE, or APICAL_DENDRITE

        new_type: enum
            The target compartment type.
            Use one of the following constants: SOMA, AXON, DENDRITE,
            BASAL_DENDRITE, or APICAL_DENDRITE
        """
        for seg in self.compartment_list:
            if seg[NODE_TYPE] == old_type:
                seg[NODE_TYPE] = new_type

    def apply_affine(self, aff, scale=None):
        """ Apply an affine transform to all compartments in this
        morphology. Node radius is adjusted as well.

        Format of the affine matrix is:

        [x0 y0 z0]  [tx]
        [x1 y1 z1]  [ty]
        [x2 y2 z2]  [tz]

        where the left 3x3 the matrix defines the affine rotation
        and scaling, and the right column is the translation
        vector.

        The matrix must be collapsed and stored in a list as follows:

        [x0 y0, z0, x1, y1, z1, x2, y2, z2, tx, ty, tz]

        Parameters
        ----------
        aff: 3x4 array of floats (python 2D list, or numpy 2D array)
            the transformation matrix
        """
        # In addition to transforming the locations of the morphology
        #   nodes, the radius of each node must be adjusted.
        # There are 2 ways to measure scale from a transform. Assuming
        #   an isotropic transform, the scale is the cube root of the
        #   matrix determinant. The other ways is to measure scale
        #   independently along each axis.
        # For now, the node radius is only updated based on the average
        #   scale along all 3 axes (eg, isotropic assumption), so calculate
        #   scale using the determinant
        #
        if scale is None:
            # calculate the determinant
            det0 = aff[0] * (aff[4] * aff[8] - aff[5] * aff[7])
            det1 = aff[1] * (aff[3] * aff[8] - aff[5] * aff[6])
            det2 = aff[2] * (aff[3] * aff[7] - aff[4] * aff[6])
            det = det0 + det1 + det2
            # determinant is change of volume that occurred during transform
            # assume equal scaling along all axes. take 3rd root to get
            #   scale factor
            det_scale = math.pow(abs(det), 1.0 / 3.0)
            # measure scale along each axis
            # keep this code here in case
            #scale_x = abs(aff[0] + aff[3] + aff[6])
            #scale_y = abs(aff[1] + aff[4] + aff[7])
            #scale_z = abs(aff[2] + aff[5] + aff[8])
            #avg_scale = (scale_x + scale_y + scale_z) / 3.0;
            #
            # use determinant for scaling for now as it's most simple
            scale = det_scale
        for seg in self.compartment_list:
            x = seg[NODE_X] * aff[0] + seg[NODE_Y] * \
                aff[1] + seg[NODE_Z] * aff[2] + aff[9]
            y = seg[NODE_X] * aff[3] + seg[NODE_Y] * \
                aff[4] + seg[NODE_Z] * aff[5] + aff[10]
            z = seg[NODE_X] * aff[6] + seg[NODE_Y] * \
                aff[7] + seg[NODE_Z] * aff[8] + aff[11]
            seg[NODE_X] = x
            seg[NODE_Y] = y
            seg[NODE_Z] = z
            seg[NODE_R] *= scale

    def _separate_trees(self):
        """
        construct list of independent trees (each tree has a root of -1)
        """
        trees = []
        # reset each node's tree ID to indicate that it's not assigned
        for seg in self.compartment_list:
            seg[NODE_TREE_ID] = -1
        # construct trees for each node
        # if a node is adjacent an existing tree, merge to it
        # if a node is adjacent multiple trees, merge all
        for seg in self.compartment_list:
            # see what trees this node is adjacent to
            local_trees = []
            if seg[NODE_PN] >= 0 and self.compartment_list[seg[NODE_PN]][NODE_TREE_ID] >= 0:
                local_trees.append(self.compartment_list[
                                   seg[NODE_PN]][NODE_TREE_ID])
            for child_id in seg[NODE_CHILDREN]:
                child = self.compartment_list[child_id]
                if child[NODE_TREE_ID] >= 0:
                    local_trees.append(child[NODE_TREE_ID])
            # figure out which tree to put node into
            # if there are muliple possibilities, merge all of them
            if len(local_trees) == 0:
                tree_num = len(trees)  # create new tree
            elif len(local_trees) == 1:
                tree_num = local_trees[0]   # use existing tree
            elif len(local_trees) > 1:
                # this node is an intersection of multiple trees
                # merge all trees into the first one found
                tree_num = local_trees[0]
                for j in range(1, len(local_trees)):
                    dead_tree = local_trees[j]
                    trees[dead_tree] = []
                    for node in self.compartment_list:
                        if node[NODE_TREE_ID] == dead_tree:
                            node[NODE_TREE_ID] = tree_num
            # merge node into tree
            # ensure there's space
            while len(trees) <= tree_num:
                trees.append([])
            trees[tree_num].append(seg)
            seg[NODE_TREE_ID] = tree_num
        # consolidate tree lists into class's tree list object
        self._tree_list = []
        for tree in trees:
            if len(tree) > 0:
                self._tree_list.append(tree)
        # make soma's tree be the first tree, if soma present
        # this should be the case if the file is properly ordered, but
        #   don't assume that
        soma_tree = -1
        for seg in self.compartment_list:
            if seg[NODE_TYPE] == 1:
                soma_tree = seg[NODE_TREE_ID]
                break
        if soma_tree > 0:
            # swap soma tree for first tree in list
            tmp = self._tree_list[soma_tree]
            self._tree_list[soma_tree] = self._tree_list[0]
            self._tree_list[0] = tmp
        # reset node tree_id to correct tree number
        self._reset_tree_ids()

    def _reset_tree_ids(self):
        """
        reset each node's tree_id value to the correct tree number
        """
        for i in range(len(self._tree_list)):
            for j in range(len(self._tree_list[i])):
                self._tree_list[i][j][NODE_TREE_ID] = i

    def _check_consistency(self):
        """
        internal function -- don't publish in the docs
        TODO? print warning if unrecognized types are present
        Return value: number of errors detected in file
        """
        errs = 0
        # Make sure that the parents are of proper ID range
        n = self.num_nodes
        for seg in self.compartment_list:
            if seg[NODE_PN] >= 0:
                if seg[NODE_PN] >= n:
                    print("Parent for node %d is invalid (%d)" %
                          (seg[NODE_ID], seg[NODE_PN]))
                    errs += 1
        # make sure that each tree has exactly one root
        for i in range(self.num_trees):
            tree = self.tree(i)
            root = -1
            for j in range(len(tree)):
                if tree[j][NODE_PN] == -1:
                    if root >= 0:
                        print("Too many roots in tree %d" % i)
                        errs += 1
                    root = j
            if root == -1:
                print("No root present in tree %d" % i)
                errs += 1
        # make sure each axon has at most one root
        # find type boundaries. at each axon boundary, walk back up
        #   tree to root and make sure another axon segment not
        #   encountered
        adoptees = self._find_type_boundary()
        for child in adoptees:
            if child[NODE_TYPE] == Morphology.AXON:
                par_id = child[NODE_PN]
                while par_id >= 0:
                    par = self.compartment_list[par_id]
                    if par[NODE_TYPE] == Morphology.AXON:
                        print("Branch has multiple axon roots")
                        print(child)
                        print(par)
                        errs += 1
                        break
                    par_id = par[NODE_PN]
        if errs > 0:
            print("Failed consistency check: %d errors encountered" % errs)
        return errs

    def _find_type_boundary(self):
        """
        return a list of segments who have parents that are a different type
        """
        adoptees = []
        for node in self.compartment_list:
            par = self.parent_of(node)
            if par is None:
                continue
            if node[NODE_TYPE] != par[NODE_TYPE]:
                adoptees.append(node)
        return adoptees

    # remove tree from swc's "forest"
    def delete_tree(self, n):
        """ Delete tree, and all of its compartments, from the morphology.

        Parameters
        ----------
        n: Integer
            The tree number to delete
        """
        if n < 0:
            return
        if n >= self.num_trees:
            print("Error -- attempted to delete non-existing tree (%d)" % n)
            raise ValueError
        tree = self.tree(n)
        for i in range(len(tree)):
            self.compartment_list[tree[i][NODE_ID]] = None
        del self._tree_list[n]
        self._reconstruct()
        # reset node tree_id to correct tree number
        self._reset_tree_ids()

    def _print_all_nodes(self):
        """
        debugging function. prints all nodes
        """
        for node in self.compartment_list:
            print(node)

########################################################################
class Marker(dict):
    """ Simple dictionary class for handling reconstruction marker objects. """

    SPACING = [.1144, .1144, .28]

    CUT_DENDRITE = 10
    NO_RECONSTRUCTION = 20

    def __init__(self, *args, **kwargs):
        super(Marker, self).__init__(*args, **kwargs)

        # marker file x,y,z coordinates are offset by a single image-space
        # pixel
        self['x'] -= self.SPACING[0]
        self['y'] -= self.SPACING[1]
        self['z'] -= self.SPACING[2]


def read_marker_file(file_name):
    """ read in a marker file and return a list of dictionaries """

    with open(file_name, 'r') as f:
        rows = csv.DictReader((r for r in f if not r.startswith('#')),
                              fieldnames=['x', 'y', 'z', 'radius', 'shape', 'name', 'comment',
                                          'color_r', 'color_g', 'color_b'])

        return [Marker({'x': float(r['x']),
                        'y': float(r['y']),
                        'z': float(r['z']),
                        'name': int(r['name'])}) for r in rows]
