from typing import Dict, List, Tuple, Optional
import random
import copy
import os
import glob
from itertools import permutations, combinations
from collections import Counter, defaultdict

# tools to generate and pseudo-randomize sequences of stimuli

def get_reward_info(seq_structure:Dict[str, List[int]], n:int, seed:int):
    ''' For the whole experiment (2 runs), splits sequence names into rewarded and non rewarded.
    The same underlying sequences are rewarded in both runs (A1 in run1 and A2 in run2, etc).'''

    random.seed(seed) # make sure the randomization is the same within the participant's runs
    unique_seq_struc = sorted(set([name[0] for name in seq_structure.keys()]))
    rewarded_seq = random.sample(unique_seq_struc, n)
    non_rewarded_seq = [seq for seq in unique_seq_struc if seq not in rewarded_seq]

    reward_info = dict()
    n_runs = 2
    for i in range(n_runs):
        run_id = f'run{i+1}'
        reward_info[run_id] = dict()
        reward_info[run_id]['reward'] = [seq+str(i+1) for seq in rewarded_seq]
        reward_info[run_id]['no_reward'] = [seq+str(i+1) for seq in non_rewarded_seq]
   
    return reward_info

def distribute_snd(seq_names:List[str], snd_dir:str, seed:int)-> Dict[str, str]:
    ''' Attribute a sound path to each sequence. This is done randomly for each participant.'''
    random.seed(seed)
    sounds = sorted(glob.glob(os.path.join(snd_dir, '*.wav')))
    random.shuffle(sounds)
    random.shuffle(seq_names)
    snd_mapping = {name: snd for name, snd in zip(seq_names, sounds)}
    return snd_mapping

def set_seed(subject_id:str)-> int:
    ''' Write a seed in a file and return it. The seed is based on the subject ID'''
    seed = int(subject_id)
    random.seed(seed)
    return seed

def get_cat_from_stim(stim: str)-> str:
    '''Extract the category from a stimulus path'''
    return os.path.basename(os.path.dirname(stim))

def get_stim_name(stim: str)-> str:
    '''Extract the stimulus name from a stimulus path'''
    return os.path.splitext(os.path.basename(stim))[0].split('_')[0]

def get_stim_mod(stim: str)-> str:
    '''Extract the modality from a stimulus path'''
    return os.path.splitext(os.path.basename(stim))[0].split('_')[-1]  # assuming the last part is the modality (e.g., 'img' or 'txt')

def draw_two(ignore_idx: list=None):
    ''' Returns two items and their positions in the sequence. This function has been reworked to sample only 
    one item out of the 5 left. The 1st index will always be 0. '''
    idx2 = random.sample(range(1, 6), 1)
    if ignore_idx:
        # added this condition to avoid infinite loop. Because there are 5 items and 6 questions -> 1 item will be asked twice. 
        if len(ignore_idx) >= 5:
            return (0, idx2[0])
        while any([i in ignore_idx for i in idx2]):
            idx2 = random.sample(range(1, 6), 1)
    return (0, idx2[0])

def get_stims(input_dir:str, sequence:List[str], modality:str, lang:str)-> List[str]:
    '''Return the paths to the stimuli in the sequence'''
    stim_paths = []
    for item in sequence:
        stim = glob.glob(os.path.join(input_dir, 'stims', lang, '*', f'{item}_{modality}.png'))[0]
        stim_paths.append(stim)
    return stim_paths

def count_dupes(arr:List)-> int:
    '''Count the number of duplicates in a list'''
    seen = set()
    duplicate_counter = 0
    for item in arr:
        if item in seen:
            duplicate_counter += 1
        else:
            seen.add(item)
    return duplicate_counter

def sample_n_throw(arr:List, n:int=3) -> Tuple[List, List]:
    '''Sample n items from a list and remove them'''
    sample = random.sample(sorted(set(arr)), n)
    for s in sample:
        arr.remove(s)
    return sample, arr

def sample_until_no_dupes(arr:List, n:int=3) -> Tuple[List, List]:
    '''Sample n items from a list and remove them'''
    temp = arr.copy()
    sample, new_temp = sample_n_throw(temp, n)
    if count_dupes(new_temp) > 0:
        duplicate = True
        while duplicate:
            temp = arr.copy()
            sample, new_temp = sample_n_throw(temp, n)
            if count_dupes(new_temp) == 0:
                duplicate = False
    return sample, new_temp

def distribute_target_pos(seq_names, seed):
    ''' Distributes target positions used for questions, for each sequence and modality.
    The important point is that for a given sequence (e.g. A2), the target positions must be different
    for the two modalities. 
    
    Also sample an additional position, because there are 6 questions per seq per mod, and only
    5 positions available. '''
    random.seed(seed)
    pos_list = [1, 2, 3, 4, 5] # position 0 is the cue
    target_pos = {}
    for seq in seq_names:
        target_pos[seq] = {}
        overlapping = True
        img_list = pos_list.copy()
        txt_list = pos_list.copy()
        img_list.append(random.choice(pos_list))
        txt_list.append(random.choice(pos_list))
        while overlapping:
            random.shuffle(img_list)
            random.shuffle(txt_list)
            overlap = [a for a, b in zip(img_list, txt_list) if a == b]
            if len(overlap) == 0:
                overlapping = False
                target_pos[seq]['img'] = img_list
                target_pos[seq]['txt'] = txt_list

    return target_pos

def distribute_sequences_run() -> Dict[str, List[str]]:
    ''' 
    Distribute the sequences in the runs with no repetitions inside each run.
    The sequences are organized in 6 groups of 2 sequences each, based on the first element of the sequence.

    Edit October 2025: this function has been highly simplified because we now have 6 sequences duplicated (A1, A2, B1, B2, ...).
    It is still used but there is no randomization anymore. 
    seq_structures:Dict[str, List[int]] was the input before October 2025.

    Returns a dict that looks like that : {'run1': ['A1', 'B2'], 'run2': ['A2', 'B2', ...]}
    '''
    # groups = defaultdict(list)
    # for key, seq in seq_structures.items():
    #     start_int = seq[0]
    #     groups[start_int].append(key)

    # if len(groups) != 6 or any(len(group) != 2 for group in groups.values()):
    #     raise ValueError("The sequences must be organized in 6 groups of len 2.")
    
    # run1, run2 = list(), list()
    # for group in groups.values():
    #     random.shuffle(group)  # shuffle each group of 2
    #     # alternate assigning to run1 and run2
    #     for i, key in enumerate(group):
    #         if i % 2 == 0:
    #             run1.append(key)
    #         else:
    #             run2.append(key)
    # if len(run1) != len(run2):
    #     raise ValueError("The runs must have the same number of sequences.") 

    return {'run1': ['A1', 'B1', 'C1', 'D1', 'E1', 'F1'], 'run2': ['A2', 'B2', 'C2', 'D2', 'E2', 'F2']}

