"""
Arbor class and member functions



"""

#-----------------------------------------------------------------------------
# Copyright (c) ytree development team. All rights reserved.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

from collections import \
    defaultdict
import functools
import numpy as np
import os

from unyt import \
    unyt_array, \
    unyt_quantity

from yt.funcs import \
    get_pbar, \
    TqdmProgressBar
from unyt.dimensions import \
    dimensionless, \
    length
from unyt.unit_registry import \
    UnitRegistry
from yt.utilities.cosmology import \
    Cosmology

from ytree.data_structures.detection import \
    SelectionDetector
from ytree.data_structures.fields import \
    FieldContainer, \
    FieldInfoContainer
from ytree.data_structures.io import \
    DefaultRootFieldIO, \
    TreeFieldIO
from ytree.data_structures.node_link import \
    NodeLink
from ytree.data_structures.save_arbor import \
    save_arbor
from ytree.data_structures.node_container import NodeContainer
from ytree.data_structures.tree_node import \
    TreeNode
from ytree.data_structures.tree_node_selector import \
    tree_node_selector_registry
from ytree.utilities.logger import \
    ytreeLogger, \
    fake_pbar

arbor_registry = {}
_selection_types = ("forest", "tree", "prog")

class RegisteredArbor(type):
    """
    Add to the registry of known Arbor classes to cycle
    through in the load function.
    """
    def __init__(cls, name, b, d):
        type.__init__(cls, name, b, d)
        arbor_type = name[:name.rfind("Arbor")]
        if arbor_type:
            arbor_registry[arbor_type] = cls

