#! /usr/bin/env python3

import bdt2cpp
import argparse
import sys

p = argparse.ArgumentParser()
p.add_argument('--print-tree', action='store_true', help='''print the parsed
               tree''')
p.add_argument('model', type=str, help='''filename of the BDT model you
                want to transpile. XGBoost, TMVA and pickled XGBoost BDTs are
                supported. If the file ending is `.xgb` or `.pkl`, bdt2cpp will
                use XGBoost or pickled XGBoost, respectively. If the file
                ending is `.xml`, bdt2cpp will use the TMVA parser.
                You may also specify the type via the --type argument.''')
p.add_argument('-o', '--output-directory', type=str, default='build', help='''
               output directory to which the transpiled template will be
               written. The files inside will be overwritten without warning!
               Defaults to `build`.''')
p.add_argument('-m', '--max-trees', type=int, default=None, help='''maximum
               number of trees per file''')
p.add_argument('-n', '--feature-names', nargs='*', help='''The list of all
               feature names, if the model uses named features instead of
               indexed features. This list must be of the exact same length as
               the length of the feauture vector. The names need to be
               separated by spaces like `-n f1 f2 f3`.''')
p.add_argument('-t', '--type', type=str, default=None, help='''The BDT input
               type. Can be one of ['TMVA', 'XGBoost', 'Pickle'].
               Overrides the type deduced from the file ending.''')

args = p.parse_args()

if args.print_tree:
    print(*bdt2cpp.parse_model(args.model), sep='\n\n')
    sys.exit()

bdt2cpp.main(args.model, output_dir=args.output_directory,
             trees_per_file=args.max_trees, feature_names=args.feature_names,
             bdt_type=args.type)