def distribute_sequences_block(seq_names: list, n_blocks: int) -> Dict[str, List[str]]:
    ''' Distribute the sequences in the blocks with no repetitions inside each block.
    It also controls that the 6 sequences are all presented in the first two blocks.
    Returns dict that looks like this: {'block1': ['L', 'E', 'G'], 'block2':[...], ...}
    '''

    def check_2_blocks(blocks: dict) -> bool:
        '''Check if the first two blocks contain all the sequences'''
        block1 = blocks['block1']
        block2 = blocks['block2']
        return True if count_dupes(block1+block2) == 0 else False

    are_all_seq_presented = False

    while not are_all_seq_presented:
        all_seq_left = seq_names*2
        blocks = {f"block{i+1}": [] for i in range(n_blocks)}
        for b in range(n_blocks):
            if b==2:
                seq_subset, all_seq_left = sample_until_no_dupes(all_seq_left) 
            else:
                seq_subset, all_seq_left = sample_n_throw(all_seq_left)
            blocks[f"block{b+1}"] = seq_subset
        are_all_seq_presented = check_2_blocks(blocks)

    return blocks

def generate_run_org(seed) -> Dict[str, Dict[str, List[str]]]:
    ''' Generate the organization of sequences in the runs. The function returns a dict with the blocks of each run.
    It controls that the sequences are well distributed between the two runs.
    If the experiment crashes, the splits will stay the same, but the blocks org will change. This is due to the
    non replicable randomization in distribute_sequences_block(). At least if a seed is set, the two runs will be made
    of different sequences.

    sequences = {'A': ['item1', 'item2', ...], 'B': [...], ...}
    
    Returns dict that looks like this: {'run1': {'block1': ['L', 'E', 'G'], 'block2':[...], ...}, 'run2': {...}}
    '''

    def get_pairs(string:str) -> List[str]:
        '''Get the pairs of sequences in each block'''
        return [''.join(pair) for pair in combinations(string, 2)]
    
    def get_repeated_pairs(pairs:list) -> List[str]:
        ''' Return the repeated pairs in the list'''
        normalized_items = [''.join(sorted(item)) for item in pairs] # reorder the letters in each string
        counts = Counter(normalized_items)
        repeated_items = [item for item, count in counts.items() if count > 1]
        return repeated_items

    def gen_one_run(sequences: Dict[str, List]) ->  Dict[str, List[str]]:
        '''Generate one run with only 2 repeated pairs.
        Returned dict looks like this {'block1': ['L', 'E', 'G'], 'block2':[...], ...}'''
        two_rep_pairs = False
        while not two_rep_pairs: # Keep generating blocks until we get the best case scenario
            blocks = distribute_sequences_block(sequences, 4)
            pairs = []
            for block in blocks.values():
                pairs+=get_pairs(block)
            rep_pairs = get_repeated_pairs(pairs)
            if len(rep_pairs)==2: # 2 repeated pairs (the best case scenario)
                two_rep_pairs = True
        return blocks
    
    random.seed(seed) # has to set again because the function is called multiple times and we want the same result within participant
    seq_separated = distribute_sequences_run()
    run_org = {
        'run1': gen_one_run(seq_separated['run1']),
        'run2': gen_one_run(seq_separated['run2'])
    }
    return run_org

def distribute_sequences_trial(sequence_names: List[str], n_trials: int) -> Dict[str, List[str]]:
    ''' Distribute the sequences in the trials with no repetitions of the positions of elements across trials.
    Returns a dict that looks like this: {'trial1': ['A', 'B', 'C', 'A', 'B', 'C'], 'trial2': [...], ...}
    '''
    orders = generate_orders_trial(sequence_names, n_trials)
    trials = {f"trial{i+1}": None for i in range(n_trials)}
    for i, _ in enumerate(trials):
        trials[f"trial{i+1}"] =orders[i]*2
    return trials

def distribute_mod_seq(n_block:int, seed:int)-> Dict[str, List[str]]:
    ''' Distribute the modality of the first sequence of each trial between the blocks
    Returns a dict that looks like this: {'block1': ['img', 'txt', 'img'], 'block2': [...], ...}
    '''
    random.seed(seed)
    block_layouts = [(2, 1)] * 2 + [(1, 2)] * 2
    random.shuffle(block_layouts) # shuffle block layouts once
    block_org = {f'block{i + 1}': None for i in range(n_block)}
    for i, (img_count, txt_count) in enumerate(block_layouts):
        org = ['img']*img_count + ['txt']*txt_count
        random.shuffle(org) # shuffle the modalities
        block_org[f'block{i + 1}'] = org
    return block_org


def distribute_mod_quest(first_seq_mod):
    ''' Last version of distribute_mod_quest. Goal is to generate the order of modalities for the questions in each trial.
    As the number of questions is now even (6 instead of 3), the constraints are easier to meet.
    Input is first_seq_mod, which encodes the starting modality for each trial (looks like this {'block1': ['txt', 'txt', 'img'], 'block2': ['txt', ...)
    This function uses this info to create the modality order during questions. 
    No need for seed, this is fully deterministic

    Returns a dict that looks like this: {'block1': {'trial1': ['img', 'txt', 'img', 'txt', 'img', 'txt'], 'trial2': [...], ...} ...}
    '''
    modalities = ['img', 'txt']
    quest_mod_org = dict()
    for block_id, mod_list in first_seq_mod.items():
        quest_mod_org[block_id] = dict()
        for i, starting_mod in enumerate(mod_list):
            following_mod = [m for m in modalities if m not in [starting_mod]][0]
            quest_mod_org[block_id][f'trial{i+1}'] = [starting_mod, following_mod] * 3

    return quest_mod_org

