from ast import literal_eval

import datasets

if __name__ == "__main__":
    TRAINING_TRACES = True

    gaia_ds = datasets.load_dataset("m-ric/agents_medium_benchmark_2")["train"]
    gaia_ds = gaia_ds.filter(lambda row: row["source"] == "GAIA")

    math_ds = datasets.load_dataset("HuggingFaceH4/MATH-500")["test"]

    def filter_answer_float(row):
        try:
            float(row["answer"])
            return True
        except Exception:
            return False

    math_ds = math_ds.shuffle(seed=42)

    if TRAINING_TRACES:
        selection_range = range(50, 323)
    else:
        selection_range = range(50)
    math_ds = math_ds.filter(filter_answer_float).select(selection_range)
    math_ds = math_ds.rename_column("solution", "true_reasoning")
    math_ds = math_ds.rename_column("answer", "true_answer")
    math_ds = math_ds.rename_column("problem", "question")
    math_ds = math_ds.remove_columns(["subject", "level", "unique_id"])

    math_ds = math_ds.add_column("source", ["MATH"] * len(math_ds))

    simpleqa_ds = datasets.load_dataset("basicv8vc/SimpleQA")["test"]

    def filter_answer_float(row):
        try:
            float(row["answer"])
            if row["answer"].endswith("."):
                return False
            return True
        except Exception:
            return False

    simpleqa_ds = simpleqa_ds.shuffle(seed=42)

    if TRAINING_TRACES:
        selection_range = range(50, 250)
    else:
        selection_range = range(50)
    simpleqa_ds = simpleqa_ds.filter(filter_answer_float).select(selection_range)
    simpleqa_ds = simpleqa_ds.rename_column("answer", "true_answer")
    simpleqa_ds = simpleqa_ds.rename_column("problem", "question")
    simpleqa_ds = simpleqa_ds.add_column(
        "true_reasoning", [str(literal_eval(el["metadata"])["urls"]) for el in simpleqa_ds]
    )
    simpleqa_ds = simpleqa_ds.remove_columns(["metadata"])

    simpleqa_ds = simpleqa_ds.add_column("source", ["SimpleQA"] * len(simpleqa_ds))

    if TRAINING_TRACES:
        datasets_to_add = [math_ds, simpleqa_ds]
        eval_ds = datasets.concatenate_datasets(datasets_to_add)
        eval_ds.push_to_hub("smolagents/trace-generation-tasks")
    else:
        datasets_to_add = [gaia_ds, math_ds, simpleqa_ds]
        eval_ds = datasets.concatenate_datasets(datasets_to_add)
        eval_ds.push_to_hub("m-ric/smol_agents_benchmark")
