from meta_researcher.src.reward_score import qa_search_and_format, qa_search_and_format_tool

def _default_batch_compute_score_format_answer(data_source_list, solution_str_list, ground_truth_list, extra_info_list, turn_list, inference_turn_list, think_len_threshold, think_len_belta, max_step_num):
    reward, format_reward, answer_reward, think_len, think_reward = qa_search_and_format_tool.batch_compute_score_format_answer(data_source_list, solution_str_list, ground_truth_list, extra_info_list, turn_list, inference_turn_list, think_len_threshold, think_len_belta, max_step_num)
    
    reward_list = []
    format_reward_list = []
    answer_reward_list = []
    think_len_list = []
    think_reward_list = []
    if isinstance(reward, list):
        for one_res in reward:
            if isinstance(one_res, (int, float, bool)):
                reward_list.append(float(one_res))
            else:
                reward_list.append(float(one_res[0]))

    if isinstance(format_reward, list):
        for one_res in format_reward:
            if isinstance(one_res, (int, float, bool)):
                format_reward_list.append(float(one_res))
            else:
                format_reward_list.append(float(one_res[0]))

    if isinstance(answer_reward, list):
        for one_res in answer_reward:
            if isinstance(one_res, (int, float, bool)):
                answer_reward_list.append(float(one_res))
            else:
                answer_reward_list.append(float(one_res[0]))

    if isinstance(think_len, list):
        for one_res in think_len:
            if isinstance(one_res, (int, float, bool)):
                think_len_list.append(float(one_res))
            else:
                think_len_list.append(float(one_res[0]))

    if isinstance(think_reward, list):
        for one_res in think_reward:
            if isinstance(one_res, (int, float, bool)):
                think_reward_list.append(float(one_res))
            else:
                think_reward_list.append(float(one_res[0]))

    # return reward_list, format_reward_list, answer_reward_list
    return {
            "reward_list": reward_list,
            "answer_reward_list": answer_reward_list,
            "format_reward_list": format_reward_list,
            "think_len_list": think_len_list,
            "think_reward_list": think_reward_list,
        }