import os
import json
import random

induce_data_path = os.path.join(os.path.dirname(__file__), 'raw/induce/')
eval_data_path = os.path.join(os.path.dirname(__file__), 'raw/execute/')

bb_induce_data_path = os.path.join(os.path.dirname(__file__), '../bbh/induce/')
bb_eval_data_path = os.path.join(os.path.dirname(__file__), '../bbh/execute/')
# Get a list of tasks (by looking at the names of the files in the induced directory)
tasks = [f.split('.')[0] for f in os.listdir(induce_data_path)]
#'object_counting','word_sorting', 'causal_judgement'放在instruction_induction中了
# bb_tasks = ["dyck_languages", "tense", "gender_inclusive_sentences_german", "operators"]
bbh_tasks = [
    'boolean_expressions', 'date_understanding', 'disambiguation_qa',
    'dyck_languages', 'formal_fallacies', 'geometric_shapes', 'hyperbaton',
    'logical_deduction_five_objects', 'logical_deduction_seven_objects', 'logical_deduction_three_objects',
    'movie_recommendation', 'multistep_arithmetic_two', 'navigate',
    'penguins_in_a_table', 'reasoning_about_colored_objects', 'ruin_names',
    'salient_translation_error_detection', 'snarks', 'sports_understanding', 'temporal_sequences',
    'tracking_shuffled_objects_five_objects', 'tracking_shuffled_objects_seven_objects',
    'tracking_shuffled_objects_three_objects', 'web_of_lies'
]
def load_data(type, task):
    if task in bbh_tasks:
        print(f"bbh task: {task}")
        base_path = bb_induce_data_path if type == 'induce' else bb_eval_data_path
    else:
        base_path = induce_data_path if type == 'induce' else eval_data_path
    path = base_path + task + '.json'
    with open(path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    examples = data['examples']

    msg = '' if 'msgs' not in data else data['msgs']

    num_examples = len(examples)

    inputs, outputs = [], []
    for i in range(num_examples):
        data = examples[str(i + 1)]
        if task == 'cause_and_effect':
            cause, effect = data['cause'], data['effect']
            # Pick an order randomly
            if random.random() < 0.5:
                input_ = f'Sentence 1: {cause} Sentence 2: {effect}'
            else:
                input_ = f'Sentence 1: {effect} Sentence 2: {cause}'
            output_ = [cause]
        elif task == 'common_concept':
            items = data['items']
            # Make comma separated list of items
            input_ = ', '.join(items[:-1])
            output_ = data['all_common_concepts']
        elif task == 'rhymes':
            input_, output_ = data['input'], data['other_rhymes']
        elif 'translation' in task and task != 'salient_translation_error_detection':
            input_, output_ = data['input'], data['possible_translations']
        elif task == 'gender_inclusive_sentences_german':
            input_, output_ = data['input'], data['output']
        else:
            input_, output_ = data['input'], [data['output']]
        inputs.append(input_)
        outputs.append(output_)
    return inputs, outputs, msg
