rereference_bipolar

import re
import numpy as np

def _reref_electrodes(signals, signal_list, method='laplacian'):
    """ references given signals """
    # https://www.ncbi.nlm.nih.gov/pmc/articles/PMC6495648/
    # signals: 2d array of time-series data (rows = electrodes as defined and indexed by signal list, columns = time)
    # signal_list = electrode labels, in order of rows in signals variable; must be the same length
    # method = referencing method
           
    signal_dict = _cluster_contacts(signal_list) 

    sig_list_new = []
    signals_new = []
    if method == 'unipolar':
        signals_new = signals
    else:
        for sig in signal_dict:
            if method == 'mean':
                channels = [signal_list.index(e) for e in signal_list if sig in e]
                sig_ref = signals[channels, :] - np.mean(signals[channels,:], axis=0)
                signals_new.append(sig_ref)
            
            else:
                minn = np.min(signal_dict[sig])
                maxn = np.max(signal_dict[sig])
                for num in signal_dict[sig]:
                    if method == 'bipolar':
                        if num >= minn and num < maxn:
                            adj = 1
                            if num + adj <= maxn:
                                while not '{}{}'.format(sig, num+adj) in signal_list and num + adj <= maxn:
                                    adj += 1
                                sig_list_new.append('{}{}'.format(sig, num))
                                sig_ref_i = signal_list.index('{}{}'.format(sig, num))
                                sig_ref_j = signal_list.index('{}{}'.format(sig, num + adj))
                                sig_ref = signals[sig_ref_i, :] - signals[sig_ref_j, :]
                                signals_new.append(sig_ref)

                    elif method == 'laplacian':
                        if num > minn and num < maxn:
                            hi = 1
                            lo = 1
                            above = '{}{}'.format(sig, num + hi)
                            below = '{}{}'.format(sig, num - lo)
                            while above not in signal_list or below not in signal_list and num + hi <= maxn:
                                if above not in signal_list:
                                    hi += 1
                                if below not in signal_list:
                                    lo += 1
                                above = '{}{}'.format(sig, num + hi)
                                below = '{}{}'.format(sig, num - lo)

                            sig_ref_i = signal_list.index(below)
                            sig_ref_j = signal_list.index(above)
                            sig_ref_0 = signal_list.index('{}{}'.format(sig, num))

                            print('{} = {} - {}'.format('{}{}'.format(sig, num), below, above))

                            sig_list_new.append('{}{}'.format(sig, num))
                            sig_ref = signals[sig_ref_0, :] - (signals[sig_ref_i, :] + signals[sig_ref_j, :])/2
                            signals_new.append(sig_ref)
                        elif num == minn:
                            hi = 1
                            above = '{}{}'.format(sig, num + hi)
                            while above not in signal_list and num + hi <= maxn:
                                hi += 1
                                above = '{}{}'.format(sig, num + hi)

                            sig_ref_j = signal_list.index(above)
                            sig_ref_0 = signal_list.index('{}{}'.format(sig, num))

                            print('{} = {} - {}'.format('{}{}'.format(sig, num), '{}{}'.format(sig, num), above))

                            sig_list_new.append('{}{}'.format(sig, num))
                            sig_ref = signals[sig_ref_0, :] - signals[sig_ref_j, :]
                            signals_new.append(sig_ref)                                

    if len(sig_list_new) == 0:
        sig_list_new = signal_list

    del signals, signal_list # maybe saves space? 
    signals_new = np.vstack(signals_new)

    return signals_new, sig_list_new

def _cluster_contacts(signal_list):
    """ return electrode-contact hierarchy """

    signal_dict = {}

    for sig in signal_list:
        sig_grps = re.search('([A-Z|0-9]{1,}[-| ])?([A-Z]{1,})([0-9]{1,})', sig, re.IGNORECASE)
        if sig_grps:
            n_grps = len(sig_grps.groups())
            electrode = ''.join(filter(None,[sig_grps.group(i) for i in range(1, n_grps)]))
            num = sig_grps.group(n_grps)
            if electrode in signal_dict.keys():
                assert int(num) not in signal_dict[electrode]
                signal_dict[electrode].append(int(num))
            else:
                signal_dict[electrode] = [int(num)]

    return signal_dict