def sleep_stage_dict(mode="decode"):
    """
    Get dictionary to encode/decode sleep stage
    
    Parameters
    ----------
    mode : {"decode", "encode"}, default "decode"
        If "decode", return dict num -> str
        If "encode", return dict str -> num

    Returns
    -------
    dict
        Dictionary for decoding/encoding sleep stages

    """
    assert mode in ["decode", "encode"]
    d = {
        "no_sleep": 0,
        "wake": 1,
        "awake": 1,
        "no_stage": 2,
        "rem": 3,
        "restless": 4,
        "light": 5,
        "asleep": 6,
        "deep": 7,
    }
    if mode == "decode":
        d = {val: key for key, val in d.items()}
    return d