from mne.time_frequency import tfr_array_morlet


def filterbank_extractor(data, sfreq, filter_banks, reshape_freqs_dim=False):
    n_cycles = filter_banks / 4
    power = tfr_array_morlet(data[None],
                            sfreq=sfreq,
                            freqs=filter_banks,
                            n_cycles=n_cycles,
                            output='avg_power',
                            verbose=False)
    # (n_ch, n_freqs, n_times)
    # remove power line noise, * f to normalize
    if reshape_freqs_dim:
        power = power.reshape((-1, power.shape[-1]))
    return power