def distribute_mod_quest_old(full_block_org: Dict[str, Dict[str, Tuple[str]]], seed: int) -> Dict[str, Dict[str, Dict[str, int]]]:
    '''
    This function takes a dict as argument : {block1:{trial1:('A', 'B', 'C', 'A', 'B', 'C'), trial2:{...}}, block2:{...}}
    and based on multiple constraints, distributes the modalities (images or text) to sequences for the questions of each trial. 
    For each trial, there are 3 items, and one question per item. 9 questions in total per block.
    Constraints:
    1. Each item (i.e., sequence name) must be presented in each modality equally across all blocks (3img, 3txt).
    2. Each block must have a maximum of 5 items in the same modality.
    3. Each trial can have a maximum of 2 items in the same modality.

    Clunky because the output needs to be reprocuced, so I avoided using random iterations until the constraints are met.
    Instead, the random processes are genrated by random.shuffle() and random.choice(), following a step by step logic.
    '''

    ########## MAIN FUNCTIONS (2) ##########

    def assign_mod_prop_block(full_block_org: Dict[str, Dict[str, Tuple[str]]]) -> Dict[str, Dict[str, List[str]]]:
        ''' Distribute modalities at the run level, i.e., for each block. It creates a dict with the number of modality per item in each block.
        It is based on constraints 1 and 2.'''
        # step 1 : create a dict to hold the distribution
        full_block_dist = gen_storage(full_block_org)
        # step 2 : define the dist for the first block following contraint 2.
        assign_initial_patterns(full_block_dist)
        # step 3 : based on what has been assigned, assign opposing pattern for the same items in other blocks (contraint 1).
        sync_pattern_with_block1(full_block_dist)
        # step 4 + 5 : find the loner item (block in which there is only 1 item with no attribution). 
        # Determine what possible modality distribution can be attributed based on other items in the block (constraint 2).
        # Also take care of the loner corresponding item in the other block (constraint 1).
        set_loner_modality(full_block_dist)     
        # step 6 : assign patterns to the two last items
        distribute_last_item_mods(full_block_dist)
        return full_block_dist
    
    def assign_mod_trial(full_block_org: Dict[str, Dict[str, Tuple[str]]], full_block_dist: Dict[str, Dict[str, Dict[str, int]]]) -> Dict[str, Dict[str, List[str]]]:
        ''' Assign the modalities to the trials based on the distribution of modalities in each block. It takes the number of modality 
        assigned to each item in each block and creates a linear list per trial indicating the order in which modalities should be picked for
        questions.
        Here, the constraints are 1 and 3. The function uses backtracking to ensure that the assignment is valid.
        The output is a dict with the same structure as the input, but with the modalities assigned to each item in each trial.
        Example output:
        {'block1': {'trial1': ['img', 'txt', 'img'], 'trial2': [...], ...} ...}
        '''

        quest_mod_org = {}
        for block, trials in full_block_org.items():
            trials_order = {trial_id:items[:3] for trial_id, items in trials.items()}
            block_quotas = full_block_dist[block]
            solution = backtrack(
                trials=trials_order, 
                trial_keys=list(trials_order.keys()), 
                trial_index=0,
                current_assignment={},
                quotas=copy.deepcopy(block_quotas)
            )
            quest_mod_org[block] = solution
        return quest_mod_org
    
    ########## HELPER FUNCTIONS ##########

    # for assign_mod_trial()
    def is_valid_assignment(trial_assignment: List[str]) -> bool:
        '''Ensure no more than 2 items have the same modality'''
        return all(trial_assignment.count(mod) <= 2 for mod in ['img', 'txt'])

    def backtrack(
            trials:Dict[str, Tuple[str]], 
            trial_keys: List[str], 
            trial_index: int, 
            current_assignment: Dict[str, List[str]], 
            quotas: Dict[str, Dict[str, int]]) -> Optional[Dict[str, List[str]]]:
        ''' 
        Parameters:
        --------
        trials: Dict[str, Tuple[str]] - a dict with trial names as keys and items as values.
        trial_keys: List[str] - a list of trial names.
        trial_index: int - the index of the current trial being processed.
        current_assignment: Dict[str, List[str]] - a dict with the current assignment of modalities to items in trials.
        quotas: Dict[str, Dict[str, int]] - a dict with the number of modalities available for each item in the current trial.
        Returns:
        -------
        Optional[Dict[str, List[str]]] - a dict with the assignment of modalities to items in trials if a valid assignment is found, None otherwise.
        '''
        
        # base case: if all trials have been processed, return the current assignment
        if trial_index == len(trial_keys):
            return current_assignment
        # get the current trial name and items
        trial_name = trial_keys[trial_index]
        items_in_trial = trials[trial_name]

        # recursive function to assign modalities to items in the current trial
        def assign_items(item_index: int, trial_assignment: List[str], updated_quotas: Dict[str, Dict[str, int]]) -> Optional[Dict[str, List[str]]]:
            ''' Recursively assign modalities to items in the current trial.
            Parameters:
            --------
            item_index: int - the index of the current item being processed in the trial.
            trial_assignment: List[str] - a list of modalities assigned to items in the current trial.
            updated_quotas: Dict[str, Dict[str, int]] - a dict with the updated number of modalities available for each item in the current trial.
            Returns:
            -------
            Optional[Dict[str, List[str]]] - a dict with the assignment of modalities to items in trials if a valid assignment is found, None otherwise.
            '''
            if item_index == len(items_in_trial):
                if is_valid_assignment(trial_assignment):
                    next_assignment = copy.deepcopy(current_assignment)
                    next_assignment[trial_name] = trial_assignment
                    return backtrack(trials, trial_keys, trial_index + 1, next_assignment, updated_quotas)
                else:
                    return None

            item = items_in_trial[item_index]
            mods = ['img', 'txt']
            random.shuffle(mods)
            for mod in mods:
                if updated_quotas[item][mod] > 0:
                    updated_quotas_copy = copy.deepcopy(updated_quotas)
                    updated_quotas_copy[item][mod] -= 1
                    result = assign_items(item_index + 1, trial_assignment + [mod], updated_quotas_copy)
                    if result:
                        return result
            return None

        return assign_items(0, [], quotas)
    
    ########################

    # for assign_mod_prop_block()
    def distribute_last_item_mods(full_block_dist):
        ''' At that point, the distribution is almost complete, but there are two items that have no attribution. It does not matter
        what pattern is assiged to them as long as they are opposing patterns in the same block: {block2 : {item 1: (2, 1), item 2: (1, 2)...}} and
        opposing accross blocks: {block2 : {item 1: (1, 2)...}, block3 : {item 1: (2, 1)...}}'''
        last_items = set()
        for block, dist in full_block_dist.items():
            for item, mods in dist.items():
                if 0 in mods.values():
                    last_items.add(item)

        last_items = sorted(list(last_items))
        patterns = [(2, 1), (1, 2)]
        random.shuffle(patterns)
        # last items with opposing patterns so they can be used in the same block without violating the constraints
        pairs = [list(zip(last_items, patterns)), list(zip(last_items, patterns[::-1]))]
        # find the blocks that have unfilled items
        unfilled_blocks = [
            block for block, dist in full_block_dist.items()
            if any(v['img'] == 0 and v['txt'] == 0 for v in dist.values())
        ]

        # update these blocks with corresponding predefined pairs
        for block, pair in zip(unfilled_blocks, pairs):
            for item, (n_img, n_txt) in pair:
                if item in full_block_dist[block]:
                    full_block_dist[block][item]['img'] = n_img
                    full_block_dist[block][item]['txt'] = n_txt

    def set_loner_modality(full_block_dist):
        ''' Because of the run structure, after the 2 first steps, there is always one block that has only one item with no attribution.
        This function finds that item and assigns a pattern to it, then assigns the opposing pattern to the corresponding item in the other block.'''

        for block, dist in full_block_dist.items():
            empty_items = [item for item, mods in dist.items() if 0 in mods.values()]
            if len(empty_items) == 1:
                loner = empty_items[0]
                n_img = sum(mods['img'] for mods in dist.values())
                n_txt = sum(mods['txt'] for mods in dist.values())

                if n_img == n_txt:
                    patterns = [(2, 1), (1, 2)]
                elif n_img > n_txt:
                    patterns = [(1, 2)]
                else:
                    patterns = [(2, 1)]

                pattern = random.choice(patterns)
                opposing_pattern = (pattern[1], pattern[0])

                # assign pattern to loner
                full_block_dist[block][loner]['img'] = pattern[0]
                full_block_dist[block][loner]['txt'] = pattern[1]

                # assign opposing pattern to other occurrence of loner
                for other_block, other_dist in full_block_dist.items():
                    if other_block == block:
                        continue
                    if loner in other_dist and 0 in other_dist[loner].values():
                        full_block_dist[other_block][loner]['img'] = opposing_pattern[0]
                        full_block_dist[other_block][loner]['txt'] = opposing_pattern[1]
                break  # no need to continue after processing one loner block

    def sync_pattern_with_block1(full_block_dist):
        ''' Sync the patterns of the items in other blocks with the first block.'''
        block1_items = full_block_dist['block1']
        for block, dist in full_block_dist.items():
            if block == 'block1':
                continue
            # only process items present in both block1 and current block
            common_items = dist.keys() & block1_items.keys()
            for item in common_items:
                mods_ = block1_items[item]
                full_block_dist[block][item]['img'] = mods_['txt']
                full_block_dist[block][item]['txt'] = mods_['img']

    def assign_initial_patterns(full_block_dist):
        ''' Assign the initial patterns for the first block.'''
        patterns = [(2, 1), (1, 2)] # (n_img, n_txt)
        random.shuffle(patterns)
        b1_patterns = [patterns[0], patterns[0], patterns[1]] # randomly chose one pattern to be the dominant
        random.shuffle(b1_patterns) 
        for pattern, item in zip(b1_patterns, full_block_dist['block1'].keys()):
            n_img, n_txt = pattern  
            full_block_dist['block1'][item]['img'] = n_img
            full_block_dist['block1'][item]['txt'] = n_txt

    def gen_storage(full_block_org):
        ''' Generate a storage dict for the distribution of modalities.
        {block1: {item1: {'img': 0, 'txt': 0}, item2: {'img': 0, 'txt': 0}, ...}}, items being sequence names.'''
        full_block_dist = {}
        for block, trials in full_block_org.items():
            full_block_dist[block] = {} 
            for items in trials.values():
                for i in items[:3]: # only 3 because repetition
                    full_block_dist[block][i] = {'img':0, 'txt':0}
        return full_block_dist
    

    random.seed(seed)
    full_block_dist = assign_mod_prop_block(full_block_org)
    quest_mod_org = assign_mod_trial(full_block_org, full_block_dist)
    return quest_mod_org


