dualmicro_functions.py

    cond7_idx = event_labels0 == "imitate, follower"
    cond67_idx = cond6_idx+cond7_idx
    micro_data[cond67_idx] = np.concatenate([micro_data2[cond67_idx],
                                            micro_data1[cond67_idx]],axis=1)
    # Transform data to correct shape
    arr_shape = micro_data.shape # get shape
    micro_data = micro_data.swapaxes(1,2) # swap ch and time axis
    micro_data = micro_data.reshape(arr_shape[0]*arr_shape[2],arr_shape[1]) # reshape by combining epochs and times
    # Filter the data in alpha
    # Notice it is done only after combining epochs and time, as filter length
    # would be too long for 1s epochs. The filter function wants time on the last axis
    micro_data = micro_data.transpose()
    micro_data_alpha = mne.filter.filter_data(micro_data, sfreq, freq_range[0], freq_range[1])
    # Reverse the shape
    micro_data_alpha = micro_data_alpha.transpose()
    
    return micro_data_alpha

def pseudo_pair_dualmicro_backfitting(micro_data, prototype_map, events, n_maps, sfreq):
    """ Backfit microstate labels based on prototype map previously determined """
    assert prototype_map.shape[0] == n_maps, "Template map size is not equal to n_maps"
    # Load micro data
    micro_data0 = micro_data
    n_ch = micro_data.shape[1]
    # Estimate global explained variance (GEV) by the averaged template maps
    # on the global field potential peaks in our data
    # Find GFPs
    gfp = np.std(micro_data0, axis=1)
    gfp_peaks = locmax(gfp)
    gfp_values = gfp[gfp_peaks]
    gfp2 = np.sum(gfp_values**2) # normalizing constant in GEV
    # Calculate spatial correlation
    # Using absolute value of the topographies to obtain polarity invariance
    # Since we are working on two-person microstates, we also need to test
    # the different configurations. There are 4 in total, but we only need to try 2
    tmp_maps = np.split(prototype_map, [n_ch//2,n_ch], axis=1)
    all_polarity_combinations = [np.concatenate([tmp_maps[0],tmp_maps[1]],axis=1),
                                 np.concatenate([tmp_maps[0],-tmp_maps[1]],axis=1)]
    C_arr = np.zeros((micro_data0.shape[0],n_maps,len(all_polarity_combinations)))
    for p in range(len(all_polarity_combinations)):
        C_arr[:,:,p] = np.dot(micro_data0, all_polarity_combinations[p].T)
        # rescale cov
        C_arr[:,:,p] /= (n_ch*np.outer(gfp, np.std(prototype_map, axis=1)))
    # Get C as the highest correlation independent of polarity
    # Take max for all polarity configurations and then argmax to find label
    C = np.sqrt(np.max(C_arr**2,axis=2)) # notice sign is lost here, but it is only used as C^2 later so it is fine
    L = np.argmax(C**2, axis=1)

    L_gfp = L[gfp_peaks]
    C_gfp = C[gfp_peaks]

    gev = np.zeros(n_maps)
    for k in range(n_maps):
        r = L_gfp==k
        gev[k] = np.sum(gfp_values[r]**2 * C_gfp[r,k]**2)/gfp2

    # Reshape labels back to epoch, time
    L_reshaped = L.reshape(len(events[0]),sfreq)

    return L_reshaped, gev

def dualmicro_fit_all_pseudo_pair_feature_computation(i, n_maps, backfit_results, sfreq,
                                                      event_id, collapsed_event_id):
    """
    Overview of common microstate features:
        1. Average duration a given microstate remains stable (Dur)
        2. Frequency occurrence, independent of individual duration (Occ)
            Average number of times a microstate becomes dominant per second
        3. Ratio of total Time Covered (TCo)
        4. Transition probabilities (TMx)
        5. Ratio of shannon entropy relative to theoretical max chaos (Ent)
    
    Parameters
    ----------
    i : int
        The index.
    n_maps : int
        The number of maps (clusters) used.
    backfit_results : list
        The estimated back-fitted microstates.
    sfreq : int
        The sampling frequency.
    event_id : dict
        The mappings between event id and condition.
    collapsed_brain_event_id : dict
        The renamed mappings between event id and collapsed conditions

    Returns
    -------
    m_labels : list of two np.array
        The microstate sequence (labels) time series in the format (epoch, time).
    events : list of two pd.DataFrame
        The events corresponding to each epoch.
    MFeatures : list
        List of arrays of each microstate feature.

    """
    m_labels = backfit_results[1][i]
    events = backfit_results[3][i]
    
    # Since they are synchronized already, just take the first and fix
    # epoch idx column with the new idx for the synchronized labels
    events = events[0]
    
    # Update in v7
    # # Collapse event_id 6 and 7 to 4 and 5
    # events.loc[events["Event_id"] == 6,["Event_id","Event_label"]] = [4, "observe, actor"]
    # events.loc[events["Event_id"] == 7,["Event_id","Event_label"]] = [5, "imitate, leader"]
    # The microstate clustering was performed on flipped (collapsed) events,
    # but I will compute the features on the 8 trials to avoid the flipping
    # effect and then collapse by averaging afterwards
    
    # Reshape m_labels to (epoch, time)
    m_labels = m_labels.reshape(len(events),sfreq)
    
    # Remove pre-trial epochs
    pre_trial_epochs = events["Trial_start_time"] < 0
    m_labels = m_labels[np.invert(pre_trial_epochs)]
    events = events.loc[np.invert(pre_trial_epochs)].reset_index(drop=True)
    events["Epoch_idx"] = events.index
    
    # Pre-allocate memory
    Dur_arr = np.zeros((len(event_id),n_maps)); Dur_arr.fill(np.nan)
    Occ_arr = np.zeros((len(event_id),n_maps)); Occ_arr.fill(np.nan)
    TCo_arr = np.zeros((len(event_id),n_maps)); TCo_arr.fill(np.nan)
    TMx_arr = np.zeros((len(event_id),n_maps,n_maps))
    Ent_arr = np.zeros(len(event_id))
    for e in range(len(event_id)):
        ev_idx = list(event_id.values())[e]
        ep_idx = events["Epoch_idx"][events["Event_id"] == ev_idx]
        trial_numbers0 = np.unique(events["Trial_number"][ep_idx])
        if len(trial_numbers0) > 8:
            # Compute the first 8 trials and the last 8 separately, before
            # taking the average to be consistent with asymmetric trials
            Dur_arr0 = np.zeros((2,n_maps))
            Occ_arr0 = np.zeros((2,n_maps))
            TCo_arr0 = np.zeros((2,n_maps))
            TMx_arr0 = np.zeros((2,n_maps,n_maps))
            Ent_arr0 = np.zeros(2)
            trial_numbers_split = np.array_split(trial_numbers0, 2)
            # Array_split is used in cases where a trial might have been
            # dropped, hence the split is only 8/7
            for s in range(len(trial_numbers_split)):
                ep_idx0 = events.loc[(events["Event_id"] == ev_idx)&
                                    (events["Trial_number"].isin(trial_numbers_split[s])),
                                    "Epoch_idx"]
                m_labels0 = m_labels[ep_idx0,:]
                m_labels0_flat = m_labels0.reshape(m_labels0.shape[0]*m_labels0.shape[1])
                # Compute duration, occurrence and time covered
                l_, d_ = microstate_run_length_encoding(m_labels0_flat)
                # For each microstate
                for ii in range(n_maps):
                    if np.isnan(np.nanmean(d_[l_==ii])):
                        # The specific microstate did not occur at all for this event
                        Dur_arr0[s,ii] = 0
                        Occ_arr0[s,ii] = 0
                        TCo_arr0[s,ii] = 0
                    else:
                        Dur_arr0[s,ii] = np.mean(d_[l_==ii]) * 1000/sfreq # convert to ms
                        Occ_arr0[s,ii] = len(d_[l_==ii])/len(d_) * sfreq
                        TCo_arr0[s,ii] = np.sum(d_[l_==ii])/np.sum(d_)
                # Compute transition matrix
                TMx_arr0[s] = T_empirical(m_labels0_flat, n_maps)
                # Compute Shannon Entropy relative to max
                Ent_arr0[s] = H_1(m_labels0_flat, n_maps)/np.log(float(n_maps))
            # Average over the splits
            Dur_arr[e] = np.mean(Dur_arr0, axis=0)
            Occ_arr[e] = np.mean(Occ_arr0, axis=0)
            TCo_arr[e] = np.mean(TCo_arr0, axis=0)
            TMx_arr[e] = np.mean(TMx_arr0, axis=0)
            Ent_arr[e] = np.mean(Ent_arr0, axis=0)
            
        else:
            m_labels0 = m_labels[ep_idx,:]
            m_labels0_flat = m_labels0.reshape(m_labels0.shape[0]*m_labels0.shape[1])
            # Compute duration, occurrence and time covered
            l_, d_ = microstate_run_length_encoding(m_labels0_flat)
            # For each microstate
            for ii in range(n_maps):
                if np.isnan(np.nanmean(d_[l_==ii])):
                    # The specific microstate did not occur at all for this event
                    Dur_arr[e,ii] = 0
                    Occ_arr[e,ii] = 0
                    TCo_arr[e,ii] = 0
                else:
                    Dur_arr[e,ii] = np.mean(d_[l_==ii]) * 1000/sfreq # convert to ms
                    Occ_arr[e,ii] = len(d_[l_==ii])/len(d_) * sfreq
                    TCo_arr[e,ii] = np.sum(d_[l_==ii])/np.sum(d_)
            # Compute transition matrix
            TMx_arr[e] = T_empirical(m_labels0_flat, n_maps)
            # Compute Shannon Entropy relative to max
            Ent_arr[e] = H_1(m_labels0_flat, n_maps)/np.log(float(n_maps))
    # Combine all features in a list
    MFeatures = [Dur_arr,Occ_arr,TCo_arr,TMx_arr,Ent_arr]
    
    # Update in v7 - Collapse after computations in single events
    Dur_arr2 = np.zeros((len(collapsed_event_id),n_maps))
    Occ_arr2 = np.zeros((len(collapsed_event_id),n_maps))
    TCo_arr2 = np.zeros((len(collapsed_event_id),n_maps))
    TMx_arr2 = np.zeros((len(collapsed_event_id),n_maps,n_maps))
    Ent_arr2 = np.zeros(len(collapsed_event_id))
    MFeatures2 = [Dur_arr2,Occ_arr2,TCo_arr2,TMx_arr2,Ent_arr2]
    
    for f in range(len(MFeatures2)):
        tmp_feat = MFeatures[f]
        tmp_feat2 = MFeatures2[f]
        for e in range(len(collapsed_event_id)):
            ee = list(collapsed_event_id.keys())[e]
            if (ee == 'observer_actor'):
                old_ev_idx1 = list(event_id.keys()).index("observe, actor")
                old_ev_idx2 = list(event_id.keys()).index("observe, observer")
                tmp_feat2[e] = np.mean([tmp_feat[old_ev_idx1],tmp_feat[old_ev_idx2]], axis=0)
            elif ee == 'follower_leader':
                old_ev_idx1 = list(event_id.keys()).index("imitate, leader")
                old_ev_idx2 = list(event_id.keys()).index("imitate, follower")
                tmp_feat2[e] = np.mean([tmp_feat[old_ev_idx1],tmp_feat[old_ev_idx2]], axis=0)
            else:
                new_ev_idx = list(collapsed_event_id.values())[e]
                old_ev_idx = list(event_id.values()).index(new_ev_idx)
                tmp_feat2[e] = tmp_feat[old_ev_idx]
    
    return m_labels, events, MFeatures2

def label_to_walk_direction(labels):
    directions = np.zeros(len(labels))
    directions[pd.Series(labels).isin([0,1,2,3])] = 1 # A/B/C/D assigned to 1
    directions[pd.Series(labels).isin([4,5,6,7])] = -1 # E/F/G/H assigned to -1
    assert np.sum(directions==0)==0
    return directions

def compute_dualmicro_DFA(i, microstate_results, trialinfo_list, sfreq, window_sizes, event_id, collapsed_event_id, overlap=True):
    """
    See Hardstone et al, 2012 for more info
    Perform DFA
        1 Compute cumulative sum of time series to create signal profile
        2 Define set of window sizes (see below)
        3 Remove the linear trend using least-squares for each window
        4 Calculate standard deviation for each window and take the mean
        5 Plot fluctuation function (Standard deviation) as function
          for all window sizes, on double logarithmic scale
        6 The DFA exponent alpha correspond to Hurst exponent
          f(L) = sd = L^alpha (with alpha as linear coefficient in log plot)
      
    Parameters
    ----------
    i : int
        The index.
    microstate_results : list
        The estimated microstates.
    trialinfo_list : list
        List with trial informations.
    sfreq : int
        The sampling frequency.
    window_sizes : np.array
        Window sizes to compute DFA over.
    event_id : dict
        The mappings between event id and condition.
    collapsed_brain_event_id : dict
        The renamed mappings between event id and collapsed conditions
    overlap : bool, optional
        Boolean to determine whether to use overlapping windows. The default is True.

    Returns
    -------
    dfa_array : np.array
        The Hurst Exponents.
    fluctuations : np.array
        The fluctuations estimated at different window sizes.

    """
    # Load all microstate results    
    pair_idx = microstate_results[5]
    pair_indices = pair_idx[i], pair_idx[i+1]
    m_labels = microstate_results[1][pair_indices[0]:pair_indices[1]]

    events = trialinfo_list[2][i]
    # Since they are synchronized already, just take the first and fix
    # epoch idx column with the new idx for the synchronized labels
    events = events[0]
    
    # Reshape m_labels to (epoch, time)
    m_labels = m_labels.reshape(len(events),sfreq)
    
    # Remove pre-trial epochs
    pre_trial_epochs = events["Trial_start_time"] < 0
    m_labels = m_labels[np.invert(pre_trial_epochs)]
    events = events.loc[np.invert(pre_trial_epochs)].reset_index(drop=True)
    events["Epoch_idx"] = events.index
    
    # Pre-allocate memory
    dfa_array = np.zeros((len(event_id)))
    dfa_array[:] = np.nan
    fluctuations = np.zeros((len(event_id),len(window_sizes)))
    fluctuations[:] = np.nan
    # As I am working with linear regression on log-values, I can just take
    # the avg of the x, y coordinates and the coefficients will correspond
    # to the avg for each separate linear regression

    for e in range(len(event_id)):
        ev_idx = list(event_id.values())[e]
        ep_idx = events["Epoch_idx"][events["Event_id"] == ev_idx]
        trial_numbers0 = np.unique(events["Trial_number"][ep_idx])
        if len(trial_numbers0) > 8:
            # Compute the first 8 trials and the last 8 separately, before
            # taking the average to be consistent with asymmetric trials
            dfa0 = []
            fluct0 = np.zeros((2,len(window_sizes)))
            trial_numbers_split = np.array_split(trial_numbers0, 2)
            # Array_split is used in cases where a trial might have been
            # dropped, hence the split is only 8/7
            for s in range(len(trial_numbers_split)):
                ep_idx0 = events.loc[(events["Event_id"] == ev_idx)&
                                    (events["Trial_number"].isin(trial_numbers_split[s])),
                                    "Epoch_idx"]
                m_labels0 = m_labels[ep_idx0,:]
                m_labels0_flat = m_labels0.reshape(m_labels0.shape[0]*m_labels0.shape[1])
                # Convert the labels to the directons for computation of cumulative
                # sum in order to get the signal profile
                data = label_to_walk_direction(m_labels0_flat)
                # plt.plot(np.arange(1,len(data)+1,1),np.cumsum(data)) # plot signal profile
                # Estimate and save DFA
                dfa, fluctuations_2d = nolds.dfa(data,nvals=window_sizes,overlap=overlap,debug_data=True)
                dfa0.append(dfa)
                fluct0[s] = fluctuations_2d[1]
            # Average the two DFA estimations
            dfa_array[e] = np.mean(dfa0)
            fluctuations[e,:] = np.mean(fluct0, axis=0)
            
        else:
            m_labels0 = m_labels[ep_idx,:]
            m_labels0_flat = m_labels0.reshape(m_labels0.shape[0]*m_labels0.shape[1])
            # Convert the labels to the directons for computation of cumulative
            # sum in order to get the signal profile
            data = label_to_walk_direction(m_labels0_flat)
            # Estimate and save DFA
            dfa_array[e],fluctuations_2d = nolds.dfa(data,nvals=window_sizes,overlap=overlap,debug_data=True)
            fluctuations[e,:] = fluctuations_2d[1]
    
    # Update in DFA v2 - Collapse after computation of DFA
    dfa_array2 = np.zeros((len(collapsed_event_id)))
    dfa_array2[:] = np.nan
    fluctuations2 = np.zeros((len(collapsed_event_id),len(window_sizes)))
    fluctuations2[:] = np.nan
    for e in range(len(collapsed_event_id)):
        ee = list(collapsed_event_id.keys())[e]
        if ee == 'observer_actor':
            old_ev_idx1 = list(event_id.keys()).index("observe, actor")
            old_ev_idx2 = list(event_id.keys()).index("observe, observer")
            dfa_array2[e] = np.mean([dfa_array[old_ev_idx1],dfa_array[old_ev_idx2]])
            fluctuations2[e] = np.mean([fluctuations[old_ev_idx1],fluctuations[old_ev_idx2]], axis=0)
        elif ee == 'follower_leader':
            old_ev_idx1 = list(event_id.keys()).index("imitate, leader")
            old_ev_idx2 = list(event_id.keys()).index("imitate, follower")
            dfa_array2[e] = np.mean([dfa_array[old_ev_idx1],dfa_array[old_ev_idx2]])
            fluctuations2[e] = np.mean([fluctuations[old_ev_idx1],fluctuations[old_ev_idx2]], axis=0)
        else:
            new_ev_idx = list(collapsed_event_id.values())[e]
            old_ev_idx = list(event_id.values()).index(new_ev_idx)
            dfa_array2[e] = dfa_array[old_ev_idx]
            fluctuations2[e] = fluctuations[old_ev_idx]
    
    return dfa_array2, fluctuations2


def compute_dualmicro_DFA_pseudo(i, backfit_results, sfreq, window_sizes, event_id, collapsed_event_id, overlap=True):
    """
    See Hardstone et al, 2012 for more info
    Perform DFA
        1 Compute cumulative sum of time series to create signal profile
        2 Define set of window sizes (see below)
        3 Remove the linear trend using least-squares for each window
        4 Calculate standard deviation for each window and take the mean
        5 Plot fluctuation function (Standard deviation) as function
          for all window sizes, on double logarithmic scale
        6 The DFA exponent alpha correspond to Hurst exponent
          f(L) = sd = L^alpha (with alpha as linear coefficient in log plot)
      
    Parameters
    ----------
    i : int
        The index.
    backfit_results : list
        The estimated back-fitted microstates.
    sfreq : int
        The sampling frequency.
    window_sizes : np.array
        Window sizes to compute DFA over.
    event_id : dict
        The mappings between event id and condition.
    collapsed_brain_event_id : dict
        The renamed mappings between event id and collapsed conditions
    overlap : bool, optional
        Boolean to determine whether to use overlapping windows. The default is True.

    Returns
    -------
    dfa_array : np.array
        The Hurst Exponents.
    fluctuations : np.array
        The fluctuations estimated at different window sizes.

    """
    m_labels = backfit_results[1][i]
    events = backfit_results[3][i]
    
    # Since they are synchronized already, just take the first and fix
    # epoch idx column with the new idx for the synchronized labels
    events = events[0]
    
    # Update in DFA v2
    # # Collapse event_id 6 and 7 to 4 and 5
    # events.loc[events["Event_id"] == 6,["Event_id","Event_label"]] = [4, "observe, actor"]
    # events.loc[events["Event_id"] == 7,["Event_id","Event_label"]] = [5, "imitate, leader"]
    # The microstate clustering was performed on flipped (collapsed) events,
    # but I will compute the DFA on the 8 trials to avoid the flipping
    # effect on the linear detrending, and then collapse by averaging
    # after DFA computation

    # Reshape m_labels to (epoch, time)
    m_labels = m_labels.reshape(len(events),sfreq)
    
    # Remove pre-trial epochs
    pre_trial_epochs = events["Trial_start_time"] < 0
    m_labels = m_labels[np.invert(pre_trial_epochs)]
    events = events.loc[np.invert(pre_trial_epochs)].reset_index(drop=True)
    events["Epoch_idx"] = events.index
    
    # Pre-allocate memory
    dfa_array = np.zeros((len(event_id)))
    dfa_array[:] = np.nan
    fluctuations = np.zeros((len(event_id),len(window_sizes)))
    fluctuations[:] = np.nan
    # As I am working with linear regression on log-values, I can just take
    # the avg of the x, y coordinates and the coefficients will correspond
    # to the avg for each separate linear regression

    for e in range(len(event_id)):
        ev_idx = list(event_id.values())[e]
        ep_idx = events["Epoch_idx"][events["Event_id"] == ev_idx]
        trial_numbers0 = np.unique(events["Trial_number"][ep_idx])
        if len(trial_numbers0) > 8:
            # Compute the first 8 trials and the last 8 separately, before
            # taking the average to be consistent with asymmetric trials
            dfa0 = []
            fluct0 = np.zeros((2,len(window_sizes)))
            trial_numbers_split = np.array_split(trial_numbers0, 2)
            # Array_split is used in cases where a trial might have been
            # dropped, hence the split is only 8/7
            for s in range(len(trial_numbers_split)):
                ep_idx0 = events.loc[(events["Event_id"] == ev_idx)&
                                    (events["Trial_number"].isin(trial_numbers_split[s])),
                                    "Epoch_idx"]
                m_labels0 = m_labels[ep_idx0,:]
                m_labels0_flat = m_labels0.reshape(m_labels0.shape[0]*m_labels0.shape[1])
                # Convert the labels to the directons for computation of cumulative
                # sum in order to get the signal profile
                data = label_to_walk_direction(m_labels0_flat)
                # plt.plot(np.arange(1,len(data)+1,1),np.cumsum(data)) # plot signal profile
                # Estimate and save DFA
                dfa, fluctuations_2d = nolds.dfa(data,nvals=window_sizes,overlap=overlap,debug_data=True)
                dfa0.append(dfa)
                fluct0[s] = fluctuations_2d[1]
            # Average the two DFA estimations
            dfa_array[e] = np.mean(dfa0)
            fluctuations[e,:] = np.mean(fluct0, axis=0)
            
        else:
            m_labels0 = m_labels[ep_idx,:]
            m_labels0_flat = m_labels0.reshape(m_labels0.shape[0]*m_labels0.shape[1])
            # Convert the labels to the directons for computation of cumulative
            # sum in order to get the signal profile
            data = label_to_walk_direction(m_labels0_flat)
            # Estimate and save DFA
            dfa_array[e],fluctuations_2d = nolds.dfa(data,nvals=window_sizes,overlap=overlap,debug_data=True)
            fluctuations[e,:] = fluctuations_2d[1]
    
    # Update in DFA v2 - Collapse after computation of DFA
    dfa_array2 = np.zeros((len(collapsed_event_id)))
    dfa_array2[:] = np.nan
    fluctuations2 = np.zeros((len(collapsed_event_id),len(window_sizes)))
    fluctuations2[:] = np.nan
    for e in range(len(collapsed_event_id)):
        ee = list(collapsed_event_id.keys())[e]
        if ee == 'observer_actor':
            old_ev_idx1 = list(event_id.keys()).index("observe, actor")
            old_ev_idx2 = list(event_id.keys()).index("observe, observer")
            dfa_array2[e] = np.mean([dfa_array[old_ev_idx1],dfa_array[old_ev_idx2]])
            fluctuations2[e] = np.mean([fluctuations[old_ev_idx1],fluctuations[old_ev_idx2]], axis=0)
        elif ee == 'follower_leader':
            old_ev_idx1 = list(event_id.keys()).index("imitate, leader")
            old_ev_idx2 = list(event_id.keys()).index("imitate, follower")
            dfa_array2[e] = np.mean([dfa_array[old_ev_idx1],dfa_array[old_ev_idx2]])
            fluctuations2[e] = np.mean([fluctuations[old_ev_idx1],fluctuations[old_ev_idx2]], axis=0)
        else:
            new_ev_idx = list(collapsed_event_id.values())[e]
            old_ev_idx = list(event_id.values()).index(new_ev_idx)
            dfa_array2[e] = dfa_array[old_ev_idx]
            fluctuations2[e] = fluctuations[old_ev_idx]
    
    return dfa_array2, fluctuations2

def similar_interbrain_microstates(x, y):
    assert len(x) == len(y)
    ratio_in_similar_state = sum(x == y)/len(x)
    return ratio_in_similar_state

def shift_label_time_series(x, y, shift):
    assert len(x) == len(y)
    x = pd.Series(x)
    y = pd.Series(y)
    # Shift the data
    y = y.shift(shift)
    # Drop the NaN
    drop_idx = np.where(y.isna())[0]
    y = y.drop(drop_idx)
    x = x.drop(drop_idx)
    assert len(y) == len(x)
    assert sum(y.isna()) == 0
    return x.to_numpy().astype(int), y.to_numpy().astype(int)

def shifted_interbrain_microstate_feature_computation(i, n_maps, microstate_results,
       trialinfo_list, sfreq, event_id, collapsed_event_id, lag_search_range, lag_interval):
    """
    Compute inter-brain features, but for time-lagged microstate label time series.
    Similar to cross-correlation, but maximizing the amount of synchronized
    (similar) microstates at a given lag
    
    Parameters
    ----------
    i : int
        The index.
    n_maps : int
        The number of maps (clusters) used.
    microstate_results : list
        The estimated microstates.
    trialinfo_list : list
        List with trial informations.
    sfreq : int
        The sampling frequency.
    event_id : dict
        The mappings between event id and condition.
    collapsed_brain_event_id : dict
        The renamed mappings between event id and collapsed conditions
    lag_search_range : float
        The lag in both directions the data is being shifted
    lag_interval : np.array
        The sample points corresponding to the lags the data is being shifted
    
    Returns
    -------
    m_labels : list of two np.array
        The microstate sequence (labels) time series in the format (epoch, time).
    events : list of two pd.DataFrame
        The events corresponding to each epoch.
    MFeatures : list
        List of arrays of each microstate feature.
    shift_info : list
        List containing the information regarding the shift, e.g. the cross-
        similarity at different time lags and the optimal time lag
    
    """
    
    # Here i refers to the pair in range(n_subjects//2)
    sub_idx = microstate_results[5]
    
    # Get the microstate labels and events for participant 1
    subject_indices1 = sub_idx[2*i], sub_idx[2*i+1]
    m_labels1 = microstate_results[1][subject_indices1[0]:subject_indices1[1]]
    # Get the trialinfo with conditions
    Subject1, trialinfo1 = trialinfo_list[2*i]
    # Convert to dataframe
    event_id_inv = {v: k for k, v in event_id.items()} # Inverse the event id
    events1 = pd.DataFrame(trialinfo1,columns=["Event_id","Trial_number","Trial_start_time"])
    events1["Event_label"] = events1["Event_id"].replace(event_id_inv)
    events1 = events1.reset_index().rename(columns={"index":"Epoch_idx"})
    # Reshape m_labels to (epoch, time)
    m_labels1 = m_labels1.reshape(len(trialinfo1),sfreq)
    
    # Get the microstate labels and events for participant 2
    subject_indices2 = sub_idx[2*i+1], sub_idx[2*i+2]
    m_labels2 = microstate_results[1][subject_indices2[0]:subject_indices2[1]]
    # Get the trialinfo with conditions
    Subject2, trialinfo2 = trialinfo_list[2*i+1]
    # Convert to dataframe
    events2 = pd.DataFrame(trialinfo2,columns=["Event_id","Trial_number","Trial_start_time"])
    events2["Event_label"] = events2["Event_id"].replace(event_id_inv)
    events2 = events2.reset_index().rename(columns={"index":"Epoch_idx"})
    # Reshape m_labels to (epoch, time)
    m_labels2 = m_labels2.reshape(len(trialinfo2),sfreq)
    
    # check that the participants from a pair was loaded
    assert Subject2-1 == Subject1
    # Check that the same amount of event types are present
    assert trialinfo1[-1,1] == trialinfo2[-1,1]
    # Synchronize the events from the pair based on timing info
    # By trimming the epochs to only include the epochs that are present
    # in both participants of the pair
    # Initialize with an array filled with a unique number not in use
    sync_m_labels1 = np.zeros_like(m_labels1); sync_m_labels1.fill(9999)
    sync_m_labels2 = np.zeros_like(m_labels2); sync_m_labels2.fill(9999)
    for t in np.unique(trialinfo1[:,1]):
        t_idx1 = np.where(trialinfo1[:,1] == t)[0]
        t_idx2 = np.where(trialinfo2[:,1] == t)[0]
        # Get the timings for the epochs for the specific trial t
        t_timings1 = trialinfo1[t_idx1,2]
        t_timings2 = trialinfo2[t_idx2,2]
        timings_intersect = np.intersect1d(t_timings1,t_timings2)
        # Get the indices where the timings matches
        t_idx_match1 = t_idx1[pd.Series(t_timings1).isin(timings_intersect)]
        t_idx_match2 = t_idx2[pd.Series(t_timings2).isin(timings_intersect)]
        # Get the actual values from the synchronized epochs
        sync_m_labels1[t_idx_match1] = m_labels1[t_idx_match1]
        sync_m_labels2[t_idx_match2] = m_labels2[t_idx_match2]
    # Find the epochs that were asynchronous, which have to be trimmed
    asynch_epochs1 = np.unique(np.where(sync_m_labels1==9999)[0])
    asynch_epochs2 = np.unique(np.where(sync_m_labels2==9999)[0])
    # Trim/delete the asynchronous epochs
    sync_m_labels1 = np.delete(sync_m_labels1,asynch_epochs1,axis=0)
    sync_m_labels2 = np.delete(sync_m_labels2,asynch_epochs2,axis=0)
    assert len(sync_m_labels1) == len(sync_m_labels2) # check the amount of synchronized epochs are equal
    # Fix events
    sync_events1 = events1.drop(asynch_epochs1,axis=0).reset_index(drop=True)
    sync_events2 = events2.drop(asynch_epochs2,axis=0).reset_index(drop=True)
    
    # Since they are synchronized already, just take the first and fix
    # epoch idx column with the new idx for the synchronized labels
    events = sync_events1
    events["Epoch_idx"] = events.index
    # Notice that for the intrabrain fit all alpha v6 I only corrected
    # ppn2. So by taking ppn1 I get the original event labels.
    
    # Remove pre-trial epochs
    pre_trial_epochs = events["Trial_start_time"] < 0
    sync_m_labels1 = sync_m_labels1[np.invert(pre_trial_epochs)]
    sync_m_labels2 = sync_m_labels2[np.invert(pre_trial_epochs)]
    events = events.loc[np.invert(pre_trial_epochs)].reset_index(drop=True)
    events["Epoch_idx"] = events.index
    
    # Determine the optimal lag for highest interbrain state synchrony
    # For each condition - using the collapsed dictionary
    # To get one value of optimal lag for e.g. follower-leader and observer-actor
    # Collapse event_id 6 and 7 to 4 and 5
    events_collapsed = events.copy()
    events_collapsed.loc[events_collapsed["Event_id"] == 6,["Event_id","Event_label"]] = [4, "observe, actor"]
    events_collapsed.loc[events_collapsed["Event_id"] == 7,["Event_id","Event_label"]] = [5, "imitate, leader"]
    normal_events_to_collapsed_map = {"1":1, "2":2, "3":3, "4":4, "5":5, "6":4, "7":5, "8":8}
    
    shift_info = [0]*len(collapsed_event_id)
    for e in range(len(collapsed_event_id)):
        ev_idx = list(collapsed_event_id.values())[e]
        ep_idx = events_collapsed["Epoch_idx"][events_collapsed["Event_id"] == ev_idx]
        trial_numbers0 = np.unique(events_collapsed["Trial_number"][ep_idx])
        
        lts1 = sync_m_labels1[ep_idx,:]
        lts2 = sync_m_labels2[ep_idx,:]
        # Flatten the labels
        lts1_flat = pd.Series(lts1.ravel())
        lts2_flat = pd.Series(lts2.ravel())
        # Compute cross similarity
        cross_similarity = [lts1_flat.corr(lts2_flat.shift(lag), method=similar_interbrain_microstates) for lag in lag_interval]
        # Original ratio of time in similar states (time = 0)
        t_zero_ratio_in_similar_microstate = cross_similarity[lag_search_range]
        # The optimal lag
        opt_lag = lag_interval[np.argmax(cross_similarity)]
        opt_ratio_in_similar_microstate = cross_similarity[np.argmax(cross_similarity)]
        # Shift the label time series with the optimal lag prior to computation of interbrain features
        # Save the features
        shift_info[e] = [t_zero_ratio_in_similar_microstate, opt_lag, opt_ratio_in_similar_microstate, cross_similarity]
    
    # Pre-allocate memory
    Dur_arr = np.zeros((len(event_id),n_maps+1)); Dur_arr.fill(np.nan)
    Occ_arr = np.zeros((len(event_id),n_maps+1)); Occ_arr.fill(np.nan)
    TCo_arr = np.zeros((len(event_id),n_maps+1)); TCo_arr.fill(np.nan)
    TMx_arr = np.zeros((len(event_id),n_maps+1,n_maps+1))
    Ent_arr = np.zeros(len(event_id))
    for e in range(len(event_id)):
        ev_idx = list(event_id.values())[e]
        ep_idx = events["Epoch_idx"][events["Event_id"] == ev_idx]
        trial_numbers0 = np.unique(events["Trial_number"][ep_idx])
        # Get the opt_lag from collapsed events
        ev_collapsed_idx = normal_events_to_collapsed_map[str(ev_idx)]
        shift_info_idx = list(collapsed_event_id.values()).index(ev_collapsed_idx)
        opt_lag = shift_info[shift_info_idx][1]
        
        if len(trial_numbers0) > 8:
            # Compute the first 8 trials and the last 8 separately, before
            # taking the average to be consistent with asymmetric trials
            Dur_arr0 = np.zeros((2,n_maps+1))
            Occ_arr0 = np.zeros((2,n_maps+1))
            TCo_arr0 = np.zeros((2,n_maps+1))
            TMx_arr0 = np.zeros((2,n_maps+1,n_maps+1))
            Ent_arr0 = np.zeros(2)
            trial_numbers_split = np.array_split(trial_numbers0, 2)
            # Array_split is used in cases where a trial might have been
            # dropped, hence the split is only 8/7
            for s in range(len(trial_numbers_split)):
                ep_idx0 = events.loc[(events["Event_id"] == ev_idx)&
                                    (events["Trial_number"].isin(trial_numbers_split[s])),
                                    "Epoch_idx"]
                # Get the microstate labels
                m_labels10 = sync_m_labels1[ep_idx0,:]
                m_labels20 = sync_m_labels2[ep_idx0,:]
                # Flatten the labels
                m_labels10_flat = m_labels10.reshape(m_labels10.shape[0]*m_labels10.shape[1])
                m_labels20_flat = m_labels20.reshape(m_labels20.shape[0]*m_labels20.shape[1])
                # Shift the label time-series to maximize interbrain synchrony
                m_labels10_shifted, m_labels20_shifted = shift_label_time_series(m_labels10_flat, m_labels20_flat, opt_lag)
                
                # Compute average duration of common microstate
                # Output: label and duration of common microstate. Label -1 is used
                # for not common microstate
                l_, d_ = interbrain_microstate_run_length_encoding(m_labels10_shifted,m_labels20_shifted)
                # For each microstate
                for ii in range(n_maps+1):
                    if np.isnan(np.nanmean(d_[l_==ii-1])):
                        # The specific microstate did not occur at all for this event
                        Dur_arr0[s,ii] = 0
                        Occ_arr0[s,ii] = 0
                        TCo_arr0[s,ii] = 0
                    else:
                        Dur_arr0[s,ii] = np.mean(d_[l_==ii-1]) * 1000/sfreq # convert to ms
                        Occ_arr0[s,ii] = len(d_[l_==ii-1])/len(d_) * sfreq
                        TCo_arr0[s,ii] = np.sum(d_[l_==ii-1])/np.sum(d_)
                # Compute transition matrix
                TMx_arr0[s] = interbrain_T_matrix(m_labels10_shifted,m_labels20_shifted,n_maps)
                # Compute Joint Shannon Entropy relative to max
                # Max is the sum of max individual entropies
                Ent_arr0[s] = H_2(m_labels10_shifted,m_labels20_shifted,n_maps)/(2*np.log(float(n_maps)))
            # Average over the splits
            Dur_arr[e] = np.mean(Dur_arr0, axis=0)
            Occ_arr[e] = np.mean(Occ_arr0, axis=0)
            TCo_arr[e] = np.mean(TCo_arr0, axis=0)
            TMx_arr[e] = np.mean(TMx_arr0, axis=0)
            Ent_arr[e] = np.mean(Ent_arr0, axis=0)
            
        else:
            # Get the microstate labels
            m_labels10 = sync_m_labels1[ep_idx,:]
            m_labels20 = sync_m_labels2[ep_idx,:]
            # Flatten the labels
            m_labels10_flat = m_labels10.reshape(m_labels10.shape[0]*m_labels10.shape[1])
            m_labels20_flat = m_labels20.reshape(m_labels20.shape[0]*m_labels20.shape[1])
            # Shift the label time-series to maximize interbrain synchrony
            m_labels10_shifted, m_labels20_shifted = shift_label_time_series(m_labels10_flat, m_labels20_flat, opt_lag)
            
            # Compute average duration of common microstate
            # Output: label and duration of common microstate. Label -1 is used
            # for not common microstate
            l_, d_ = interbrain_microstate_run_length_encoding(m_labels10_shifted,m_labels20_shifted)
            # For each microstate
            for ii in range(n_maps+1):
                if np.isnan(np.nanmean(d_[l_==ii-1])):
                    # The specific microstate did not occur at all for this event
                    Dur_arr[e,ii] = 0
                    Occ_arr[e,ii] = 0
                    TCo_arr[e,ii] = 0
                else:
                    Dur_arr[e,ii] = np.mean(d_[l_==ii-1]) * 1000/sfreq # convert to ms
                    Occ_arr[e,ii] = len(d_[l_==ii-1])/len(d_) * sfreq
                    TCo_arr[e,ii] = np.sum(d_[l_==ii-1])/np.sum(d_)
            # Compute transition matrix
            TMx_arr[e] = interbrain_T_matrix(m_labels10_shifted,m_labels20_shifted,n_maps)
            # Compute Joint Shannon Entropy relative to max
            # Max is the sum of max individual entropies
            Ent_arr[e] = H_2(m_labels10_shifted,m_labels20_shifted,n_maps)/(2*np.log(float(n_maps)))
            
    # Combine all features in a list
    MFeatures = [Dur_arr,Occ_arr,TCo_arr,TMx_arr,Ent_arr]
    
    # Update in v7 - Collapse after computations in single events
    Dur_arr2 = np.zeros((len(collapsed_event_id),n_maps+1))
    Occ_arr2 = np.zeros((len(collapsed_event_id),n_maps+1))
    TCo_arr2 = np.zeros((len(collapsed_event_id),n_maps+1))
    TMx_arr2 = np.zeros((len(collapsed_event_id),n_maps+1,n_maps+1))
    Ent_arr2 = np.zeros(len(collapsed_event_id))
    MFeatures2 = [Dur_arr2,Occ_arr2,TCo_arr2,TMx_arr2,Ent_arr2]
    
    for f in range(len(MFeatures2)):
        tmp_feat = MFeatures[f]
        tmp_feat2 = MFeatures2[f]
        for e in range(len(collapsed_event_id)):
            ee = list(collapsed_event_id.keys())[e]
            if (ee == 'observer_actor'):
                old_ev_idx1 = list(event_id.keys()).index("observe, actor")
                old_ev_idx2 = list(event_id.keys()).index("observe, observer")
                tmp_feat2[e] = np.mean([tmp_feat[old_ev_idx1],tmp_feat[old_ev_idx2]], axis=0)
            elif ee == 'follower_leader':
                old_ev_idx1 = list(event_id.keys()).index("imitate, leader")
                old_ev_idx2 = list(event_id.keys()).index("imitate, follower")
                tmp_feat2[e] = np.mean([tmp_feat[old_ev_idx1],tmp_feat[old_ev_idx2]], axis=0)
            else:
                new_ev_idx = list(collapsed_event_id.values())[e]
                old_ev_idx = list(event_id.values()).index(new_ev_idx)
                tmp_feat2[e] = tmp_feat[old_ev_idx]
    
    return [sync_m_labels1, sync_m_labels2], [sync_events1, sync_events2], MFeatures2, shift_info