class Arbor(metaclass=RegisteredArbor):
    """
    Base class for all Arbor classes.

    Loads a merger tree output file or a series of halo catalogs
    and create trees, stored in an array in
    :func:`~ytree.data_structures.arbor.Arbor.trees`.
    Arbors can be saved in a universal format with
    :func:`~ytree.data_structures.arbor.Arbor.save_arbor`. Also, provide some
    convenience functions for creating unyt_arrays and unyt_quantities and
    a cosmology calculator.
    """

    _field_info_class = FieldInfoContainer
    _root_field_io_class = DefaultRootFieldIO
    _tree_field_io_class = TreeFieldIO
    _default_dtype = np.float64

    ### attributes required for generating a TreeNode object
    ### for a given Arbor class.
    ### We store these in arrays and use them to generate TreeNodes
    ### when they are needed.
    ## attributes required for constructing TreeNodes
    _node_con_attrs = ('uid',)
    ## attributes we may not have, but would be nice if we did
    _node_too_attrs = ('_tree_size',)
    ## attributes specific to an Arbor class for facilitating io
    _node_io_attrs = ()

    ### tree node attributes for all Arbor types.
    ### These facilitate walking the tree, getting fields, etc.
    ### We keep track of these for resetting TreeNodes and
    ### deciding when they are setup or grown.
    _reset_attrs = ("_tfi", "_pfi")
    _setup_attrs = ("_desc_uids", "_uids")
    _grow_attrs = ("_link_storage", "_link")

    omega_matter = None
    omega_lambda = None
    omega_radiation = 0

    def __init__(self, filename):
        """
        Initialize an Arbor given an input file.
        """

        self._set_paths(filename)
        self._parse_parameter_file()
        self._set_units()
        self._setup_io()
        self._get_data_files()
        self._setup_fields()
        self._set_default_selector()

    def _set_paths(self, filename):
        """
        Set data paths.
        """

        self.filename = filename
        if isinstance(filename, (list, tuple)):
            fn = filename[0]
        else:
            fn = filename
        self._set_parameter_filename(fn)
        self.basename = os.path.basename(fn)
        dn = os.path.dirname(fn)
        self.directory = dn if dn else '.'

    def _set_parameter_filename(self, filename):
        """
        Set the name of the file from which to get metadata.
        """

        self.parameter_filename = filename

    def _parse_parameter_file(self):
        """
        Read relevant parameters from parameter file or file header
        and detect fields.
        """
        raise NotImplementedError

    def _set_units(self):
        """
        Set "cm" units for explicitly comoving.
        Note, we are using comoving units all the time since
        we are dealing with data at multiple redshifts.
        """
        for my_unit in ["m", "pc", "AU"]:
            new_unit = f"{my_unit}cm"
            self.unit_registry.add(
                new_unit, self.unit_registry.lut[my_unit][0],
                length, self.unit_registry.lut[my_unit][3])

    @functools.cached_property
    def cosmology(self):
        """
        Cosmology calculator with self-consistent parameters and units.

        See :ref:`yt:cosmology-calculator` for more information.
        """

        missing = []
        for attr in ["hubble_constant",
                     "omega_matter",
                     "omega_lambda"]:
            if getattr(self, attr) is None:
                missing.append(attr)

        if missing:
            ytreeLogger.warning(
                "Cannot create cosmology calculator without missing attributes: " +
                str(missing))
            return None

        else:
            return Cosmology(
                hubble_constant=self.hubble_constant,
                omega_matter=self.omega_matter,
                omega_lambda=self.omega_lambda,
                omega_radiation=self.omega_radiation,
                unit_registry=self.unit_registry)

    def _setup_io(self):
        """
        Create field io objects.
        """
        self._node_io = self._tree_field_io_class(
            self, default_dtype=self._default_dtype)
        self._root_io = self._root_field_io_class(
            self, default_dtype=self._default_dtype)

    def _get_data_files(self):
        """
        Get all files that hold field data and make them known
        to the i/o system.
        """
        pass

    def _setup_fields(self):
        """
        Setup field containers and definitions.
        """
        self.field_data = FieldContainer(self)
        self.derived_field_list = []
        # This may already be created by ytree frontend.
        if not hasattr(self, "analysis_field_list"):
            self.analysis_field_list = []
        self.field_info.setup_known_fields()
        self.field_info.setup_aliases()
        self.field_info.setup_derived_fields()
        self.field_info.setup_vector_fields()

    def _set_default_selector(self):
        """
        Set the default tree node selector as maximum mass.
        """
        self.set_selector("max_field_value", "mass")

    @property
    def is_planted(self):
        """
        Determine if trees have been planted.
        """
        return self._node_info_storage is not None

    def _plant_trees(self):
        """
        Create arrays to construct root nodes.
        """
        raise NotImplementedError

    _node_info_storage = None
    @property
    def _node_info(self):
        """
        The dict of arrays for storing node information.
        """

        if self._node_info_storage is not None:
            return self._node_info_storage
        self._initialize_node_info()
        return self._node_info_storage

    def _initialize_node_info(self):
        """
        Initialize the node_info arrays.
        """

        attrs = self._node_con_attrs + \
          self._node_io_attrs

        self._node_info_storage = \
          dict((attr, np.empty(self._size, dtype=np.int64))
               for attr in attrs)
        # initialize the target of opportunity attributes
        self._node_info_storage.update(
            dict((attr, -np.ones(self._size, dtype=np.int64))
                for attr in self._node_too_attrs))

    def is_setup(self, tree_node):
        """
        Return True if arrays of uids and descendent uids have
        been read in. Setup has also completed if tree is already
        grown.
        """
        return self.is_grown(tree_node) or \
          tree_node._uids is not None

    def _setup_tree(self, tree_node, **kwargs):
        """
        Create arrays of uids and desc_uids and attach them to the
        root node.
        """

        # skip if this is not a root or if already setup
        if self.is_setup(tree_node):
            return

        idtype = np.int64
        fields, _ = \
          self.field_info.resolve_field_dependencies(["uid", "desc_uid"])
        halo_id_f, desc_id_f = fields
        dtypes = {halo_id_f: idtype, desc_id_f: idtype}
        # Note to self, we call _read_fields and not _get_fields to
        # avoid recursion issues.
        field_data  = self._node_io._read_fields(tree_node, fields,
                                                 dtypes=dtypes, **kwargs)

        tree_node._uids      = field_data[halo_id_f]
        tree_node._desc_uids = field_data[desc_id_f]
        tree_node._tree_size = tree_node._uids.size
        tree_node.field_data["uid"] = tree_node._uids
        tree_node.field_data["desc_uid"] = tree_node._desc_uids

    def is_grown(self, tree_node):
        """
        Return True if a tree has been fully assembled, i.e.,
        the hierarchy of ancestor tree nodes has been built.
        """
        return tree_node.root != -1

    def _grow_tree(self, tree_node, **kwargs):
        """
        Construct the hierarchy of ancestors and descendents
        for all nodes in the tree.
        """

        # skip this if not a root or if already grown
        if self.is_grown(tree_node):
            return

        self._setup_tree(tree_node, **kwargs)
        size      = tree_node.tree_size
        uids      = tree_node.uids
        desc_uids = tree_node.desc_uids
        links     = np.empty(size, dtype=object)

        # Make a dict mapping uids to index of storage array.
        # First, try to get indices out as the dict is constructed
        # since the dict will be smaller at first.
        uidmap = {}
        not_found = []
        for i, (uid, desc_uid) in enumerate(zip(uids, desc_uids)):
            node = NodeLink(i)
            uidmap[uid] = i
            desc_index = uidmap.get(desc_uid)
            if desc_index is None:
                not_found.append((node, desc_uid))
            else:
                desc = links[desc_index]
                desc.add_ancestor(node)
            links[i] = node

        # Make any additional links missed on the first pass.
        for node, desc_uid in not_found:
            if desc_uid == -1:
                continue
            desc = links[uidmap[desc_uid]]
            desc.add_ancestor(node)

        tree_node.root = tree_node
        tree_node._link = links[0]
        tree_node._link_storage = links

    _attr_map = None
    def _build_attr(self, attr, tree_node):
        """
        Call the correct function for building a given attribute.
        """

        if self._attr_map is None:
            self._attr_map = \
              dict([(attr, self._setup_tree)
                    for attr in self._setup_attrs] +
                   [(attr, self._grow_tree)
                    for attr in self._grow_attrs])

        self._attr_map[attr](tree_node)

    def reset_node(self, tree_node):
        """
        Reset all data structures for a single node.

        The goal is to clear as many data structures as
        possible without rendering the node object useless,
        if they are intended to be kept around.

        Calling reset_node on a non-root node should not make
        the non-root node useless.

        Calling reset_node on a root node will render any
        generated non-root nodes useless.
        """

        tree_node.clear_fields()
        attrs = self._reset_attrs

        if tree_node.is_root:
            if self.is_grown(tree_node):
                attrs += self._grow_attrs
                tree_node.root = -1
            if self.is_setup(tree_node):
                attrs += self._setup_attrs

        for attr in attrs:
            setattr(tree_node, attr, None)

    @property
    def ytds(self):
        raise NotImplementedError(
            "Only ytree data can be loaded with yt. "
            "Save data with save_arbor and then reload.")

    def _node_io_loop(self, func, *args, **kwargs):
        """
        Call the provided function over a list of nodes.

        If possible, group nodes by common data files to speed
        things up.

        Parameters
        ----------
        func : function
            Function to be called on an array of nodes.
        pbar : optional, string or yt.funcs.TqdmProgressBar
            A progress bar to be updated with each iteration.
            If a string, a progress bar will be created and the
            finish function will be called. If a progress bar is
            provided, the finish function will not be called.
            Default: None (no progress bar).
        root_nodes : optional, array of root TreeNodes
            Array of nodes over which the function will be called.
            If None, the list will be self[:] (i.e., all
            root_nodes).
            Default: None.

        Returns
        -------
        rvals : list
            return values from calling func on each node.
            These will have the same order as the original node list.
        """

        self._plant_trees()

        pbar = kwargs.pop("pbar", None)
        root_nodes = kwargs.pop("root_nodes", None)

        data_files, node_list, return_order = \
          self._node_io_loop_prepare(root_nodes)
        nnodes = sum([nodes.size for nodes in node_list])

        finish = True
        if pbar is None:
            pbar = fake_pbar("", nnodes)
        elif not isinstance(pbar, TqdmProgressBar):
            pbar = get_pbar(pbar, nnodes)
        else:
            finish = False

        rvals = []
        c = 0
        for data_file, nodes in zip(data_files, node_list):
            self._node_io_loop_start(data_file)

            # if we're doing all of them, just give the indices
            if root_nodes is None:
                my_nodes = nodes
            else:
                my_nodes = root_nodes[nodes]

            for node in self._yield_root_nodes(my_nodes):
                rval = func(node, *args, **kwargs)
                rvals.append(rval)
                c += 1
                pbar.update(c)

            self._node_io_loop_finish(data_file)

        if finish:
            pbar.finish()

        if return_order is not None:
            rvals = [rvals[i] for i in return_order]

        return rvals

    def _node_io_loop_start(self, data_file):
        pass

    def _node_io_loop_finish(self, data_file):
        pass

    def _node_io_loop_prepare(self, nodes):
        """
        This is called at the beginning of _node_io_loop.

        In different frontends, this can be used to group nodes by
        common data files. If nodes is None, we want all root
        nodes in the Arbor.

        Below is the default behavior, which does the bare minimum
        of returning:
        list of [None] : meaning all nodes come in a single group
            associated with no particular data file.
        list containing array of all provided nodes: meaning there
            is no specific grouping to be done
        None : meaning the nodes do not have to be reordered after
            being processed.

        See the implementation in individual frontends for
        more informative examples.

        Returns
        -------
        data_files : list
            list of data files that will be used
        index_list : list of arrays
            indices of the provided array of nodes associated
            with each of the data files
        return_order : int array
            array of indices used to reorder the return values
            to the order of the provided nodes
        """

        self._plant_trees()

        if nodes is None:
            my_size = self.size
        else:
            my_size = nodes.size
        indices = np.arange(my_size)

        return [None], [indices], None

    def __iter__(self):
        """
        Iterate over all trees in the arbor.
        """

        self._plant_trees()
        for node in self._yield_root_nodes(range(self.size)):
            yield node

    def __repr__(self):
        return self.basename

    def __getitem__(self, key):
        return self.query(key)

    def query(self, key):
        """
        If given a string, return an array of field values for the
        roots of all trees.
        If given an integer, return a tree from the list of trees.
        """

        if isinstance(key, str):
            if key in _selection_types:
                raise SyntaxError("Argument must be a field or integer.")
            self._root_io.get_fields(self, fields=[key])
            return self.field_data[key]
        return self._generate_root_nodes(key)

    def _generate_root_nodes(self, key):
        """
        Create root nodes given an index or slice from uid array.
        """

        self._plant_trees()
        if isinstance(key, (int, np.integer)):
            return self._generate_root_node(key)
        elif isinstance(key, slice) or isinstance(key, np.ndarray):
            indices = np.arange(self.size)[key]
            return self._yield_root_nodes(indices)
        else:
            raise ValueError('Cannot generate nodes from argument: ', key)

    def _yield_root_nodes(self, indices):
        """
        Root node generator.
        """

        # If we've been given an array of TreeNodes,
        # just yield them back.
        if getattr(indices, 'dtype', None) == object:
            for index in indices:
                yield index
            return

        for index in indices:
            node = self._generate_root_node(index)
            yield node

    def _generate_root_node(self, index):
        """
        Create a root node given its index in the array of uids.
        """

        args = tuple(self._node_info[attr][index]
                      for attr in self._node_con_attrs)
        my_node = TreeNode(*args, arbor=self, root=True)
        my_node._arbor_index = index

        for attr in self._node_io_attrs:
            setattr(my_node, attr, self._node_info[attr][index])

        for attr in self._node_too_attrs:
            val = self._node_info[attr][index]
            if val != -1:
                setattr(my_node, attr, self._node_info[attr][index])

        return my_node

    def _generate_tree_node(self, root_node, node_link):
        """
        Create a non-root node in a tree.
        """

        tree_id = node_link.tree_id
        if tree_id == 0:
            return root_node
        uid        = root_node.uids[tree_id]
        node       = TreeNode(uid, arbor=self, root=False)
        node.root  = root_node
        node._link = node_link
        return node

    def _store_node_info(self, tree_node, attr):
        """
        Store a TreeNode attribute an array for retrieval later.
        """

        self._node_info[attr][tree_node._arbor_index] = \
          getattr(tree_node, attr)

    _field_info = None
    @property
    def field_info(self):
        """
        A dictionary containing information for each available field.
        """
        if self._field_info is None and \
          self._field_info_class is not None:
            self._field_info = self._field_info_class(self)
        return self._field_info

    _size = None
    @property
    def size(self):
        """
        Return total number of trees.
        """
        if self._size is None:
            self._plant_trees()
        return self._size

    def __len__(self):
        """
        Return total number of trees.
        """
        return self.size

    _unit_registry = None
    @property
    def unit_registry(self):
        """
        Unit system registry.
        """
        if self._unit_registry is None:
            self._unit_registry = UnitRegistry()
        return self._unit_registry

    @unit_registry.setter
    def unit_registry(self, value):
        self._unit_registry = value
        self._arr = None
        self._quan = None

    _hubble_constant = None
    @property
    def hubble_constant(self):
        """
        Value of the Hubble parameter.
        """
        return self._hubble_constant

    @hubble_constant.setter
    def hubble_constant(self, value):
        self._hubble_constant = value
        if value is None:
            return
        # reset the unit registry lut while preserving other changes
        self.unit_registry = UnitRegistry.from_json(
            self.unit_registry.to_json())
        if 'h' in self.unit_registry:
            self.unit_registry.modify("h", self.hubble_constant)
        else:
            self.unit_registry.add(
                'h', self.hubble_constant, dimensionless)

    _box_size = None
    @property
    def box_size(self):
        """
        The simulation box size.
        """
        return self._box_size

    @box_size.setter
    def box_size(self, value):
        self._box_size = value
        # set unitary as soon as we know the box size
        self.unit_registry.add(
            "unitary", float(self.box_size.in_base()), length)

    def set_selector(self, selector, *args, **kwargs):
        r"""
        Sets the tree node selector to be used.

        This sets the manner in which halo progenitors are
        chosen from a list of ancestors.  The most obvious example
        is to select the most massive ancestor.

        Parameters
        ----------
        selector : string
            Name of the selector to be used.

        Any additional arguments and keywords to be provided to
        the selector function should follow.

        Examples
        --------

        >>> import ytree
        >>> a = ytree.load("rockstar_halos/trees/tree_0_0_0.dat")
        >>> a.set_selector("max_field_value", "mass")

        """
        self.selector = tree_node_selector_registry.find(
            selector, *args, **kwargs)

    _arr = None
    @property
    def arr(self):
        """
        Create a unyt_array using the Arbor's unit system.

        Parameters
        ----------

        input_array : Iterable
            A tuple, list, or array to attach units to
        units: String unit specification, unit symbol or astropy object
            The units of the array. Powers must be specified using python syntax
            (cm**3, not cm^3).
        dtype : string or NumPy dtype object
            The dtype of the returned array data

        Examples
        --------

        >>> import ytree
        >>> a = ytree.load("consistent_trees/tree_0_0_0.dat")
        >>> a.arr([1,2,3], "kg")
        unyt_array([1, 2, 3], 'kg')

        """
        if self._arr is not None:
            return self._arr
        self._arr = functools.partial(unyt_array,
                                      registry=self.unit_registry)
        return self._arr

    _quan = None
    @property
    def quan(self):
        """
        Create a unyt_quantity using the Arbor's unit system.

        Parameters
        ----------

        input_array : Iterable
            A tuple, list, or array to attach units to
        units: String unit specification, unit symbol or astropy object
            The units of the array. Powers must be specified using python syntax
            (cm**3, not cm^3).
        dtype : string or NumPy dtype object
            The dtype of the returned array data

        Examples
        --------

        >>> import ytree
        >>> a = ytree.load("consistent_trees/tree_0_0_0.dat")
        >>> a.quan(34, "cm")
        unyt_quantity(34, 'cm')

        """
        if self._quan is not None:
            return self._quan
        self._quan = functools.partial(unyt_quantity,
                                       registry=self.unit_registry)
        return self._quan

    @functools.cached_property
    def container(self):
        """
        Helper function for creating
        :class:`~ytree.data_structures.node_container.NodeContainer` objects.

        See :class:`~ytree.data_structures.node_container.NodeContainer` for more
        information.
        """
        nc = functools.partial(NodeContainer, arbor=self)
        nc.__doc__ = NodeContainer.__doc__
        return nc

    def select_halos(self, criteria, trees=None,
                     select_from=None, fields=None):
        """
        Select halos from the arbor based on a set of criteria given as a string.

        Halos matching the criteria will be returned through a generator. Matches
        are returned as soon as they are found, allowing you to begin working
        with them before the search has completed. The progress bar will update
        to report the number of matches found as the search progresses.

        Parameters
        ----------

        criteria : string
            A string that will eval to a Numpy-like selection operation
            performed on a TreeNode object called "tree".
            Example: 'tree["tree", "redshift"] > 1'
        trees : optional, list or array of TreeNodes
            A list or array of TreeNode objects in which to search. If none given,
            the search is performed over the full arbor.
        select_from : deprecated, do not use
            This keyword is no longer required and using it does nothing.
        fields : deprecated, do not use
            This keyword is no longer required and using it does nothing.

        Returns
        -------

        halos : :class:`~ytree.data_structures.tree_node.TreeNode` generator
            A generator yielding all TreeNodes meeting the criteria.

        Examples
        --------

        >>> import ytree
        >>> a = ytree.load("tree_0_0_0.dat")
        >>> for halo in a.select_halos('tree["tree", "redshift"] > 1'):
        ...     print (halo["mass"])
        >>>
        >>> halos = list(a.select_halos('tree["prog", "mass"].to("Msun") >= 1e10'))
        >>> print (len(halos))

        """

        if select_from is not None:
            import warnings
            from numpy.exceptions import VisibleDeprecationWarning
            warnings.warn(
                "The \"select_from\" keyword is deprecated and no longer does anything.",
                VisibleDeprecationWarning, stacklevel=2)

        if fields is not None:
            import warnings
            from numpy.exceptions import VisibleDeprecationWarning
            warnings.warn(
                "The \"fields\" keyword is deprecated and no longer does anything.",
                VisibleDeprecationWarning, stacklevel=2)

        tree = SelectionDetector(self)
        eval(criteria)
        if len(tree.selectors) > 1:
            raise ValueError(
                f"Selection criteria must only use one selector: \"{criteria}\".\n"
                f"    Selection criteria uses {len(tree.selectors)} selectors: "
                f"{tree.selectors}.")
        selector = tree.selectors[0]

        if trees is None:
            trees = self

        found = 0
        pbar = get_pbar(f"Selecting halos ({found} found)", trees.size)
        for i, tree in enumerate(trees):
            imatches = np.where(eval(criteria))[0]
            if imatches.size > 0:
                found += imatches.size
                if isinstance(pbar, TqdmProgressBar):
                    pbar._pbar.set_description_str(f"Selecting halos (found {found})")
            pbar.update(i+1)

            for imatch in imatches:
                yield tree.get_node(selector, imatch)

        pbar.finish()

    def add_analysis_field(self, name, units, dtype=None, default=0):
        r"""
        Add an empty field to be filled by analysis operations.

        Parameters
        ----------
        name : string
            Field name.
        units : string
            Field units.
        dtype : optional, type
            Data type for field values. If None, the default data type
            of the arbor is used.
            Default: None.
        default: optional, numeric
            Default field value when field is initialized.
            Default: 0.

        Examples
        --------

        >>> import ytree
        >>> a = ytree.load("tree_0_0_0.dat")
        >>> a.add_analysis_field("robots", "Msun * kpc")
        >>> # Set field for some halo.
        >>> my_tree = a[0]
        >>> my_tree["tree"][7]["robots"] = 1979.816
        """

        self.field_info.add_analysis_field(
            name, units,
            dtype=dtype, default=default)

    def add_alias_field(self, alias, field, units=None,
                        force_add=True):
        r"""
        Add a field as an alias to another field.

        Parameters
        ----------
        alias : string
            Alias name.
        field : string
            The field to be aliased.
        units : optional, string
            Units in which the field will be returned.
        force_add : optional, bool
            If True, add field even if it already exists and warn the
            user and raise an exception if dependencies do not exist.
            If False, silently do nothing in both instances.
            Default: True.

        Examples
        --------

        >>> import ytree
        >>> a = ytree.load("tree_0_0_0.dat")
        >>> # "Mvir" exists on disk
        >>> a.add_alias_field("mass", "Mvir", units="Msun")
        >>> print (a["mass"])

        """

        self.field_info.add_alias_field(
            alias, field, units=units, force_add=force_add)

    def add_derived_field(self, name, function,
                          units=None, dtype=None, description=None,
                          vector_components=None, force_add=True):
        r"""
        Add a field that is a function of other fields.

        Parameters
        ----------
        name : string
            Field name.
        function : callable
            The function to be called to generate the field.
            This function should take two arguments, the
            arbor and the data structure containing the
            dependent fields.  See below for an example.
        units : optional, string
            The units in which the field will be returned.
        dtype : optional, type
            The data type of the field array. If none, use the
            default type set by Arbor._default_dtype.
        description : optional, string
            A short description of the field.
        vector_components : optional, list of strings
            If given, the field is assumed to be a vector field with
            the x/y/z component fields provided here.
            Default: None.
        force_add : optional, bool
            If True, add field even if it already exists and warn the
            user and raise an exception if dependencies do not exist.
            If False, silently do nothing in both instances.
            Default: True.

        Examples
        --------

        >>> import ytree
        >>> a = ytree.load("tree_0_0_0.dat")
        >>> def _redshift(field, data):
        ...     return 1. / data["scale"] - 1
        ...
        >>> a.add_derived_field("redshift", _redshift)
        >>> print (a["redshift"])

        """

        self.field_info.add_derived_field(
            name, function,
            units=units, dtype=dtype, description=description,
            vector_components=vector_components, force_add=force_add)

    def add_vector_field(self, name, vector_components=None):
        """
        Add vector fields for a set of x,y,z component fields.

        This will add a general vector field that returns the combined
        x, y, z components as a single Nx3 array. A <field>_magnitude
        field with the quadrature sum of the components is also added.

        Parameters
        ----------
        name : string
            The name of the field. Component x,y,z fields must exist.
        vector_components : list of strings
            Used to explicitly provide the x/y/z component fields. If not
            given, a naming convention of <name>_<xyz> is assumed.
            Default: None.

        Examples
        --------

        >>> import ytree
        >>> a = ytree.load("tree_0_0_0.dat")
        >>> for ax in 'xyz':
        >>>     a.add_analysis_field(f"thing_{ax}")
        >>> fn = a.save_arbor()
        >>> a_new = ytree.load(fn)
        >>> a_new.add_vector_field("thing")
        >>> print (a_new["thing"])
        >>> print (a_new["thing_magnitude"])

        >>> import ytree
        >>> a = ytree.load("AHF_100_tiny/GIZMO-NewMDCLUSTER_0047.snap_128.parameter")
        >>> a.add_vector_field("mean_z", vector_components=["mean_z_gas", "mean_z_star"])

        """

        self.field_info.add_vector_field(name, cfields=vector_components)

    def get_yt_selection(self, *args, **kwargs):
        raise NotImplementedError(
            "This function is only implemented for ytree arbors."
            "Use save_arbor to save your data in the correct format.")

    def get_nodes_from_selection(self, *args, **kwargs):
        raise NotImplementedError(
            "This function is only implemented for ytree arbors."
            "Use save_arbor to save your data in the correct format.")

    def reload_arbor(self):
        raise NotImplementedError(
            "This function is only implemented for ytree arbors."
            "Use save_arbor to save your data in the correct format.")

    @classmethod
    def _is_valid(cls, *args, **kwargs):
        """
        Check if input file works with a specific Arbor class.
        This is used with :func:`~ytree.data_structures.arbor.load` function.
        """
        return False

    def save_arbor(self, **kwargs):
        r"""
        Save the arbor to a file.

        The saved arbor can be re-loaded as an arbor.

        Parameters
        ----------
        filename : optional, string
            Output file keyword.  If filename ends in ".h5",
            the main header file will be just that.  If not,
            filename will be <filename>/<basename>.h5.
            Default: "arbor".
        fields : optional, list of strings
            The fields to be saved.  If not given, all
            fields will be saved.
        trees : optional, list or array of TreeNodes
            If given, only save trees stemming from these nodes.
            If not provide, all trees will be saved.
        save_in_place : optional, bool or None
            If True, analysis fields will be saved to the original
            arbor, even if only a subset of all trees is provided
            with the trees keyword. This will essentially "update"
            the arbor in place. If False and only a subset of
            all trees is provided, a new arbor will be created
            containing only the trees provided. If set to None,
            behavior is determined by the type of arbor loaded.
            If the arbor is a YTreeArbor (i.e., saved with
            save_arbor), save_in_place will be set to True. If
            not of this type, it will be set to False.
            Default: None
        save_roots_only : optional, bool
            If True, only field values of each node are saved.
            If False, field data for the entire tree stemming
            from that node are saved.
            Default:  False.
        max_file_size : optional, float
            The maximum number of nodes saved to a single file.
            Smaller numbers will result in more files. Performance
            may change somewhat with different values.
            Default: 524288 (2^19).

        Returns
        -------
        header_filename : string
            The filename of the saved arbor.

        Examples
        --------

        >>> import ytree
        >>> a = ytree.load("rockstar_halos/trees/tree_0_0_0.dat")
        >>> fn = a.save_arbor()
        >>> # reload it
        >>> a2 = ytree.load(fn)

        """

        fn = save_arbor(self, **kwargs)
        return fn