def generate_orders_trial(sequence_names: list, n_trials: int) -> List[Tuple[str]]:
    ''' Return 3 orders of sequence names: The position of each unique elements is unique.
    Returns this kind of list [('A', 'B', 'C'), ('B', 'C', 'A'), ('C', 'A', 'B')]
    '''
    orders = list(permutations(sequence_names, n_trials))
    pos_tracker = {key:0 for key in range(len(sequence_names))}
    unique_pos_seq = list()
    for order in orders:
        unique_pos = 0
        for i, seq in enumerate(order): # check if the position of the elements is unique
            if pos_tracker[i] == seq:
                continue
            else:
                unique_pos += 1
        if unique_pos == 3: # if all the elements had unique positions
            unique_pos_seq.append(order)
            for i, seq in enumerate(order):
                pos_tracker[i] = seq # update the position tracker

    return unique_pos_seq

def generate_sequences(input_dir:str, seq_structures:Dict, lang:str, seed:int)-> Dict[str, List[str]]:
    ''' Generate 12 unique amodal sequences. They are based on the fixed structures in seq_structures.
    The sequences are returned in a dict {'A1':[item1, 'item2', ...], ...}
    '''
    random.seed(seed) # has to set again because the function is called multiple times and we want the same result within participant
    all_cat = sorted(os.listdir(os.path.join(input_dir, 'stims', lang)))
    all_cat = [cat for cat in all_cat if not cat.startswith('.')] # remove .DS_store
    all_stims = {}
    for cat in all_cat:
        cat_stims = glob.glob(os.path.join(input_dir, 'stims', lang, cat, '*img.png'))
        cat_stims = [os.path.basename(stim).split('_')[0] for stim in cat_stims]
        random.shuffle(cat_stims)
        all_stims[cat] = cat_stims

    sequences = {}
    for i, (seq_name, seq_order) in enumerate(seq_structures.items()):
        seq = list()
        ordered_cat = [all_cat[j] for j in seq_order] # reorder the categries (strings format)
        for cat in ordered_cat:
            item = all_stims[cat][i] # extract the ith item in each category
            seq.append(item)
        sequences[seq_name] = seq
    
    return sequences

def generate_modalities(start_with_img=True):
    ''' Generate the order of modalities for a block'''
    if start_with_img:
        return ['img', 'txt']*3
    else:
        return ['txt', 'img']*3

def check_nstims(categories, input_dir, lang):
    'Check if there is the same number of stim per class, raise error if not'
    counter = 0
    for i, cat in enumerate(categories.keys()):
        nstim = len(glob.glob(f"{input_dir}/stims/{lang}/{cat}/*.png"))
        if i == 0:
            counter = nstim
        if nstim != counter:
            raise ValueError(f"Number of stim for category {cat} is different from the others")
        
def check_img_txt(input_dir, lang):
    ''' Verifiy that for each text stim, there is a associated image stim.'''
    all_img = sorted(glob.glob(f"{input_dir}/stims/{lang}/*/*img.png"))
    all_txt = sorted(glob.glob(f"{input_dir}/stims/{lang}/*/*txt.png"))
    for img in all_img:
        txt = img.replace('img', 'txt')
        if txt not in all_txt:
            raise ValueError(f"Missing image for text stim {img}")

#############################################
 #             Question functions           #
#############################################

def get_slot_pos(y_pos:float)-> List[Tuple[float]]:
    ''' Return the positions of the slots in the trial question'''
    return [(-0.45, y_pos), (-0.15, y_pos), (0.15, y_pos), (0.45, y_pos), (0.75, y_pos)]

def get_response_distance(correct_idx:int, response_idx:int, rt:float)-> int:
    ''' Return the distance between the correct and the response index. If rt is 'NA', return 'NA'.'''
    if rt == 'NA':
        return rt
    return abs(correct_idx - response_idx)

def get_feedback_args(distance:int, lang:str)-> Tuple[str, str, bool, int]:
    '''Return the feedback text and color based on the distance between the correct and the response index.
    The function also returns a boolean indicating if the response was correct'''
    if lang == 'fr':
        txt = 'Trop lent!'
    else:
        txt = 'Too slow!'
    feedback_map = {
        'NA': (txt, "red", False, 0),
        0: ("Correct!\n+ 3pt", "green", True, 3),
    }
    return feedback_map.get(distance, ("Incorrect!\n+ 0pt", "red", False, 0))

def get_trial_feedback(n_points:int, max_points:int, lang:str)-> str:
    '''Return the feedback text based on the number of points obtained'''
    messages = {
        'en': [
            (0, f"Too bad! You didn't win any points out of {max_points}."),
            (0.4, f"Not bad, you got {n_points} questions out of {max_points}. Keep trying!"),
            (0.7, f"Well done! You got {n_points} questions out of {max_points}."),
            (1.0, f"Bravo! Perfect score: {n_points} out of {max_points} !")
        ],
        'fr': [
            (0, f"Dommage! Vous n'avez réussi aucune question sur {max_points}."),
            (0.4, f"Pas mal, vous avez réussi {n_points} question(s) sur {max_points}. Continuez à essayer !"),
            (0.7, f"Bien joué! Vous avez réussi {n_points} questions sur {max_points}."),
            (1.0, f"Bravo! Score excellent : {n_points} sur {max_points} !")
        ]
    }
    lang = lang if lang in messages else 'fr'  # default french
    score_ratio = n_points / max_points if max_points > 0 else 0
    for threshold, message in messages[lang]:
        if score_ratio <= threshold:
            return message

    return messages[lang][-1][1] # if issue, return the last message

def get_fb_sound(reward_sounds, n_points:int):
    ''' Return the reward sound (psychopy.sound) based on the number of points obtained for this question.'''
    if n_points < 3:
        return None
    elif n_points == 3:
        return reward_sounds[2]
    else:
        raise ValueError(f"Invalid number of points: {n_points}")

