#!/usr/bin/env python
# -*- coding: utf-8 -*-
import seutils
import os.path as osp

parser = seutils.ParserMultipleLFNs()
parser.add_argument('-t', '--tree', type=str, help='Name of the tree to count')
parser.add_argument(
    '-f', '--force', action='store_true',
    help='Forces a recount of events (i.e. cache is ignored)'
    )
parser.add_argument(
    '-n', '--nocache', action='store_true',
    help='Disables using the cache (both reading and writing)'
    )
args = parser.parse_args()

def main():
    per_dirname = {}
    for rootfile in args.lfns:
        dirname = osp.dirname(rootfile)
        if dirname not in per_dirname: per_dirname[dirname] = []
        per_dirname[dirname].append(rootfile)

    if args.nocache:
        read_cache = False
        write_cache = False
    else:
        read_cache = not(args.force)
        write_cache = True

    counts_per_tree = {}
    for dirname, rootfiles_in_dir in per_dirname.items():
        counts_per_rootfile = seutils.root.count_dataset(
            dirname, rootfiles=rootfiles_in_dir,
            read_cache=read_cache, write_cache=write_cache
            )
        for rootfile, counts in counts_per_rootfile.items():
            print(rootfile)
            for tree, count in counts.items():
                if (args.tree and tree != args.tree): continue
                print('  {:5} {}'.format(count, tree))
                counts_per_tree.setdefault(tree, 0)
                counts_per_tree[tree] += count

    for tree, count in sorted(counts_per_tree.items()):
        if (args.tree and tree != args.tree): continue
        print(
            '\033[31mTTree {0} ({1} entries)\033[0m'
            .format(tree, count)
            )

if __name__ == '__main__':
    main()