# -*- coding: utf-8 -*-

########################################################################
#
# License: BSD
# Created: November 8, 2014
# Author:  Alistair Muldal - alimuldal@gmail.com
#
# $Id$
#
########################################################################

"""This utility prints the contents of an HDF5 file as a tree.

Pass the flag -h to this for help on usage.

"""
from __future__ import absolute_import
from __future__ import print_function

import tables
import numpy as np
import os
import argparse
from collections import defaultdict, deque
import warnings
from six.moves import zip


def _get_parser():
    parser = argparse.ArgumentParser(
        description='''
        `pttree` is designed to give a quick overview of the contents of a
        PyTables HDF5 file by printing a depth-indented list of nodes, similar
        to the output of the Unix `tree` function.

        It can also display the size, shape and compression states of
        individual nodes, as well as summary information for the whole file.

        For a more verbose output (including metadata), see `ptdump`.
        ''')

    parser.add_argument(
        '-L', '--max-level', type=int, dest='max_depth',
        help='maximum branch depth of tree to display (-1 == no limit)',
    )
    parser.add_argument(
        '-S', '--sort-by', type=str, dest='sort_by',
        help='artificially order nodes, can be either "size", "name" or "none"'
    )
    parser.add_argument(
        '--print-size', action='store_true', dest='print_size',
        help='print size of each node/branch',
    )
    parser.add_argument(
        '--no-print-size', action='store_false', dest='print_size',
    )
    parser.add_argument(
        '--print-shape', action='store_true', dest='print_shape',

        help='print shape of each node',
    )
    parser.add_argument(
        '--no-print-shape', action='store_false', dest='print_shape',
    )
    parser.add_argument(
        '--print-compression', action='store_true', dest='print_compression',
        help='print compression library(level) for each compressed node',
    )
    parser.add_argument(
        '--no-print-compression', action='store_false',
  
      dest='print_compression',
    )
    parser.add_argument(
        '--print-percent', action='store_true', dest='print_percent',
        help='print size of each node as a %% of the total tree size on disk',
    )
    parser.add_argument(
        '--no-print-percent', action='store_false',
        dest='print_percent',
    )
    parser.add_argument(
        '--use-si-units', action='store_true', dest='use_si_units',
        help='report sizes in SI units (1 MB == 10^6 B)',
    )
    parser.add_argument(
        '--use-binary-units', action='store_false', dest='use_si_units',
        help='report sizes in binary units (1 MiB == 2^20 B)',
    )

    parser.add_argument('src', metavar='filename[:nodepath]',
                        help='path to the root of the tree structure')

    parser.set_defaults(max_depth=1, sort_by="size", print_size=True,
                        print_percent=True, print_shape=False,
                        print_compression=False, use_si_units=False)

    return parser


def main():

    parser = _get_parser()
    args = parser.parse_args()

    # Catch the files passed as the last arguments
    src = args.__dict__.pop('src').split(':')
    if len(src) == 1:
        filename, nodename = src[0], "/"
    else:
        filename, nodename = src
        if nodename == "":
            # case where filename == "filename:" instead of "filename:/"
            nodename = "/"

    with tables.open_file(filename, 'r') as f:
        tree_str = get_tree_str(f, nodename, **args.__dict__)
        print(tree_str)

    pass


def get_tree_str(f, where='/', max_depth=-1, print_class=True,
                 print_size=True, print_percent=True, print_shape=False,
                 print_compression=False, print_total=True, sort_by=None,
                 use_si_units=False):
    """
    Generate the ASCII string representing the tree structure, and the summary
    info (if requested)
    """

    root = f.get_node(where)
    root._g_check_open()
    start_depth = root._v_depth
    if max_depth < 0:
        max_depth = os.sys.maxint

    b2h = bytes2human(use_si_units)

    # we will pass over each node in the tree twice

    # on the first pass we'll start at the root node and recurse down the
    # branches, finding all of the leaf nodes and calculating the total size
    # over all tables and arrays
    total_in_mem = 0
    total_on_disk = 0
    total_items = 0

    # defaultdicts for holding the cumulative branch sizes at each node
    in_mem = defaultdict(lambda: 0.)
    on_disk = defaultdict(lambda: 0.)
    leaf_count = defaultdict(lambda: 0)

    # keep track of node addresses within the HDF5 file so that we don't count
    # nodes with multiple references (i.e. hardlinks) multiple times
    ref_count = defaultdict(lambda: 0)
    ref_idx = defaultdict(lambda: 0)
    hl_addresses = defaultdict(lambda: None)
    hl_targets = defaultdict(lambda: '')

    stack = deque(root)
    leaves = deque()

    while stack:

        node = stack.pop()

        if isinstance(node, tables.link.Link):
            # we treat links like leaves, except we don't dereference them to
            # get their sizes or addresses
            leaves.append(node)
            continue

        path = node._v_pathname
        addr, rc = node._get_obj_info()
        ref_count[addr] += 1
        ref_idx[path] = ref_count[addr]
        hl_addresses[path] = addr

        if isinstance(node, tables.Leaf):

            # only count the size of a hardlinked leaf the first time it is
            # visited
            if ref_count[addr] == 1:

                try:
                    m = node.size_in_memory
                    d = node.size_on_disk

                    # size of this node
                    in_mem[path] += m
                    on_disk[path] += d
                    leaf_count[path] += 1

                    # total over all nodes
                    total_in_mem += m
                    total_on_disk += d
                    total_items += 1

                    # arbitrarily treat this node as the 'target' for all other
                    # hardlinks that point to the same address
                    hl_targets[addr] = path

                except NotImplementedError as e:
                    # size_on_disk is not implemented for VLArrays
                    warnings.warn(e.message)

            # push leaf nodes onto the stack for the next pass
            leaves.append(node)

        elif isinstance(node, tables.Group):

            # don't recurse down the same hardlinked branch multiple times!
            if ref_count[addr] == 1:
                stack.extend(list(node._v_children.values()))
                hl_targets[addr] = path

            # if we've already visited this group's address, treat it as a leaf
            # instead
            else:
                leaves.append(node)


    # on the second pass we start at each leaf and work upwards towards the
    # root node, computing the cumulative size of each branch at each node, and
    # instantiating a PrettyTree object for each node to create an ASCII
    # representation of the tree structure

    # this will store the PrettyTree objects for every node we're printing
    pretty = {}

    stack = leaves

    while stack:

        node = stack.pop()
        path = node._v_pathname

        parent = node._v_parent
        parent_path = parent._v_pathname

        # cumulative size at parent node
        in_mem[parent_path] += in_mem[path]
        on_disk[parent_path] += on_disk[path]
        leaf_count[parent_path] += leaf_count[path]

        depth = node._v_depth - start_depth

        # if we're deeper than the max recursion depth, we print nothing
        if not depth > max_depth:

            # create a PrettyTree representation of this node
            name = node._v_name
            if print_class:
                name += " (%s)" % node.__class__.__name__

            labels = []
            pct = 100 * on_