"""
Report benchmark results from CSV files in Markdown format.
"""

import csv
import itertools
import operator
import string


def unbreak(s):
    return s.replace(' ', '\N{NO-BREAK SPACE}')


def concat_files(csv_files):
    for csv_file in csv_files:
        with open(csv_file) as f:
            yield from f


def read_rows(csv_rows):
    # CSV Format:
    # benchmark_module, benchmark_name, params, revision_name, pyversion, flags, bm_time, diff
    reader = csv.reader(csv_rows)

    # Sort by benchmark name.
    rows = sorted(reader, key=operator.itemgetter(0, 1, 2))
    return rows


def time_in_seconds(time_string):
    units = {"nsec": 1e-9, "usec": 1e-6, "msec": 1e-3, "sec": 1.0}
    number, unit = time_string.split()
    return float(number) * units[unit]


def warn_difference(reference, value, max_margin):
    difference = abs(value - reference) / reference
    is_better = value < reference

    if difference < max_margin:
        return ''
    if difference < max_margin * 2:
        return ' \N{BLACK MEDIUM DOWN-POINTING TRIANGLE}' if is_better else ' \N{UP-POINTING RED TRIANGLE}'
    if difference < max_margin * 3:
        return ' \N{LARGE GREEN CIRCLE}' if is_better else ' \N{LARGE RED CIRCLE}'
    else:
        return ' \N{LARGE GREEN SQUARE}' if is_better else ' \N{LARGE RED SQUARE}'


def format_timings(tmin, diff, *, master_data=None, warn_margin=.1/3.):
    warn = warn_difference(time_in_seconds(master_data[0]), time_in_seconds(tmin), warn_margin) if master_data else ''
    diff_str = f" ({unbreak(diff.strip(' ()'))})" if diff else ''
    return f"{unbreak(tmin)}{diff_str}{warn}"


def format_sizes(size, diff, *, master_data=None, warn_margin=.01):
    warn = warn_difference(int(master_data[0]), int(size), warn_margin) if master_data else ''
    diff_str = f" ({unbreak(diff.strip(' ()'))})" if diff else ''
    return f"{size}{diff_str}{warn}"


def build_table(rows, title, data_formatter):
    # Collect all revision names and Python versions, keeping their original order.
    # (The set may not be the same for all benchmarks.)
    revisions = list({row[3]: 1 for row in rows})
    python_versions = list({row[4]: 1 for row in rows})

    # Prepare table column mapping and header.
    pos = itertools.count(1)
    column_map = {
        (pyversion, revision, flags):  next(pos)
        for pyversion in python_versions
        for flags in ('', 'L')
        for revision in revisions
    }
    header = [title] + [
        f"Py{pyversion}{'-' if flags else ''}{flags}: {revision.replace('origin/', '')[:16]}"
        for (pyversion, revision, flags) in column_map
    ]
    row_template = [''] * len(header)

    # For each benchmark, report all timings in separate columns.
    table = []
    empty_column_indices = set(column_map.values())
    for (bm_module, bm_name, bm_params), bm_rows in itertools.groupby(rows, key=operator.itemgetter(0, 1, 2)):
        row = row_template[:]
        table.append(row)

        bm_rows = [
            (pyversion, revision_name, flags, pyversion + flags, data)
            for _, _, _, revision_name, pyversion, flags, *data in bm_rows
        ]
        master_data_seen = {
            version_key: data
            for _, revision_name, _, version_key, data in bm_rows
            if 'master' in revision_name
        }

        bm_module = bm_module[6:] if bm_module.startswith('bench_') else bm_module
        row[0] = f"{bm_module}:{bm_name}({bm_params.strip()})"
        for pyversion, revision_name, flags, version_key, data in bm_rows:
            column_index = column_map[(pyversion, revision_name, flags)]
            empty_column_indices.discard(column_index)
            master_data = master_data_seen.get(version_key) if 'HEAD' in revision_name else None
            row[column_index] = data_formatter(*data, master_data=master_data)

    # Strip empty columns, highest to lowest.
    for column_index in sorted(empty_column_indices, reverse=True):
        del header[column_index]
        for row in table:
            del row[column_index]

    return header, table


def generate_markdown(header, table):
    # Size the table columns.
    column_lengths = [
        max(map(len, map(operator.itemgetter(i), itertools.chain([header], table))))
        for i in range(len(header))
    ]

    # Generate Markdown formatted table lines.
    row_format = ("| {:<%ds}" + " | {:>%ds}" * (len(column_lengths) - 1) + " |\n") % tuple(column_lengths)
    format_row = row_format.format

    yield format_row(*header)
    yield format_row(*['-' * length for length in column_lengths])
    yield from itertools.starmap(format_row, table)


def parse_options(args):
    from argparse import ArgumentParser, RawDescriptionHelpFormatter
    parser = ArgumentParser(
        description="Report benchmark numbers as markdown tables.",
        formatter_class=RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "-t", "--type",
        dest="type", default='timings', choices=['timings', 'sizes'],
        help="The type of report.",
    )
    parser.add_argument(
        "csv_files",
        nargs="*", default=[],
        help="The CSV files to collect data from.",
    )

    return parser.parse_args(args)


def main(args):
    options = parse_options(args)

    rows = read_rows(concat_files(options.csv_files))

    if options.type == 'timings':
        title = "Benchmark timings"
        data_formatter = format_timings
    else:
        title = 'Module sizes'
        data_formatter = format_sizes

    header, table = build_table(rows, title, data_formatter)
    for line in generate_markdown(header, table):
        print(line, end='')


if __name__ == "__main__":
    import sys
    main(sys.argv[1:])