def run_question(tools:dict, slots:dict, start_item, end_item, rt_clock, global_clock, t_act:float, key_dict:dict, trig_dict:dict)-> Tuple[int, float]:
    '''Run a question where the participant has to place the second item in the correct position.
    NB : it adds 1 to the returned index to takes into account the first item of the sequence (which is not
    selectable)'''

    demo = False # added after the demo was implemented
    if 'demo' in tools:
        demo = True

    left_key = key_dict['left']
    right_key = key_dict['right']
    confirm_key = key_dict['confirm']
    wait_fun = tools['wait_fun']
    event_fun = tools['adapt_getKeys']
    clear_event_fun = tools['adapt_clearEvents']
    if not demo:
        trig_fun = tools['trig_fun']

    def draw_all():
        ''' Draw the slots, the start item and the background'''
        tools['background'].draw()
        for slot in slots:
            slot["rect"].draw()
            slot["highlight"].draw()
        start_item.draw()
        start_item.draw()
        end_item.draw()

    def reset_highlight(slots):
        ''' Reset the highlight of all slots before the next question'''
        for slot in slots:
            slot["highlight"].opacity = 0

    iterations = 0
    current_index = 0
    highlight_onset = None
    running = True
    global_clock.reset()
    rt_clock.reset()

    while running:

        draw_all()
        tools['win'].flip()

        if iterations == 0:
            wait_fun(0.01)
            clear_event_fun()
        else:
            if highlight_onset is None:
                slots[current_index]["highlight"].opacity = 1
                highlight_onset = True

            keys = event_fun()

            if left_key in keys:
                current_index = move_highlight(slots, current_index=current_index, direction="left")
                if not demo:
                    trig_fun(trig_dict['misc']['left'])
                    tools['logger'].info(f"key pressed: {left_key}")
                    tools['logger'].info(f"key trigger (left): {trig_dict['misc']['left']}")
                    tools['logger'].info(f"current index: {current_index+1}")

            if right_key in keys:
                current_index = move_highlight(slots, current_index=current_index, direction="right")
                if not demo:
                    trig_fun(trig_dict['misc']['right'])
                    tools['logger'].info(f"key pressed: {right_key}")
                    tools['logger'].info(f"key trigger (right): {trig_dict['misc']['right']}")
                    tools['logger'].info(f"current index: {current_index+1}")

            if confirm_key in keys:
                if not demo:
                    trig_fun(trig_dict['misc']['confirm'])
                resp_time = rt_clock.getTime()
                global_clock.reset() # reset the global clock to avoid time out
                reset_highlight(slots)
                end_item.pos = slots[current_index]["rect"].pos # move the end item to the selected slot
                draw_all()
                tools['win'].flip()
                if not demo:
                    tools['logger'].info(f"key pressed: {confirm_key}")
                    tools['logger'].info(f"key trigger (confirm): {trig_dict['misc']['confirm']}")
                    tools['logger'].info(f"index selected: {current_index+1}")
                    tools['logger'].info(f"RT: {resp_time}")
                wait_fun(1)
                running = False

            if "escape" in keys:
                running = False

            if global_clock.getTime() > t_act:
                if not demo:
                    trig_fun(trig_dict['misc']['timeout']) 
                    tools['logger'].info('time out')
                    tools['logger'].info('key pressed: None')
                    tools['logger'].info(f"trigger (timeout): {trig_dict['misc']['timeout']}")
                    tools['logger'].info(f"index selected: {current_index+1}")
                    tools['logger'].info("RT: NA")
                    tools['logger'].info(f"clock time: {global_clock.getTime()}")
                running = False
                resp_time = 'NA'
                reset_highlight(slots)
        
        iterations += 1

    return current_index+1, resp_time

def move_highlight(slots:dict, current_index:int, direction:bool=None)-> int:
    ''' Move the highlight to the next or previous slot depending on direction
    
    Parameters
    ----------
    slots : list
        List of slot objects
    current_index : int
        The current index of the highlighted slot
    direction : str
        The direction of the movement, either "left" or "right"

    Returns
    -------
    int
        The new index of the highlighted slot
    '''
    # Reset all highlights
    for slot in slots:
        slot["highlight"].opacity = 0
    if direction == "left":
        if current_index > 0:
            current_index -= 1
        elif current_index == 0:
            current_index = current_index # removed the possibility to jump from position 1 to 5 
        else:
            current_index = len(slots) - 1
    elif direction == "right":
        if current_index < len(slots) - 1:
            current_index += 1
        else:
            current_index = current_index # removed the possibility to jump from position 5 to 1
    # Set the highlight on the current slot
    slots[current_index]["highlight"].opacity = 1
    return current_index
        
def fade_out(tools, obj, clock, f_dur):
    ''' Fade out the object'''
    clock.reset()
    while clock.getTime() < f_dur:
        elapsed_time = clock.getTime()
        opacity = 1.0 - (elapsed_time / f_dur)  # linear fade-out
        obj.opacity = max(0.0, opacity)  # ensure opacity doesn't go below 0
        tools['background'].draw()
        obj.draw()
        tools['win'].flip()

#############################################
 #        Questionnaire functions           #
#############################################

def move_slider(slider, slider_positions:List[float], y:float, current_pos:int, keys:List[str], key_dict:Dict[str, str]):
    ''' To move the slider left or right based on the keys pressed.'''
    if key_dict['left'] in keys:
        current_pos -= 1 if current_pos > 0 else 0
    elif key_dict['right'] in keys:
        current_pos += 1 if current_pos < len(slider_positions) - 1 else 0
    pos = slider_positions[current_pos]
    slider.pos = (pos, y)
    return current_pos

def rotate(w_imgs:list, current_idx:int):
    ''' Rotate the wheel images to the next one'''
    if current_idx == len(w_imgs)-1:
        current_idx = 0
    else:
        current_idx += 1
    return current_idx

def display_wheel(tools, wclock, rclock, tw, tr, wh_imgs):
    ''' Display spinning wheel for tw seconds, rotating every tr seconds.'''
    current_idx = 0
    rclock.reset()
    wclock.reset()
    while wclock.getTime() < tw:
        t = rclock.getTime()
        if t > tr:
            current_idx = rotate(wh_imgs, current_idx)
            rclock.reset()
        tools['background'].draw()
        wh_imgs[current_idx].draw()
        tools['win'].flip()
    return


"""
********   TEST FUNCTIONS   *********
*                                     *
*  .----.                    .---.    *
* '---,  `.________________.'  _  `.  *
*      )   ________________   <_>  :  *
* .---'  .'                `.     .'  *
*  `----'                    `---'    *
*                                     *
***************************************
"""

# TODO: write a small file that tests all the important functions and run it before the experiment