class SegmentedArbor(Arbor):
    """
    Arbor subclass for multi-file datasets where an entire merger tree
    is contained within a file (i.e., no overlap). This permits the
    definition of a useful _node_io_loop_prepare function.
    """

    # Data formats organized similar to below can use this class.
    # _fi - file index, i.e., which data file is it in
    # _si - start index, the array index where this tree starts
    _node_io_attrs = ('_fi', '_si')

    def _node_io_loop_start(self, data_file):
        data_file.open()

    def _node_io_loop_finish(self, data_file):
        data_file.close()

    def _node_io_loop_prepare(self, nodes):
        if nodes is None:
            nodes = np.arange(self.size)
            fi = self._node_info['_fi']
            si = self._node_info['_si']
        elif nodes.dtype == object:
            fi = np.array(
                [node._fi if node.is_root else node.root._fi
                 for node in nodes])
            si = np.array(
                [node._si if node.is_root else node.root._si
                 for node in nodes])
        else: # assume an array of indices
            fi = self._node_info['_fi'][nodes]
            si = self._node_info['_si'][nodes]

        # the order they will be processed
        io_order = np.lexsort((si, fi))
        fi = fi[io_order]
        # array to return them to original order
        return_order = np.empty_like(io_order)
        return_order[io_order] = np.arange(io_order.size)

        ufi = np.unique(fi)
        data_files = [self.data_files[i] for i in ufi]
        index_list = [io_order[fi == i] for i in ufi]

        return data_files, index_list, return_order

