"""
This module includes pipeline elements and processors for the preprocessing of the n-back task.
"""

import re
import pandas as pd

def extract_events(df, drop_old_cols=False):
  """
  Pipeline element to extract N-back events from a list of Behaverse json events.
  """
  pattern = re.compile("N[Bb]ack\..*")

  def find_head_event(lst):
    evts = [*map(lambda x: x.split('.',1)[1], filter(pattern.match, lst))]
    return evts[0]

  df['nback_event_type'] = df.types.apply(lambda t: find_head_event(t))
  if drop_old_cols: df = df.drop(['types'],axis=1)
  return df