def test_distribute_mod_quest(full_block_org, seed):

    q_mod = distribute_mod_quest(full_block_org, seed)

    # Constraint 1: Global item balance (3 img, 3 txt)
    items_d = {b_n:{t:items[:3] for t, items in b.items()} for b_n,b in full_block_org.items()}
    all_items = [item for block in items_d.values() for trial in block.values() for item in trial]
    global_item_counts =  {i:{'img':0, 'txt':0} for i in set(all_items)}  # unique items across all trials
    
    for (bname_mod, block_mod), (bname_items, block_items) in zip(q_mod.items(), items_d.items()):
        assert bname_mod == bname_items, f"Block names do not match: {bname_mod} != {bname_items}"
        for mods, items in zip(block_mod.values(), block_items.values()):
            for i, m in zip(items, mods):
                global_item_counts[i][m] += 1
    for item, counts in global_item_counts.items():
        assert counts['img'] == 3, f"Item {item} has {counts['img']} img assignments (should be 3)"
        assert counts['txt'] == 3, f"Item {item} has {counts['txt']} txt assignments (should be 3)"

    # Constraint 2: Block balance (max 5 same modality per block). 
    # Also check that two blocks are img heavy and two blocks are txt heavy (should be naturally imposed by the distribution of items)
    heavy_mods = {'img': 0, 'txt': 0}
    for block_name, block in q_mod.items():
        block_mod_counts = {'img': 0, 'txt': 0}
        for trial in block.values():
            for mod in trial:
                block_mod_counts[mod] += 1
        if block_mod_counts['img'] > block_mod_counts['txt']:
            heavy_mods['img'] += 1
        else:
            heavy_mods['txt'] += 1
        assert block_mod_counts['img'] <= 5, f"Block {block_name} has {block_mod_counts['img']} img assignments (max 5)"
        assert block_mod_counts['txt'] <= 5, f"Block {block_name} has {block_mod_counts['txt']} txt assignments (max 5)"
    assert heavy_mods['img'] == 2, f"Expected 2 img heavy blocks, found {heavy_mods['img']}"
    assert heavy_mods['txt'] == 2, f"Expected 2 txt heavy blocks, found {heavy_mods['txt']}"

    # Constraint 3: Trial balance (max 2 same modality per trial)
    for block_name, block in q_mod.items():
        for trial_name, trial in block.items():
            for mod in ['img', 'txt']:
                assert trial.count(mod) <= 2, f"{block_name}-{trial_name} has more than 2 {mod}"

    print("All superficial constraints for test_distribute_mod_quest() are respected.")


def test_draw_two():
    seq = [0, 1, 2, 3, 4, 5]
    ignore_idx = [0, 1, 3, 4]
    indices = draw_two(seq, ignore_idx)
    assert indices not in [[0, 1], [1, 0], [3, 4], [4, 3]], "Indices should not be in ignore_idx"
    print("\o/ All tests passed for draw_two()")

def test_generate_sequences(input_dir, seq_structures):

    def check_no_dupe(randomize, n=100):
        ''' Check multiple times that there are no duplicates in the 6 generated sequences'''
        for i in range(n):
            seqs = generate_sequences(input_dir, seq_structures, randomize=randomize)
            keys = list(seqs.keys())
            all_items = [item for sublist in seqs.values() for item in sublist]
            assert len(keys) == 6, f"randomize={randomize}; Wrong number of sequences: {keys}"
            assert len(keys) == len(set(keys)), f"randomize={randomize}; Duplicate keys in {keys}, iteration {i}"
            assert len(all_items) == len(set(all_items)), f"randomize={randomize}; Duplicate items in {all_items}, iteration {i}"
    
    for randomize in [True, False]:
        check_no_dupe(randomize)

    def get_2_items_same_pos(randomize, n=10):
        ''' Fetch two items in the same position (same index of same sequence) in different iterations'''
        all_seqs = []
        for i in range(n):
            seqs = generate_sequences(input_dir, seq_structures, randomize=randomize)
            all_seqs.append(seqs)
        items1 = []
        items2 = []
        for i, seq1 in enumerate(all_seqs):
            if i == 0:
                continue
            seq2 = all_seqs[i-1]
            for j, key1 in enumerate(seq1.keys()):
                key2 = list(all_seqs[i-1].keys())[j]
                for k in range(len(seq1[key1])):
                    item1 = seq1[key1][k]
                    item2 = seq2[key2][k]
                    items1.append(item1)
                    items2.append(item2)
        return zip(items1, items2)
    
    n = 10
    n_pairs = (n-1) * 6 * 6
    pairs = get_2_items_same_pos(randomize=False, n=n)
    for item1, item2 in pairs:
        assert item1 == item2, "When param randomize=False, the sequences should be similar between iterations"
    
    pairs = get_2_items_same_pos(randomize=True, n=n)
    n_similar_items = sum([item1 == item2 for item1, item2 in pairs])
    assert n_similar_items != n_pairs, "When param randomize=True, the sequences should be different between iterations"
    print(f'When randomize=True, the number of similar items: {n_similar_items} out of {n_pairs} pairs ({1/(n_pairs/n_similar_items)*100:.2f}%), should be around 16.66%')
    print("\o/ All tests passed for generate_sequences()")
 
 
'''
*** LEGACY FUNCTIONS ****
*      \_______/        *
*  `.,-'\_____/`-.,'    *
*   /`..'\ _ /`.,'\     *
*  /  /`.,' `.,'\  \    *
* /__/__/     \__\__\__ *
* \  \  \     /  /  /   * 
*  \  \,'`._,'`./  /    *
*   \,'`./___\,'`./     *
*  ,'`-./_____\,-'`.    *
*      /       \        *
*************************   
'''


def get_reward_info_old(two_run_org:Dict[str, Dict[str, List[str]]], seed:int, n:int=3)-> Dict[str, List[str]]:
    ''' For the whole experiment (2 runs), splits sequence names into rewarded and non rewarded. 
    The function needs a dict that looks like this: 
    {'run1': {'block1': ['L', 'E', 'G'], 'block2':[...], ...}, 'run2': {...}}. 
    All the sequences names are in the blocks 1 and 2 of each run (in the two last blocks, they are repeated).
    The function uses set() to remove duplicates in the list of sequences.
    '''
    random.seed(seed) # make sure the randomization is the same within the participant's runs
    reward_info = {}
    for run, blocks in two_run_org.items(): # for key, small dict in bigger dict
        reward_info[run] = {}
        all_seq = [] # use a list to be able to use += below
        for block in blocks.values():
            all_seq += block
        all_seq = sorted(list(set(all_seq))) # remove duplicates
        reward_info[run]['reward'] = random.sample(all_seq, n)
        reward_info[run]['no_reward'] = [seq for seq in all_seq if seq not in reward_info[run]['reward']]

    return reward_info

def jitter_isi(value:float=0.2)-> float:
    '''Return a random jittered inter-stimulus interval'''
    # not used anymore
    return random.uniform(-value, value)

# was used when I needed to distribute 2*6 sequences. Current version works for 12 unique sequences. 
def old_generate_run_org(input_dir, seq_structures) -> dict:
    ''' Generate the organization of sequences in the runs. The function returns a dict with the blocks of each run.
    It controls that the sequences are well distributed between the two runs, with minimal overlap'''

    def get_pairs(string:str) -> List[str]:
        '''Get the pairs of sequences in each block'''
        return [''.join(pair) for pair in combinations(string, 2)]
    
    def get_repeated_pairs(pairs:list) -> List[str]:
        ''' Return the repeated pairs in the list'''
        normalized_items = [''.join(sorted(item)) for item in pairs] # reorder the letters in each string
        counts = Counter(normalized_items)
        repeated_items = [item for item, count in counts.items() if count > 1]
        return repeated_items
    
    def get_unrepeated_seq(pairs:list, repeated_pairs) -> List[str]:
        ''' Return sequences that are not in repeated pairs'''
        pair_as_str = "".join(pairs)
        rep_pairs_as_str = "".join(repeated_pairs)
        unrepeated_seq = []
        for pair in pair_as_str:
            if (pair not in rep_pairs_as_str) and (pair not in unrepeated_seq):
                unrepeated_seq.append(pair)
        return unrepeated_seq

    def gen_one_run(sequences: dict) -> Tuple[dict, List[str]]:
        '''Generate one run with only 2 repeated pairs. Return the blocks, repeated pairs and unrepeated sequences'''
        two_rep_pairs = False
        while not two_rep_pairs: # Keep generating blocks until we get the best case scenario
            blocks = distribute_sequences_block(sequences, 4)
            pairs = []
            for block in blocks.values():
                pairs+=get_pairs(block)
            rep_pairs = get_repeated_pairs(pairs)
            unrep_seq = get_unrepeated_seq(pairs, rep_pairs)
            if len(rep_pairs)==2: # 2 repeated pairs (the best case scenario)
                two_rep_pairs = True
        return blocks, rep_pairs, unrep_seq
    
    sequences = generate_sequences(input_dir, seq_structures)
    # Generate two runs with unique repeated pairs between them (e.g. run1 AB, CF and run2 ED, FA)
    running = True
    while running:
        blocks1, rep_pairs1, unrep_seq1 = gen_one_run(sequences)
        blocks2, rep_pairs2, unrep_seq2 = gen_one_run(sequences)
        # condition 1 : the standalone sequences are different in the two runs
        four_standalone_seq = True if len(set(unrep_seq1+unrep_seq2))==4 else False
        # condition 2 : the repeated pairs are unique between the two runs
        counter = 0
        for pair in rep_pairs1:
            if pair not in rep_pairs2:
                counter+=1
        unique_repetitions = True if counter==2 else False
        if unique_repetitions and four_standalone_seq:
            running = False

    run_org = {'run1': blocks1, 'run2': blocks2}
    return run_org

