def get_unit_waveforms(recording, sorting, unit_ids=None, grouping_property=None, start_frame=None, end_frame=None,
                     ms_before=3., ms_after=3., dtype=None, max_num_waveforms=np.inf, filter=False,
                     bandpass=[300, 6000], save_as_features=True, verbose=False, compute_property_from_recording=False):
    '''
    This function returns the spike waveforms from the specified unit_ids from t_start and t_stop
    in the form of a numpy array of spike waveforms.

    Parameters
    ----------
    recording
    sorting
    unit_ids
    grouping_property
    start_frame
    end_frame
    ms_before
    ms_after
    dtype
    max_num_waveforms
    filter
    bandpass
    save_as_features
    verbose
    compute_property_from_recording

    Returns
    -------
    waveforms

    '''
    if isinstance(unit_ids, (int, np.integer)):
        unit_ids = [unit_ids]
    elif unit_ids is None:
        unit_ids = sorting.get_unit_ids()
    elif not isinstance(unit_ids, (list, np.ndarray)):
        raise Exception("unit_ids is not a valid in valid")

    if dtype is None:
        dtype = np.float32

    waveform_list = []
    if grouping_property is not None:
        if grouping_property not in recording.get_channel_property_names():
            raise ValueError("'grouping_property' should be a property of recording extractors")
        if compute_property_from_recording:
            compute_sorting_group = True
        elif grouping_property not in sorting.get_unit_property_names():
            print(grouping_property, ' not in sorting extractor. Computing it from the recording extractor')
            compute_sorting_group = True
        else:
            compute_sorting_group = False
        print("Waveforms by property: ", grouping_property)

        if not compute_sorting_group:
            rec_list, rec_props = se.get_sub_extractors_by_property(recording, grouping_property, return_property_list=True)
            sort_list, sort_props = se.get_sub_extractors_by_property(sorting, grouping_property, return_property_list=True)
            if len(rec_props) != len(sort_props):
                print('Different' + grouping_property + ' numbers: using largest number of ' + grouping_property)
                if len(rec_props) > len(sort_props):
                    for i_r, rec in enumerate(rec_props):
                        if rec not in sort_props:
                            print('Inserting None for property ', rec)
                            sort_list.insert(i_r, None)
                else:
                    for i_s, sort in enumerate(sort_props):
                        if sort not in rec_props:
                            rec_list.insert(i_s, None)
            else:
                assert len(rec_list) == len(sort_list)

            for i_list, (rec, sort) in enumerate(zip(rec_list, sort_list)):
                for i, unit_id in enumerate(unit_ids):
                    # ts_ = time.time()
                    if sort is not None and rec is not None:
                        if unit_id in sort.get_unit_ids():
                            if not filter:
                                rec = rec
                            else:
                                rec = st.preprocessing.bandpass_filter(recording=rec, freq_min=bandpass[0],
                                                                       freq_max=bandpass[1])

                            fs = rec.get_sampling_frequency()
                            n_pad = [int(ms_before * fs / 1000), int(ms_after * fs / 1000)]

                            if verbose:
                                print('Waveform ' + str(i + 1) + '/' + str(len(unit_ids)))
                            waveforms, indices = _get_random_spike_waveforms(recording=rec,
                                                                             sorting=sort,
                                                                             unit=unit_id,
                                                                             max_num=max_num_waveforms,
                                                                             snippet_len=n_pad)
                            waveforms = waveforms.swapaxes(0, 2)
                            waveforms = waveforms.swapaxes(1, 2)
                            waveforms = waveforms.astype(dtype)

                            if save_as_features:
                                if len(indices) < len(sort.get_unit_spike_train(unit_id)):
                                    if 'waveforms' not in sorting.get_unit_spike_feature_names(unit_id):
                                        features = np.array([None] * len(sorting.get_unit_spike_train(unit_id)))
                                    else:
                                        features = np.array(sorting.get_unit_spike_features(unit_id, 'waveforms'))
                                    for i, ind in enumerate(indices):
                                        features[ind] = waveforms[i]
                                else:
                                    features = waveforms
                                sorting.set_unit_spike_features(unit_id, 'waveforms', features)
                            waveform_list.append(waveforms)
        else:
            for i, unit_id in enumerate(unit_ids):
                # ts_ = time.time()
                if unit_id in sorting.get_unit_ids():
                    rec_groups = np.array([recording.get_channel_property(ch, grouping_property)
                                           for ch in recording.get_channel_ids()])
                    if not filter:
                        rec = recording
                    else:
                        rec = st.preprocessing.bandpass_filter(recording=recording, freq_min=bandpass[0],
                                                               freq_max=bandpass[1])

                    fs = rec.get_sampling_frequency()
                    n_pad = [int(ms_before * fs / 1000), int(ms_after * fs / 1000)]

                    if verbose:
                        print('Waveform ' + str(i + 1) + '/' + str(len(unit_ids)))
                    waveforms, indices = _get_random_spike_waveforms(recording=recording,
                                                                     sorting=sorting,
                                                                     unit=unit_id,
                                                                     max_num=max_num_waveforms,
                                                                     snippet_len=n_pad)
                    waveforms = waveforms.swapaxes(0, 2)
                    waveforms = waveforms.swapaxes(1, 2)
                    waveforms = waveforms.astype(dtype)
                    mean_waveforms = np.squeeze(np.mean(waveforms, axis=0))
                    max_amp_elec = np.unravel_index(mean_waveforms.argmin(), mean_waveforms.shape)[0]
                    group = recording.get_channel_property(recording.get_channel_ids()[max_amp_elec], grouping_property)
                    elec_group = np.where(rec_groups == group)
                    waveforms = np.squeeze(waveforms[:, elec_group, :])
                    if save_as_features:
                        if len(indices) < len(sorting.get_unit_spike_train(unit_id)):
                            if 'waveforms' not in sorting.get_unit_spike_feature_names(unit_id):
                                features = np.array([None] * len(sorting.get_unit_spike_train(unit_id)))
                            else:
                                features = np.array(sorting.get_unit_spike_features(unit_id, 'waveforms'))
                            for i, ind in enumerate(indices):
                                features[ind] = waveforms[i]
                        else:
                            features = waveforms
                        sorting.set_unit_spike_features(unit_id, 'waveforms', features)
                    waveform_list.append(waveforms)
            return waveform_list
    else:
        for i, unit_id in enumerate(unit_ids):
            # ts_ = time.time()
            if unit_id not in sorting.get_unit_ids():
                raise Exception("unit_ids is not in valid")

            if filter:
                recording = st.preprocessing.bandpass_filter(recording=recording, freq_min=bandpass[0],
                                                             freq_max=bandpass[1]).get_traces(start_frame=start_frame, end_frame=end_frame)

            fs = recording.get_sampling_frequency()
            n_pad = [int(ms_before * fs / 1000), int(ms_after * fs / 1000)]

            if verbose:
                print('Waveform ' + str(i + 1) + '/' + str(len(unit_ids)))
            waveforms, indices = _get_random_spike_waveforms(recording=recording,
                                                             sorting=sorting,
                                                             unit=unit_id,
                                                             max_num=max_num_waveforms,
                                                             snippet_len=n_pad)
            # print('extract wf: ', time.time() - ts_)
            waveforms = waveforms.swapaxes(0, 2)
            waveforms = waveforms.swapaxes(1, 2)
            waveforms = waveforms.astype(dtype)
            # print('swap wf: ', time.time() - ts_)
            if save_as_features:
                if len(indices) < len(sorting.get_unit_spike_train(unit_id)):
                    if 'waveforms' not in sorting.get_unit_spike_feature_names(unit_id):
                        features = np.array([None] * len(sorting.get_unit_spike_train(unit_id)))
                    else:
                        features = np.array(sorting.get_unit_spike_features(unit_id, 'waveforms'))
                    for i, ind in enumerate(indices):
                        features[ind] = waveforms[i]
                else:
                    features = waveforms
                sorting.set_unit_spike_features(unit_id, 'waveforms', features)
            # print('append feats: ', time.time() - ts_)
            waveform_list.append(waveforms)
            # print('append wf: ', time.time() - ts_)

        if len(waveform_list) == 1:
            return waveform_list[0]
        else:
            return waveform_list