from itertools import product
import json
import random

def new_task_generator(min_sum, max_sum, digit_ranges):
    tasks = [[], []]
    num_digits = len(digit_ranges)

    digit_combinations = product(*(range(min_val, max_val + 1) for min_val, max_val in digit_ranges))

    for digits in digit_combinations:
        result = sum(digits)
        task = ' + '.join(map(str, digits))

        if min_sum <= result <= max_sum:
            tasks[0].append(task)
            tasks[1].append(result)

    return tasks

task1 = new_task_generator(1,9,[(1,9),(1,9)])
task2 = new_task_generator(1,20,[(1,9),(1,9),(1,9)])
task3 = new_task_generator(10,40,[(1,9),(1,9),(10,19)])
task4 = new_task_generator(10,50,[(1,9),(10,19),(10,19)])
task5 = new_task_generator(10,60,[(10,19),(10,19),(10,19)])
task6 = new_task_generator(10,120,[(10,100),(10,100),(10,100),(10,100)])

with open(f'task_universe/task1.json','w') as outfile:
    json.dump(task1, outfile)
with open(f'task_universe/task2.json','w') as outfile:
    json.dump(task2, outfile)
with open(f'task_universe/task3.json','w') as outfile:
    json.dump(task3, outfile)
with open(f'task_universe/task4.json','w') as outfile:
    json.dump(task4, outfile)
with open(f'task_universe/task5.json','w') as outfile:
    json.dump(task5, outfile)
with open(f'task_universe/task6.json','w') as outfile:
    json.dump(task6, outfile)

# Load tasks and results from JSON files.
def load_tasks_from_json(file_path):
    with open(file_path, 'r') as file:
        return json.load(file)

task1 = load_tasks_from_json('task_universe/task1.json')
task2 = load_tasks_from_json('task_universe/task2.json')
task3 = load_tasks_from_json('task_universe/task3.json')
task4 = load_tasks_from_json('task_universe/task4.json')
task5 = load_tasks_from_json('task_universe/task5.json')
task6 = load_tasks_from_json('task_universe/task6.json')

ability_rounds = 8
gen_rounds = 10
total_rounds = ability_rounds+gen_rounds
x = (ability_rounds+gen_rounds)  # Number of samples
random.seed(23)
sample_1 = random.sample(range(len(task1[0])), k=x)
sample_2 = random.sample(range(len(task2[0])), k=x)
sample_3 = random.sample(range(len(task3[0])), k=x)
sample_4 = random.sample(range(len(task4[0])), k=x)
sample_5 = random.sample(range(len(task5[0])), k=2*x)
sample_6 = random.sample(range(len(task6[0])), k=4*x)

task_list = [task1,task2,task3,task4,task5,task6]
sample_list = [sample_1,sample_2,sample_3,sample_4,sample_5,sample_6]

def math_task_storer(task_list,sample_list,tot_rounds=total_rounds):
    compiled_list = []
    for _round in range(tot_rounds):
        label_list = []
        results_list = []
        for i in range(6):
            tasks = task_list[i]
            ids = sample_list[i]
            if i < 4:
                task1 = _round
                label_list.append(tasks[0][ids[task1]])
                results_list.append(tasks[1][ids[task1]])
            elif i == 4:
                task1 = _round*2
                label_list.append(tasks[0][ids[task1]])
                results_list.append(tasks[1][ids[task1]])
                label_list.append(tasks[0][ids[task1+1]])
                results_list.append(tasks[1][ids[task1+1]])
            elif i == 5:
                task1 = _round*2
                label_list.append(tasks[0][ids[task1]])
                results_list.append(tasks[1][ids[task1]])
                label_list.append(tasks[0][ids[task1+1]])
                results_list.append(tasks[1][ids[task1+1]])
                label_list.append(tasks[0][ids[task1+2]])
                results_list.append(tasks[1][ids[task1+2]])
                label_list.append(tasks[0][ids[task1+3]])
                results_list.append(tasks[1][ids[task1+3]])
        round_list = [label_list, results_list]
        compiled_list.append(round_list)
    return compiled_list

################ COMMENTED OUT #######################
#new_compiled_list = math_task_storer(task_list,sample_list,tot_rounds=total_rounds)
#first_part = new_compiled_list[:ability_rounds]
#second_part = new_compiled_list[ability_rounds:ability_rounds+5]
#third_part = new_compiled_list[ability_rounds+5:]
#with open(f'AbilityMath/static/Addition_Increasing_Difficulty/task_list_{len(first_part)}_rounds.json','w') as outfile:
#    json.dump(first_part, outfile)
#with open(f'AttentionAllocation/static/Addition_Increasing_Difficulty/task_list_{len(second_part)}_rounds.json','w') as outfile:
#    json.dump(second_part, outfile)
#with open(f'Treatment/static/Addition_Increasing_Difficulty/task_list_{len(third_part)}_rounds.json','w') as outfile:
#    json.dump(third_part, outfile)
######################################################