from synthegrator.code_problem_builders import make_simple_method_completion_problem
from synthegrator.few_shotting import Bm25FewShotLibrary


def test_bm25_few_shot_library():
    def make_prob(text):
        return make_simple_method_completion_problem(
            prompt=text, signature="def foo(x):", test_cases=[]
        )

    problems = [
        make_prob("Return the average of a list of numbers."),
        make_prob("Return the sum of a list of numbers."),
        make_prob("Return the maximum of a list of numbers."),
        make_prob("Return the minimum of a list of numbers."),
        make_prob("Return the median of a list of numbers."),  # -5
        make_prob("Return the median of a list."),  # -4
        make_prob("Return the median of a list of."),  # -3
        make_prob("Return a number that makes happy"),  # -2
        make_prob("Return a number that makes very happy"),  # -1
    ]
    library = Bm25FewShotLibrary(problems)
    sample = library.sample_problems(1, problems[-1])
    assert len(sample) == 1
    assert sample[0] == problems[-2]
    print(sample[0].working_directory.files.get_only_file().content_str)

    library = Bm25FewShotLibrary(problems)
    sample = library.sample_problems(2, problems[-4])
    assert len(sample) == 2
    assert sample[0] in (problems[-3], problems[-5])
    assert sample[1] in (problems[-3], problems[-5])
