from brainglobe_atlasapi.core import Atlas

def _get_connected_structures(self: Atlas, structure, hierarchy_lev = None, as_acronym = False ):
    """
    Return the set of connected structure IDs or acronyms for a given structure.
    
    Parameters
    ----------
    structure : str or int
        The structure key or ID in self.structures.
    hierarchy_lev : int or None, optional
        The hierarchy level to extract. If None, returns all levels.
    as_acronym : bool, optional
        If True, return acronyms instead of IDs.
    
    Returns
    -------
    set or list
        Set of structure IDs or list of acronyms at the specified hierarchy level.
    """
    # Input check for hierarchy_lev
    if not (hierarchy_lev is None or int(hierarchy_lev) == hierarchy_lev):
        raise ValueError(f"hierarchy_lev must be an int or None, got {type(hierarchy_lev)}")
    input_id = self.structures[structure]['id']
    input_id_leaves = self.structures.tree.leaves(input_id)
    if hierarchy_lev is None:
        conn_ids      = set(np.concatenate([self.structures[leaf.identifier]['structure_id_path'][:] for leaf in input_id_leaves]))
        conn_acronyms = [self.structures[ss]['acronym'] for ss in conn_ids]
    else:
        try:
            conn_ids      = set([self.structures[leaf.identifier]['structure_id_path'][hierarchy_lev] for leaf in input_id_leaves])
            conn_acronyms = [self.structures[ss]['acronym'] for ss in conn_ids]
        except IndexError:
            raise ValueError(f'Chosen structure { self.structures[structure]["acronym"]} does not have any connected nodes at hierarchy level {hierarchy_lev}')
    if as_acronym:
        return conn_acronyms
    return conn_ids

Atlas.get_connected_structures = _get_connected_structures