#!/usr/bin/env python3
#
# Author: Yipeng Sun <syp at umd dot edu>
# License: BSD 2-clause
# Last Change: Sat Oct 03, 2020 at 12:43 AM +0800

import yaml

from argparse import ArgumentParser
from io import StringIO
from pyBabyMaker.io.TupleDump import PyTupleDump


#################################
# Command line arguments parser #
#################################

def parse_input():
    parser = ArgumentParser(description='''
dump ntuple tree-branch structure, and branch datatype.''')

    parser.add_argument('input',
                        help='''
specify path to input ntuple file.''')

    parser.add_argument('output',
                        nargs='?',
                        default=None,
                        help='''
specify path to output YAML file.''')

    parser.add_argument('-f', '--output-format',
                        dest='formatter',
                        nargs='?',
                        choices=['yaml', 'plain'],
                        default='yaml',
                        help='''
select output format.''')

    return parser.parse_args()


####################
# Output formatter #
####################

def yaml_formatter(data, output):
    yaml.dump(data, output, default_flow_style=False)


########
# Main #
########

if __name__ == '__main__':
    args = parse_input()

    dumper = PyTupleDump(args.input)
    data = dumper.dump()
    content = StringIO()

    if args.formatter == 'yaml':
        yaml_formatter(data, content)

    if args.output:
        with open(args.output, 'w') as f:
            print(content.getvalue(), file=f)
    else:
        print(content.getvalue())
