import numpy as np

import mne
from mne.datasets import sample

data_path = sample.data_path()

fname = data_path + '/MEG/sample/sample_audvis-ave.fif'

n_picks = 30  # change to bigger number to see difference in speed

# Reading evoked data
condition = 'Left Auditory'
evoked = mne.read_evokeds(fname, condition=condition, baseline=(None, 0),
                          proj=False)

del evoked.info['projs'][0]
evoked.pick_types(meg=False, eeg=True)
ch_names = evoked.info['ch_names']
evoked.pick_channels(ch_names[:n_picks])
ch_names = ch_names[:n_picks]

anode, cathode = ch_names[1:], ch_names[:-1]

n_channels, n_times = evoked.data.shape
n_virtual_channels = len(anode)
virtual_ch_names = [f'{a}-{c}' for (a, c) in zip(anode, cathode)]
multiplier = np.zeros((n_virtual_channels, n_channels))
for idx, (a, c) in enumerate(zip(anode, cathode)):
    multiplier[idx, ch_names.index(a)] = 1
    multiplier[idx, ch_names.index(c)] = -1

data_rereferenced = multiplier.dot(evoked.data)
evoked_rerefenced = mne.set_bipolar_reference(evoked, anode=anode, cathode=cathode)
picks = [evoked_rerefenced.ch_names.index(ch) for ch in virtual_ch_names]

data_rereferenced_master = evoked_rerefenced.data[picks]
assert np.allclose(data_rereferenced, data_rereferenced_master)