# not used in the current version of the experiment
def check_positions(blocks):
    ''' Check if the sequences are not in the same position in two blocks'''
    all_positions = {f"pos{i}": [] for i in range(len(blocks['block1']))}
    for block in blocks.values():
        for i, seq in enumerate(block):
            all_positions[f"pos{i}"].append(seq)
    for pos in all_positions.values():
        if count_dupes(pos) > 0:
            return False
    return True
    

def check_distribution(blocks):
    ''' Check if the distribution of 3 sequences blocks is not repeated.
    e.g. if ABC is in block1, it should not be in block2 whatever the order'''
     # NOT USED, its function is also passively covered by check_positions()
    all_sequences = []
    for block in blocks.values():
        com = list(permutations(block))
        com = [''.join(c) for c in com]
        all_sequences += com
    if count_dupes(all_sequences) > 0:
        return False
    return True

def check_transitions(order1: str, order2: str):
    ''' Check if the transitions between two sequences are unique'''
    if order1 == order2:
        return False
    transitions = []
    for order in (order1, order2):
        order = ''.join(order)
        for i in range(len(order)-1):
            transitions.append(order[i:i+2])
    if count_dupes(transitions) > 0:
        return False
    return True

def generate_run_order(sequences, n_blocks=4):
    ''' Generate the run order, with the constraints that each sequence is presented in two different blocks,
     the sequences should not be in the same position in the two blocks, and all the transitions should be unique'''
    continue_sampling = True
    while continue_sampling:
        blocks = distribute_sequences_block(sequences, n_blocks=n_blocks)
        block_list = list(blocks.values())
        pairs = list(combinations(block_list, 2))
        n_pairs = len(pairs)
        passed_test = 0
        for pair in pairs:
            if check_transitions(pair[0], pair[1]):
                passed_test += 1
        if (passed_test == n_pairs) and (check_positions(blocks)):
                continue_sampling = False
    return blocks



##############################################
#          Localizer functions               #
##############################################

def get_target_stims(input_dir:str, modality:str, lang:str)-> List[str]:
    '''Return the paths to the target in the sequence'''
    stim_paths = glob.glob(os.path.join(input_dir, 'loc_targets', lang, '*', f'trg_*_{modality}.png'))
    return stim_paths

def prep_stim_tgt_order(input_dir, stim_paths, modality, nblocks, n_tgt_low, n_tgt_high, lang='fr'):
    ''' Prepare stimuli for loc type B.
    Targets are specific items created for the localizer. 
    Output dict : {'block1': [{'path': 'path/to/stim1', 'label': 'target'/'valid'}, ...]}'''

    target_paths = get_target_stims(input_dir, modality=modality, lang=lang)
    random.shuffle(target_paths)
    if len(target_paths) < nblocks:
        raise ValueError(f"Not enough target stimuli for modality {modality}. Found {len(target_paths)}, but need at least {nblocks}.")
    
    reg_paths = [s for s, mod in stim_paths.items() if mod == modality]
    stim_orders = {}
    for i, tgt in enumerate(target_paths, start=1):
        reg_p = reg_paths.copy()  
        random.shuffle(reg_p)  # Shuffle valid paths for each block
        n_tgt = random.randint(n_tgt_low, n_tgt_high) # Number of targets for this block
        n_stim_total = n_tgt + len(reg_p)
        tgt_indices = gen_tgt_idx(n_tgt, n_stim_total, min_spacing=4, first_allowed=3) # generate indices for targets
        combined_paths = []
        for idx in range(n_stim_total):
            if idx in tgt_indices:
                combined_paths.append(tgt)
            else:
                combined_paths.append(reg_p.pop(0))
        stim_orders[f'block{i}'] = [{'path': p, 'label': ('target' if p == tgt else 'valid')} for p in combined_paths] 

    # return target paths to have information for the instructions -> "this is the target in this block"
    return stim_orders, target_paths


def prep_stim_supra_cat(stim_paths, modality, n_tgt_low, n_tgt_high):
    ''' Prepare stimuli for loc type B, where stimuli are separated by broader semantic supra-categories.
    Targets are items that are not in the same supra-category as block.
    Output dict : {'block1': [{'path': 'path/to/stim1', 'label': 'valid'}, ...]}'''

    anim_cat = ['animals', 'bodyparts', 'characters']
    anim_paths = [s for s, mod in stim_paths.items() if get_cat_from_stim(s) in anim_cat and mod == modality]
    inanim_paths = [s for s, mod in stim_paths.items() if get_cat_from_stim(s) not in anim_cat and mod == modality]
    anim_pool = anim_paths.copy() # items to be chosen (and popped) as targets
    inanim_pool = inanim_paths.copy()
    supra_cat_order = 3*['animate', 'inanimate']  # 3 blocks of animate, 3 blocks of inanimate. Not randomized for now. 
    stim_orders = {}
    for i, supra_cat in enumerate(supra_cat_order, start=1):
        n_tgt = random.randint(n_tgt_low, n_tgt_high)
        if supra_cat == 'animate':
            reg_paths = anim_paths.copy() * 2  # all items are presented twice
            tgt_paths = random.sample(inanim_pool, n_tgt)  # choose targets from inanimate pool
            inanim_pool = [p for p in inanim_pool if p not in tgt_paths] # remove targets from the pool so they are not used again
        else:
            reg_paths = inanim_paths.copy() * 2
            tgt_paths = random.sample(anim_pool, n_tgt)
            anim_pool = [p for p in anim_pool if p not in tgt_paths]

        random.shuffle(reg_paths)
        tgt_indices = gen_tgt_idx(n_tgt, len(reg_paths)+len(tgt_paths), min_spacing=4, first_allowed=3)
        combined_paths = []
        n_stim_total = len(tgt_indices) + len(reg_paths)
        for idx in range(n_stim_total):
            if idx in tgt_indices:
                combined_paths.append({'path': tgt_paths.pop(0), 'label': 'target'})
            else:
                combined_paths.append({'path': reg_paths.pop(0), 'label': 'valid'})
        stim_orders[f'block{i}'] = combined_paths

    # need to return stim_orders, and the order of categories for the instructions
    return stim_orders, supra_cat_order

