#!/usr/bin/env python3

from __future__ import print_function

from math import log


'''
A script to read in a list of numbers and output an ASCII histogram of them.
Various optional arguments control the format of the output.
'''


def textogram(vals, **kwargs):
    print(get_textogram(vals, **kwargs))

def linspace(a, b, n):
    dx = float(b-a) / (n-1)
    return [a + i * dx for i in range(n)]

def get_textogram(vals, bins=None, N=10, xmin=None, xmax=None,
        with_counts=True, yscale='lin', height=60, char='#',
        categorical=False):
    if len(vals) == 0:
        return 'No data.'
    if categorical:
        categories = sorted(set(vals))
        bin_counts = [len([v for v in vals if v == c]) for c in categories]
        bin_vals = hist_categorical(vals, categories)
        labels = [str(c) for c in categories]
        max_label_length = max(map(len, labels))
        labels = [('{: >' + str(max_label_length) + '}').format(l) for l in labels]
    else:
        vals = [float(v) for v in vals]
        v_min = xmin if xmin is not None else min(vals)
        v_max = xmax if xmax is not None else max(vals)
        left_edges = linspace(v_min, v_max, N+1) if bins is None else bins
        left_tail = [v for v in vals if v < left_edges[0]]
        right_tail = [v for v in vals if v > left_edges[-1]]
        vals = [v for v in vals if v >= left_edges[0] and v <= left_edges[-1]]
        bin_vals = hist(vals, left_edges)
        bin_counts = list(bin_vals)
        labels = get_bin_labels(left_edges, left_tail=len(left_tail)>0, right_tail=len(right_tail)>0)
    if yscale == 'log':
        for i in range(len(bin_vals)):
            if bin_vals[i] > 0:
                bin_vals[i] = log(bin_vals[i], 2) + 1
    n_tot = len(vals)
    s = '\n'
    bin_val_max = max(bin_vals)
    max_bar_height = height
    sep = ' - '
    pct_fmt = '{: >3d}'
    if not categorical and len(left_tail) > 0:
        s += labels.pop(0) + ' : ' + format_percentage(len(left_tail)/float(n_tot), fmt='{:d}') + ' (' + str(len(left_tail)) + ')\n'
    for i in range(len(bin_counts)):
        prefix = labels[i] + ' : '
        prefix_length = len(prefix)
        s += prefix
        bar_height = int(max_bar_height * bin_vals[i] / bin_val_max)
        s += char * bar_height + ' ' * (max_bar_height - bar_height)
        if with_counts and bin_counts[i] > 0:
            s += ' ' * 2 + format_percentage(bin_counts[i]/float(n_tot)) + ' (' + str(bin_counts[i]) + ')'
        s += '\n'
    if not categorical and len(right_tail) > 0:
        s += labels[-1] + ' : ' + format_percentage(len(right_tail)/float(n_tot), fmt='{:d}') + ' (' + str(len(right_tail)) + ')\n'
    return s

def hist_categorical(vals, categories):
    res = []
    for i in range(len(categories)):
        res.append(len([v for v in vals if v == categories[i]]))
    return res

def hist(vals, left_edges):
    res = []
    for i in range(len(left_edges)-1):
        items = [v for v in vals if v >= left_edges[i] and v < left_edges[i+1]]
        res.append(len(items))
    return res

def format_percentage(r, fmt='{: >3d}'):
    return fmt.format(int(100 * r)) + '%'

def get_bin_labels(edges, left_tail=False, right_tail=False):
    w = max(map(lambda x: len(str(int(x))), edges))
    min_step = edges[-1] - edges[0]
    for i in range(len(edges)-1):
        min_step = min(min_step, edges[i+1] - edges[i])
    if min_step < 1:
        p = -int('{:e}'.format(min_step).split('e')[1])
    else:
        p = 0
    res = []
    fmt_suffix = str(w+p) + '.' + str(p) + 'f}'
    for i in range(len(edges)-1):
        label = ('{: >' + fmt_suffix).format(edges[i])
        label += ' - '
        label += ('{: <' + fmt_suffix).format(edges[i+1])
        res.append(label)
    if left_tail:
        label = '<' + ('{:' + fmt_suffix).format(edges[0])
        label = ' ' * (len(res[0]) - len(label)) + label
        res.insert(0, label)
    if right_tail:
        label = '>' + ('{:' + fmt_suffix).format(edges[-1])
        label = ' ' * (len(res[0]) - len(label)) + label
        res.append(label)
    return res

def main(args):
    with open(args.infile, 'rb') as fp:
        vals = []
        for line in fp:
            if not args.everything:
                try:
                    vals.append(float(line.strip()))
                except ValueError:
                    continue
            else:
                vals.append(line.strip())
    kwargs = {
            'bins': args.bins,
            'categorical': args.categorical,
            'char': args.char,
            'height': args.height,
            'N': args.num_bins,
            'with_counts': args.with_counts,
            'xmin': args.min,
            'xmax': args.max,
            'yscale': args.yscale,
            }
    textogram(vals, **kwargs)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('infile', nargs='?', default='/dev/stdin')
    parser.add_argument('--bins', '-b', nargs='+')
    parser.add_argument('--categorical', '-c', action='store_true')
    parser.add_argument('--char', '-C', default='#')
    parser.add_argument('--everything', '-e', action='store_true')
    parser.add_argument('--height', type=int, default=50)
    parser.add_argument('--min', '-a', type=float)
    parser.add_argument('--max', '-z', type=float)
    parser.add_argument('--num-bins', '-n', type=int, default=10)
    parser.add_argument('--yscale', '-y', choices=['lin', 'log'], default='lin')
    parser.add_argument('--with-counts', action='store_false')
    args = parser.parse_args()
    main(args)

