#!/usr/bin/env python

import sys
import uproot
import numpy as np

from argparse import ArgumentParser
from itertools import product
from tabulate import tabulate


################################
# Command line argument parser #
################################

def parse_input(descr='print histogram content, with errors.'):
    parser = ArgumentParser(descr)

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

    parser.add_argument('-H', '--histo', required=True,
                        help='specify name of histogram.')

    parser.add_argument('-p', '--precision', default=4, type=int,
                        help='specify float print precision.')

    parser.add_argument('-t', '--tablefmt', default='fancy_grid',
                        help='specify table format.')

    parser.add_argument('-n', '--no-headers', action='store_true',
                        help='remove headers and indices.')

    parser.add_argument('-f', '--flow', action='store_true',
                        help='include under/overflow.')

    return parser.parse_args()


#####################
# Histogram printer #
#####################

def make_data(values, errors, precision):
    data = np.empty(values.shape, dtype=object)
    idx = [list(range(i)) for i in values.shape]

    for i in product(*idx):
        data[i] = '{0:,.{1}f}'.format(values[i], precision) + '±' + \
            '{0:,.{1}f}'.format(errors[i], precision)

    return data.tolist()


def print_cell_range(lower, upper, precision=2):
    return '{0:,.{1}f}'.format(lower, precision) + '\n' + '{0:,.{1}f}'.format(
        upper, precision)


def print_2d(values, errors, xbins, ybins,
             tablefmt='simple', precision=4, no_headers=False, legend='x / y'):
    header = [legend] + [print_cell_range(i, j)
                         for i, j in zip(ybins[:-1], ybins[1:])]
    index = [print_cell_range(i, j) for i, j in zip(xbins[:-1], xbins[1:])]
    data = make_data(values, errors, precision)

    if no_headers:
        print(tabulate(data, tablefmt=tablefmt))
    else:
        print(
            tabulate(data, headers=header, showindex=index, tablefmt=tablefmt))


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

if __name__ == '__main__':
    args = parse_input()
    ntp = uproot.open(args.ntp)

    try:
        histo = ntp[args.histo]
    except Exception:
        print(f'"{args.histo}" is not in {args.ntp}! Available histos:')
        for i in ntp:
            print(f'  {i}')
        sys.exit(1)

    values, errors = histo.values(args.flow), histo.errors(args.flow)
    _, *binning = histo.to_numpy(args.flow)

    if len(binning) == 1:
        print('1D histogram not supported!')

    elif len(binning) == 2:
        print()
        print_2d(values, errors, binning[0], binning[1],
                 args.tablefmt, args.precision, args.no_headers)
        print()

    elif len(binning) == 3:
        for v, e, low, high in \
                zip(values, errors,
                    binning[0][:-1], binning[0][1:]):
            if not args.no_headers:
                print()
                print(f'## x ∈ [{low}, {high})')
            print()
            print_2d(v, e, binning[1], binning[2],
                     args.tablefmt, args.precision, args.no_headers,
                     legend='y / z')
            print()

    else:
        print(f'{len(binning)}D histogram not supported!')