def dist_pairs_far_away(pairs):
    '''Take a list of pairs, and carefully flatten the list and shuffle the items inside such that the indices (0, 1) keep alternating
    and the items in the pairs are never next to each other.
    Used brute force because marre de chez marre de faire des trucs compliqués.
    Example : [['A_txt', 'A_img'], ['B_txt', 'B_img'], ['C_txt', 'C_img'], ['D_txt', 'D_img']] becomes ['A_txt', B_img', 'C_txt', 'D_img', 'B_txt',...].'''

    # start by creating two lists and make sure the items of each list are not overlapping +- index
    firsts = [(i, pair[0]) for i, pair in enumerate(pairs)]
    seconds = [(i, pair[1]) for i, pair in enumerate(pairs)]
    for _ in range(1000):
        random.shuffle(firsts) # contains items of modality 1
        random.shuffle(seconds) # modality 2
        # j = index in the split list, i = index in the original list
        good_trans_count = 0
        for j, (i, _) in enumerate(firsts):
            if j == 0:
                before = None
                after = seconds[j+1][0]
            elif j > 0 and j < len(firsts) - 1:
                before = seconds[j-1][0]
                after = seconds[j+1][0]
            elif j == len(firsts) - 1:
                before = seconds[j-1][0]
                after = None
            same = seconds[j][0]
            if i != before and i != after and i != same:
                good_trans_count += 1

        if good_trans_count == len(firsts):
            break
    # take items for each list and fill a new list with alternating items
    final = []
    for i in range(len(firsts)):
        final.append(firsts[i][1])
        final.append(seconds[i][1])

    return final
        
def check_mod_alternation(item_list):
    ''' Check if the pairs alternate between modalities. 
    Returns True if they do, False otherwise. '''
    mods = [get_stim_mod(p) for p in item_list]
    return all(mods[i] != mods[i+1] for i in range(len(mods)-1))

def prep_stim_match(stim_paths, nblocks, n_tgt_low, n_tgt_high):
    ''' Prepare stimuli for loc type C, where items from different modalities aternate. 
    Targets are 2nd items of two matching items.'''
    # step 1: pair stimuli by their root names
    pairs = pair_paths_by_root(stim_paths)
    # new version: create list of dicts, with unpaired items as 'valid' and paired items as 'target'
    # NB : only the 2nd item of the matching par is considered a target, as the first item is always a valid item.
   
    stim_orders = {}
    for i in range(nblocks):
        pairs_cpy = pairs.copy()
        #random.shuffle(pairs_cpy)  # shuffle pairs for each block
        n_tgt = random.randint(n_tgt_low+2, n_tgt_high+2) # +2 because sequences of 72 items instead of 36
        tgt_pairs = random.sample(pairs_cpy, n_tgt)  # select target pairs
        pairs_left = [p for p in pairs_cpy if p not in tgt_pairs] # remove target pairs from the pool of pairs
        # flatten the pairs to get a list of paths
        flat_items = dist_pairs_far_away(pairs_left) # Makes sure corresponding items are far from each other. + shuffle overkill
        if not check_mod_alternation(flat_items):
            raise ValueError("The pairs do not alternate between modalities. Please check the stimuli.")
        # swap half of the target pairs to create different transitions of modality (some are txt/img, some are img/txt)
        tgt_pairs = [list(pair) for pair in tgt_pairs]  # convert tuples to be able to swap items
        n_swap = len(tgt_pairs) // 2
        for j in range(n_swap):
            idx = random.randint(0, n_tgt-1)
            tgt_pairs[idx][0], tgt_pairs[idx][1] = tgt_pairs[idx][1], tgt_pairs[idx][0]

        tgt_indices = gen_tgt_idx(n_tgt, len(flat_items), min_spacing=5, first_allowed=3) # subtract 2 because to have some margin if the indices change. 
        # start by creating a list of valid items
        block_order = []
        for path in flat_items:
            block_order.append({'path': path, 'label': 'valid'})
        
        # now insert the target pairs into the block order. 
        add_idx = 0 # need to add a shift to the indices as we insert items into the block_order
        for j in tgt_indices:
            idx = j + add_idx
            tgt = tgt_pairs.pop(0)
            start_mod = get_stim_mod(tgt[0])
            # if the mod of the 1st item has the same modality as the item sitting at the idx, insert both items at the idx.
            if start_mod == get_stim_mod(block_order[idx]['path']):
                block_order.insert(idx, {'path': tgt[0], 'label': 'valid'})
                block_order.insert(idx+1, {'path': tgt[1], 'label': 'target'})
            # else move to the next position. 
            else:
                block_order.insert(idx+1, {'path': tgt[0], 'label': 'valid'})
                block_order.insert(idx+2, {'path': tgt[1], 'label': 'target'})
            add_idx += 2  # add 2 to the indices as we insert two items into the block_order
        
        stim_orders[f'block{i+1}'] = block_order

    return stim_orders

def pair_paths_by_root(stim_paths):
    ''' Pair paths by their root names.'''
    roots = defaultdict(dict)
    for path, mod in stim_paths.items():
        root = path[:-8]  # removes '_mod.png'
        roots[root][mod] = path
    pairs = []
    for paths in roots.values():
        pairs.append((paths['txt'], paths['img']))
    return pairs

def gen_tgt_idx(n_tgt, n_total, min_spacing=4, first_allowed=3):
    '''
    Return a list of indices that are spaced apart by at least `min_spacing` indices.
    Used to place targets in the sequence of valid items.
    '''
    chosen = []
    # iterative elimination:
    # pick a random slot → remove every slot within `min_spacing` distance
    pool = list(range(first_allowed, n_total))
    while len(chosen) < n_tgt:
        idx = random.choice(pool)
        chosen.append(idx)
        pool = [p for p in pool if abs(p - idx) > min_spacing]

    return sorted(chosen)

def wait_and_check_resp(tools:dict, clock, isi:float, trig_val:int, key_dict:dict) -> bool:
    ''' Wait for the ISI time, check for response. If response, send trigger and wait an aditional 1s.
    Return True if response, False otherwise.'''
    while clock.getTime() < isi:
        keys = tools['adapt_getKeys']()
        if key_dict['confirm'] in keys:
            tools['trig_fun'](trig_val)
            tools['logger'].info("Detection key pressed")
            tools['wait_fun'](1)
            tools['adapt_clearEvents']()
            return True
    return False

def feed_resp_dict(resp_dict, resp ,label, logger):
    if resp is True and label == 'target':
        resp_dict['correct'] += 1
        logger.info('response: correct target')
    elif resp is False and label == 'valid':
        resp_dict['correct'] += 1
        logger.info('response: correct valid')
    elif resp is True and label == 'valid':
        resp_dict['false_alarm'] += 1
        logger.info('response: false alarm')
    elif resp is False and label == 'target':
        resp_dict['miss'] += 1
        logger.info('response: miss')

def cpt_points_loc(resp_dict):
    ''' Calculate the points based on the responses. Returns nothing. '''
    incorrect = resp_dict['false_alarm'] + resp_dict['miss']
    total = resp_dict['correct'] + incorrect
    percent_incor = (incorrect / total) * 100 if total > 0 else 0
    return percent_incor
