import io


def latexify_results(eval_data, id_column='run_id'):
    """Take a data frame produced by `EvaluationProtocol.eval_data()` and
    produce a LaTeX table of results for this method. Will use the `run_id`
    column as an algorithm name (or to get algorithm names, if there's more
    than one algorithm present in the given data). You can override that by
    specifying the `id_column` keyword argument."""

    # Each column of the LaTeX table corresponds to a particular evaluation
    # environment, while each row corresponds to an algorithm. In contrast,
    # each row of the given Pandas frame is represents a series of rollouts by
    # one particular algorithm on one particular test configuration.

    test_envs = eval_data['test_env'].unique()
    col_names = [r'\textbf{%s}' % e for e in test_envs]
    alg_names = eval_data[id_column].unique()

    # write to buffer so we can use print()
    fp = io.StringIO()

    # prefix is just LaTeX table setup
    print(r"\centering", file=fp)
    print(r"\begin{tabular}{l@{\hspace{1em}}%s}" % ("c" * len(col_names)),
          file=fp)
    print(r"\toprule", file=fp)

    # first line: env names
    print(r'\textbf{Randomisation} & ', end='', file=fp)
    print(' & '.join(col_names), end='', file=fp)
    print('\\\\', file=fp)
    print(r'\midrule', file=fp)

    # next lines: actual results
    for alg_name in alg_names:
        alg_mask = eval_data[id_column] == alg_name
        stat_parts = []
        for env_name in test_envs:
            full_mask = alg_mask & (eval_data['test_env'] == env_name)
            relevant_rows = list(eval_data[full_mask].iterrows())
            if len(relevant_rows) != 1:
                raise ValueError(
                    f"got {len(relevant_rows)} rows corresponding to "
                    f"{id_column}={alg_name} and test_env={env_name}, but "
                    f"expected one (maybe IDs in column {id_column} aren't "
                    f"unique?)")
            (_, row), = relevant_rows
            std = row['std_score']
            stat_parts.append(f'{row["mean_score"]:.2f} ($\\pm$ {std:.2f})')
        print(r'\textbf{%s} & ' % alg_name, end='', file=fp)
        print(' & '.join(stat_parts), end='', file=fp)
        print('\\\\', file=fp)
        print(r'\bottomrule', file=fp)
        print(r'\end{tabular}', file=fp)

    return fp.getvalue()