class CatalogArbor(Arbor):
    """
    Base class for Arbors created from a series of halo catalog
    files where the descendent ID for each halo has been
    pre-determined.

    Unlike formats where tree information is stored in single file,
    halos are scattered about multiple catalog files. This requires
    us to store the root TreeNode objects and their full assemblies.
    """

    _prefix = None
    _data_file_class = None
    # does the dataset define unique ids?
    _has_uids = False

    # We will store root TreeNodes instead of generate them,
    # so we don't need to store anything here.
    _node_con_attrs = ()

    # Don't reset _ancestors or descendents because we won't be able to
    # rebuild trees without calling _plant_trees again.
    _setup_attrs = ("_desc_uids", "_uids", "_nodes", "_link_storage")
    _grow_attrs = ()

    def __init__(self, filename):
        super().__init__(filename)
        if not self._has_uids:
            if "uid" not in self.field_list:
                for field in "uid", "desc_uid":
                    self.field_list.append(field)
                    self.field_info[field] = {"units": "",
                                              "source": "arbor"}

    def _get_data_files(self):
        raise NotImplementedError

    def _generate_root_node(self, index):
        """
        Return a node from self._trees.

        These cannot be generated easily, so we keep them.
        """
        node = self._trees[index]
        if not hasattr(node, '_arbor_index'):
            node._arbor_index = index
        return node

    _trees = None
    @property
    def is_planted(self):
        """
        Determine if trees have been planted.
        """
        return self._trees is not None

    def _plant_trees(self):
        """
        Construct all trees.

        Since nodes are spread out over multiple files, we will
        plant all trees and create all ancestor/descendent links.

        The links will be held by the nodes themselves and we will
        not store the nodes in an array until _setup_tree is called.
        """

        if self.is_planted:
            return

        # This is a somewhat hacky way of catching halos with links
        # spanning more than one data set. That said, dict access
        # is much faster than I thought and perhaps this whole
        # routine needs to be refactored.
        if self._has_uids:
            all_dict = {}
            missed_connections = []

        # this can be called once with the list, but fields are
        # not guaranteed to be returned in order.
        if self._has_uids:
            id_fields = ["uid", "desc_uid"]
        else:
            id_fields = ["halo_id", "desc_id"]
        fields = \
          [self.field_info.resolve_field_dependencies([field])[0][0]
           for field in id_fields]
        halo_id_f, desc_id_f = fields
        dtypes = dict((field, np.int64) for field in fields)
        uid = 0
        trees = []
        nfiles = len(self.data_files)
        descs = lastids = None
        pbar = get_pbar("Planting trees", len(self.data_files))
        for i, dfl in enumerate(self.data_files):
            if not isinstance(dfl, list):
                dfl = [dfl]

            batches = []
            bsize = []
            hids = []
            ancs = defaultdict(list)
            for data_file in dfl:
                data = data_file._read_fields(fields, dtypes=dtypes)
                nhalos = len(data[halo_id_f])
                batch = np.empty(nhalos, dtype=object)

                for it in range(nhalos):
                    descid = data[desc_id_f][it]
                    if self._has_uids:
                        my_uid = data[halo_id_f][it]
                    else:
                        my_uid = uid
                    root = i == 0 or descid == -1
                    # The data says a descendent exists, but it's not there.
                    # This shouldn't happen, but it does sometimes.
                    # This can also happen when a descendent is more than
                    # one snapshot removed.
                    mcollect = False
                    if not root and descid not in lastids:
                        root = True
                        my_descid = descid
                        descid = data[desc_id_f][it] = -1
                        if self._has_uids:
                            mcollect = True
                    tree_node = TreeNode(my_uid, arbor=self, root=root)
                    tree_node._fi = it
                    tree_node.data_file = data_file
                    batch[it] = tree_node

                    if self._has_uids:
                        all_dict[my_uid] = tree_node

                    if root:
                        if mcollect:
                            tree_node._desc_uid = my_descid
                            missed_connections.append(tree_node)
                        else:
                            trees.append(tree_node)
                    else:
                        ancs[descid].append(tree_node)
                    uid += 1
                data_file.trees = batch
                batches.append(batch)
                bsize.append(batch.size)
                hids.append(data[halo_id_f])

            if i > 0:
                for descid, ancestors in ancs.items():
                    # this will not be fast
                    descendent = descs[descid == lastids][0]
                    descendent._ancestors = ancestors
                    for ancestor in ancestors:
                        ancestor._descendent = descendent

            if i < nfiles - 1:
                descs = np.empty(sum(bsize), dtype=object)
                lastids = np.empty(descs.size, dtype=np.int64)
                ib = 0
                for batch, hid, bs in zip(batches, hids, bsize):
                    descs[ib:ib+bs] = batch
                    lastids[ib:ib+bs] = hid
                    ib += bs
            pbar.update(i+1)
        pbar.finish()

        if self._has_uids:
            for node in missed_connections:
                my_desc_uid = node._desc_uid
                my_desc = all_dict[my_desc_uid]
                delattr(node, "_desc_uid")

                node._descendent = my_desc
                node.root = my_desc.root
                if my_desc._ancestors is None:
                    my_desc._ancestors = []
                my_desc._ancestors.append(node)

        self._trees = np.array(trees)
        self._size = self._trees.size

    def _setup_tree(self, tree_node):
        """
        Walk the tree and place all nodes into an array.

        This is required for field access.
        """

        if self.is_setup(tree_node):
            return

        nodes     = []
        uids      = []
        desc_uids = [-1]
        # This is redundant, but enables functionality that uses
        # the link storage, like TreeNode.get_node.
        links     = []
        for i, node in enumerate(tree_node._tree_nodes):
            node._tree_id = i
            node.root     = tree_node
            nodes.append(node)
            uids.append(node.uid)
            link = NodeLink(i)
            links.append(link)
            if i > 0:
                desc_uids.append(node.descendent.uid)
                desc_link = links[node.descendent.tree_id]
                desc_link.add_ancestor(link)
        tree_node._nodes     = np.array(nodes)
        tree_node._uids      = np.array(uids)
        tree_node._desc_uids = np.array(desc_uids)
        tree_node._tree_size = tree_node._uids.size
        tree_node._link_storage = np.array(links)
        # This should bypass any attempt to get this field in
        # the conventional way.
        if self.field_info["uid"].get("source") == "arbor":
            tree_node.field_data["uid"] = tree_node._uids
            tree_node.field_data["desc_uid"] = tree_node._desc_uids

    def _grow_tree(self, tree_node):
        """
        Trees are grown when they are planted.
        """
        pass
