import pandas as pd
from lmwrapper.openai_wrapper import get_open_ai_lm

from synthegrator.code_solver import LmCodeSolverAutoRegressive
from synthegrator.df_converters import solution_evals_to_df
from synthegrator.execution_threading import solve_and_evaluate_problems
from synthegrator.synthdatasets.human_eval import yield_human_eval


def main():
    problems = list(yield_human_eval())[:10]
    lm = get_open_ai_lm()
    solver = LmCodeSolverAutoRegressive(lm, include_lm_response=True)
    evals = solve_and_evaluate_problems(
        solver,
        problems,
        max_threads_solve=8,
        max_threads_eval=8,
    )
    df = solution_evals_to_df(evals)
    # make pandas no max width
    with pd.option_context("display.max_colwidth", None):
        print(df.columns)
        print(df.solution__lm_prediction__completion_logprobs)
        print(
            list(
                zip(
                    df.solution__lm_prediction__completion_logprobs.iloc[0],
                    df.solution__lm_prediction__completion_tokens.iloc[0],
                    strict=False,
                ),
            ),
        )


if __name__ == "__main__":
    main()
