import pandas as pd
import numpy as np
import re
import copy
import nltk
from nltk.tokenize import TreebankWordTokenizer
from nltk.tokenize.treebank import TreebankWordDetokenizer
from nltk.tokenize import line_tokenize, sent_tokenize, word_tokenize
nltk.download('punkt')
from rouge_score import rouge_scorer
import collections

class SC_Elements(object):
 
  def __init__(self):
    """
    """
    super(SC_Elements, self).__init__()

  @staticmethod
  def create_pretraining_RSceneO(dfsc)->pd.DataFrame:
    dfscene = dfsc.copy()
    dfscene = dfscene.groupby('Scene')['Element'].sum().reset_index()
    dfscene['wc'] = dfscene['Element'].apply(lambda x: len(word_tokenize(x)))

    cases = [] 
    for i in range(int(dfsc['wc'].sum()//15)):
      d = {}
      scenes = dfscene.sample(n=np.random.poisson(lam=4))
      if 0 < scenes['wc'].sum() < 802:
        d['input_str'] = scenes.sample(frac=1, random_state=42+i)['Element'].apply(
        lambda x: x + '</s><s>').sum()
        d['input_str'] = '<s>' + d['input_str'] + '</s>'
        label = scenes.sort_values('Scene')['Element'].apply(
          lambda x: x + '</s><s>').sum()
        d['decoder_input_str'] = '<s><s>' + label + '</s>'
        d['labels_str'] = '<s>' + label + '</s></s>'
        cases.append(d)

    return pd.DataFrame(cases).drop_duplicates().copy()

  @staticmethod
  def create_pretraining_GSentG(dfsc,
                                percent_top=0.3,
                                use_f1=False,
                                ) -> pd.DataFrame:
    dfsent = dfsc.copy()
    dfsent['Element_sent'] = dfsent['Element'].apply(lambda x: sent_tokenize(x))
    dfsent = dfsent.explode('Element_sent').reset_index()
    dfsent['Element_sent_word'] = dfsent['Element_sent'].astype('str').apply(word_tokenize)
    # Counter of sentences
    dfsent['counter_sent_word'] = dfsent['Element_sent_word'].apply(collections.Counter)
    dfsent['wc_sent'] = dfsent['counter_sent_word'].apply(lambda x: sum(x.values()))
    # Counter of the entire screenplay
    Counter_SC = dfsent['counter_sent_word'].sum()
    dfsent['counter_rest'] = dfsent['counter_sent_word'].apply(lambda x: Counter_SC - x)
    dfsent['wc_rest'] = dfsent['counter_rest'].apply(lambda x: sum(x.values()))
    # Calculate overlap between sent and the rest
    dfsent['counter_diff'] = dfsent['counter_rest'] - dfsent['counter_sent_word']
    # Filter 0 denominator
    dfsent = dfsent[(dfsent.wc_rest > 0) | dfsent.wc_rest > 0]
    # Calculate Rouge 1
    dfsent['wc_overlap_sent_to_rest'] = dfsent['wc_rest'] - dfsent['counter_diff'].apply(lambda x: sum(x.values()))
    dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest']
    dfsent['recall'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_sent']
    dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall'])

    if use_f1:
      dfsent = dfsent.sort_values('R1-F1', ascending=False)
    else:
      dfsent = dfsent.sort_values('recall', ascending=False)
    n_top_percent = int(dfsent.shape[0]*percent_top)
    idx_topn = dfsent[:n_top_percent].index
    dfsent_sorted = dfsent.sort_index()
    
    idx_all = dfsent.index
    cases = []
    for j in sorted(idx_topn):
      sent_prev, sent_next = '', ''
      d = {}
      if j-1 in idx_all:
        sent_prev = str(dfsent.loc[j-1, 'Element_sent'])
      if j+1 in idx_all:
        sent_next = str(dfsent.loc[j+1, 'Element_sent'])
      d['input_str'] = '<s>'+ '</s><s>'.join([sent_prev, '<mask>', sent_next]) +'</s>'
      d['decoder_input_str'] = '<s><s>'+ '</s><s>'.join(
        [sent_prev, dfsent.loc[j, 'Element_sent'], sent_next]) + '</s>'
      d['labels_str'] = '<s>'+ '</s><s>'.join(
        [sent_prev, dfsent.loc[j, 'Element_sent'], sent_next]) + '</s></s>'
      if d:
        cases.append(d)  
    dftrain = pd.DataFrame(cases)
    return dftrain

  @staticmethod
  def create_pretraining_GSceneG(dfsc,
                                percent_top=0.3,
                                use_f1=False,
                                ) -> pd.DataFrame:

    # GSceneG TOP 30%
    dfscene = dfsc.groupby('Scene')['Element'].sum().reset_index()
    dfscene.columns = ['Scene', 'txt']
    dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
    # Counter of the scene and wc of the scene
    dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
    dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
    # Counter of the entire screeplay text
    Counter_SC = dfscene['counter_scene'].sum()
    # Counter of the rest and wc of the rest
    dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
    dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
    # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
    dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
    # Calculate overlap between scene and the rest
    dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
    dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
    # Calculate Rouge-1 Precision, Recall and F1
    dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
    dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
    dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])
    if use_f1:
      dfscene = dfscene.sort_values('R1-F1', ascending=False)
    else:
      dfscene = dfscene.sort_values('recall', ascending=False)
    # Make Selection by R1-F1 rank with top 30%
    n_top30p = int(dfscene.shape[0] *percent_top)
    idx_top30p = dfscene.iloc[:n_top30p].index

    cases = []
    idx_all = dfscene.index
    for i in idx_top30p:
      d = {}
      prev1, prev2 = '', ''
      next1, next2 = '', ''
      curr = dfscene.loc[i, 'txt']
      if i-1 in idx_all:
        prev1 = dfscene.loc[i-1, 'txt']
      if i-2 in idx_all:
        prev2 = dfscene.loc[i-2, 'txt']
      if i+1 in idx_all:
        nex1 = dfscene.loc[i+1, 'txt']
      if i+2 in idx_all:
        nex1 = dfscene.loc[i+2, 'txt']

      if len(word_tokenize(
          ' '.join([prev2, prev1, curr, next1, next2]))) <  768:
          d['input_str'] = '<s>' + \
          '<s></s>'.join([prev2, prev1, curr, next1, next2]) + '</s>'
          d['decoder_input_str'] =  '<s><s>' + \
            '</s><s>'.join([prev2, prev1, curr, next1, next2]) + '</s>'
          d['labels_str'] = '<s>' + \
            '</s><s>'.join([prev2, prev1, curr, next1, next2]) + '</s></s>'
      elif len(word_tokenize(
          ' '.join([prev1, curr, next1]))) <  768:
          d['input_str'] = '<s>' + \
          '<s></s>'.join([ prev1, curr, next1]) + '</s>'
          d['decoder_input_str'] =  '<s><s>' + \
          '<s></s>'.join([prev1, curr, next1]) + '</s>'
          d['labels_str'] =  '<s>' + \
          '<s></s>'.join([prev1, curr, next1]) + '</s></s>'
      else: 
          d['input_str'] = '<s>' + \
          '<s></s>'.join([ prev1, curr, next1]) + '</s>'
          d['decoder_input_str'] =  '<s><s>' + curr + '</s>'
          d['labels_str'] =  '<s>' + curr + '</s></s>'
      if d:
        cases.append(d)

    return pd.DataFrame(cases)

##########

  @staticmethod
  def generate_repattern_for_topchars(
    dfsc,
    top_n_char=10,
    ):
    """
    returns 
      pat_topchars, 
    """
    # identify top appearance characters by Character
    df_chars_sorted = dfsc[dfsc.Type == 'Character'].groupby(
      'Element')['Element'].count().sort_values(
        ascending=False).rename('Characters').reset_index()
    df_chars_sorted['Element'] = df_chars_sorted['Element'].apply(
      lambda x: re.sub('\n', '', x))
    df_chars_sorted.columns = ['Character', 'count_type']
    df_chars_sorted['Character']= df_chars_sorted['Character'].apply(
      lambda x: re.sub('\(.*\)', '', x).lstrip(' ').rstrip(' ')).apply(
        lambda x: re.sub('\(|\)', '', x))
    df_chars_sorted = df_chars_sorted.groupby('Character').sum().reset_index()
    df_chars_sorted['count_AO'] = 0

    # calcualte appearances by name in Actions
    dfsc_AO = dfsc[dfsc.Type == 'Action'].copy()
    for i, row in df_chars_sorted.iterrows():
      res = re.findall(str(row['Character']),
                      dfsc_AO['Element'].dropna().sum(),
                      flags=re.IGNORECASE
                      )
      df_chars_sorted.loc[i, 'count_AO'] = len(res)

    # combine and re-sort
    df_chars_sorted['count_total'] = \
    df_chars_sorted['count_AO'] + df_chars_sorted['count_AO']
    df_chars_sorted = df_chars_sorted.sort_values(
      'count_total', ascending=False).reset_index(drop=True)

    # create pairs from Characters
    ls_chars = df_chars_sorted['Character'].tolist()[:top_n_char]
    ls_chars += ['he', 'she', 'him', 'her', 'they', 'them']
    pairs = []
    L = len(ls_chars)
    for i in range(L):
      for j in range(i+1, L):
        pairs.append((ls_chars[i], ls_chars[j]))

    # Create pat for re
    ls_pat = []

    pat_topchars_for_word_search = '$|^'.join(ls_chars)
    pat_topchars_for_word_search = '(^' + pat_topchars_for_word_search + '$)'

    for pair in pairs:
      left, right = pair
      ls_pat.append( left + '(?=[\s\.\:\,])' +'.*'+ '(?<=[\s])' + right)
      ls_pat.append(right + '(?=[\s\.\:\,])' + '.*'+ '(?<=[\s])' + left)
      #ls_pat.append(right+'.*'+left)
    pat_pair = '({})'.format('|'.join(ls_pat))

    pat_topchars_for_str_search = '{}'.format('|(?<=[\s\.\,\!\:])'.join(ls_chars))
    pat_topchars_for_str_search = '((?<=[\s\.\,\!\:])' + pat_topchars_for_str_search + ')'
    tmp = '{}'.format('(?=[\s\.\,\!\:])|'.join(ls_chars))
    tmp = '(' + tmp + '(?=[\s\.\,\!\:]))'
    pat_topchars_for_str_search = pat_topchars_for_str_search + '|' + tmp


    return pat_pair, pat_topchars_for_word_search, pat_topchars_for_str_search

  @staticmethod
  def create_pretrainig_RCAA(
    dfsc,
    split_size=1024,
    frac=0.5,
    top_n_char=10,
    shuffle_element=False,
    shuffle_scene=False,
    random_state=42,
    GA_on_scene=True
    )-> pd.DataFrame:
    """
    RECOVER CHARACTER ACROSS ACTION
    (1) Identify top n characters -> List[str]=['Tom', 'Jerry']
    
    """

    # Instantiate Detokenizer
    detokenizer = TreebankWordDetokenizer()
    
    # Get reg patterns
    pat_pair, pat_topchars_for_word_search, pat_topchars_for_str_search = \
      SC_Elements.generate_repattern_for_topchars(dfsc=dfsc, top_n_char=top_n_char)
    
    # Identify elements to mask
    def to_mask(ls_sent) -> bool:
      local_ls_sent = copy.deepcopy(ls_sent)
      num_char_in_sent = []
      for sent in local_ls_sent:
        char_match_in_sent = []
        for word in sent:
          if re.match(pat_topchars_for_word_search, word, flags=re.IGNORECASE):
            char_match_in_sent.append(1)
          else:
            char_match_in_sent.append(0)
        num_char_in_sent.append(sum(char_match_in_sent))
        if max(num_char_in_sent) > 1:
          return True
      return False
    ##
    dfsc_to_mask = dfsc[dfsc['Element_word'].apply(to_mask)].copy()
    ##

  # Apply Mask and Get Gold
    def apply_mask_to_Element_word(ls_ls_word):
      local_ls_ls_word = copy.deepcopy(ls_ls_word)
      num_char_in_sent = []
      gold = []
      for i, sent in enumerate(local_ls_ls_word):
        char_match_in_sent = []
        for j, word in enumerate(sent):
          if re.match(pat_topchars_for_word_search, word, flags=re.IGNORECASE):
            char_match_in_sent.append(1)
          else:
            char_match_in_sent.append(0)
        
        if sum(char_match_in_sent) > 1:
          for j, word in enumerate(sent):
            if re.match(pat_topchars_for_word_search, word, flags=re.IGNORECASE):
              gold.append(word)
              local_ls_ls_word[i][j] = '<mask>'

          return [local_ls_ls_word, gold]
    ##
    dfsc_to_mask['Element_word_masked'] = dfsc_to_mask['Element_word'].apply(
      lambda x: apply_mask_to_Element_word(x)[0])
    dfsc_to_mask['gold'] = dfsc_to_mask['Element_word'].apply(
      lambda x: apply_mask_to_Element_word(x)[1])
    ##

    def word_to_sent(ls_ls_word):
      ls_sent = []
      for sent in ls_ls_word:
        ls_sent.append(detokenizer.detokenize(sent))
      return ls_sent

    def sent_to_element(ls_sent):
      return ' '.join(ls_sent) + '\n'
    ##
    dfsc_to_mask['Element_masked'] = dfsc_to_mask['Element_word_masked'].apply(
        word_to_sent).apply(sent_to_element)
    ##

    idx_w_mask = dfsc_to_mask.index
    dfsc['Element_masked'] = np.NaN
    dfsc.loc[dfsc.index.isin(idx_w_mask), 'Element_masked'] = \
    dfsc_to_mask['Element_masked'].copy()

    dfsc['mgold'] = np.NaN
    dfsc.loc[dfsc.index.isin(idx_w_mask), 'mgold'] = \
    dfsc_to_mask['gold'].copy()

    # Split dfsc with split_size
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for grp in grps:

      if GA_on_scene:
        idx_boScene = grp.drop_duplicates('Scene', keep='first')
        idx_eoScene = grp.drop_duplicates('Scene', keep='last')
        grp.loc[grp.index.isin(idx_boScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_boScene), 'Element'].apply(
            lambda x: '<s>'+str(x))
        grp.loc[grp.index.isin(idx_eoScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_eoScene), 'Element'].apply(
            lambda x: str(x) + '</s>' )
      
      # calculate breakup n if decoder-input exceeds 1024
      str_masked = grp.loc[grp.index.isin(idx_w_mask), 'mgold'].apply(
        lambda x: ' '.join(str(x))).sum()
      tc_masked = len(str(str_masked).split(' '))
      if tc_masked > 1024:
        n = tc_masked//1024 + 1
      else:
        n = 1

      for i in range(n):
        agrp = grp.copy()
        agrp['mgold'] = np.NaN
        idx_agrp = grp['mgold'].dropna().sample(
          frac=1/n, random_state=random_state+i).index
        agrp.loc[agrp.index.isin(idx_agrp), 'mgold'] = \
          grp.loc[grp.index.isin(idx_agrp), 'mgold']

        d = {}
        agrp_fillna = agrp['Element_masked'].fillna(agrp['Element'])

        if shuffle_element:
          d['input_str'] = agrp_fillna.sample(
              frac=1,
              random_state=random_state+i
          ).sum()
        elif shuffle_scene:
          d['input_str'] = agrp.groupby('Scene').sample(
            frac=1,
            random_state=random_state+i
          )['Element_masked'].fillna(
            agrp['Element']).sum()
        else:
          d['input_str'] = agrp_fillna.sum()

        # if Global Attention on Scene
        if GA_on_scene:
          d['input_str'] = '<s>' + d['input_str'] + '</s>'

        if not agrp.loc[agrp.index.isin(idx_agrp), 'mgold'].empty:
          d['decoder_input_str'] = '<s>' + agrp['Element'].sum()
          d['labels_str'] = agrp['Element'].sum() + '</s>'
        else:
          d['decoder_input_str'] = '<s>'
          d['labels_str'] = '</s>'
        
        cases.append(d)
    return pd.DataFrame(cases)  

 
  @staticmethod
  def create_pretrainig_RAAC(
    dfsc,
    split_size=1024,
    frac=0.5,
    top_n_char=10,
    shuffle_element=False,
    shuffle_scene=False,
    random_state=42,
    GA_on_scene=True
    )-> pd.DataFrame:

    # Instantiate Detokenizer
    detokenizer = TreebankWordDetokenizer()
    
    # Get reg patterns
    pat_pair, pat_topchars_for_word_search, pat_topchars_for_str_search = \
      SC_Elements.generate_repattern_for_topchars(dfsc=dfsc, top_n_char=top_n_char)
    
    # Identify elements to mask
    def to_mask(ls_sent) -> bool:
      local_ls_sent = copy.deepcopy(ls_sent)
      num_char_in_sent = []
      for sent in local_ls_sent:
        char_match_in_sent = []
        for word in sent:
          if re.match(pat_topchars_for_word_search, word, flags=re.IGNORECASE):
            char_match_in_sent.append(1)
          else:
            char_match_in_sent.append(0)
        num_char_in_sent.append(sum(char_match_in_sent))
        if max(num_char_in_sent) > 1:
          return True
      return False
    ##
    dfsc_to_mask = dfsc[dfsc['Element_word'].apply(to_mask)].copy()
    ##

  # Apply Mask and Get Gold
    def apply_mask_to_Element_word(ls_ls_word):
      local_ls_ls_word = copy.deepcopy(ls_ls_word)
      num_char_in_sent = []
      gold = []
      ls_sent = []
      for i, sent in enumerate(local_ls_ls_word):
        char_match_in_sent = []
        for j, word in enumerate(sent):
          if re.match(pat_topchars_for_word_search, word, flags=re.IGNORECASE):
            char_match_in_sent.append(1)
          else:
            char_match_in_sent.append(0)
        
        if sum(char_match_in_sent) > 1:
          sent_gold = []
          sent_mask = []
          j = 0
          while j < len(sent):
            if not re.match(pat_topchars_for_word_search, sent[j], flags=re.IGNORECASE):
              sent_gold.append(sent[j])
              if not sent_mask:
                sent_mask.append('<mask>')
              elif sent_mask and sent_mask[-1] != '<mask>':
                sent_mask.append('<mask>')
            else:
              sent_mask.append(sent[j])
            j += 1
          gold.append(' '.join(sent_gold))
          ls_sent.append(' '.join(sent_mask))
        return [' '.join(ls_sent), ' '.join(gold)]
    ##
    dfsc_to_mask['Element_masked'] = dfsc_to_mask['Element_word'].apply(
      lambda x: apply_mask_to_Element_word(x)[0])
    dfsc_to_mask['gold'] = dfsc_to_mask['Element_word'].apply(
      lambda x: apply_mask_to_Element_word(x)[1])

    idx_w_mask = dfsc_to_mask.index
    dfsc['Element_masked'] = np.NaN
    dfsc.loc[dfsc.index.isin(idx_w_mask), 'Element_masked'] = \
     dfsc_to_mask['Element_masked'].copy()

    dfsc['mgold'] = np.NaN
    dfsc.loc[dfsc.index.isin(idx_w_mask), 'mgold'] = \
     dfsc_to_mask['gold'].copy()

    # Split dfsc with split_size
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for grp in grps:

      if GA_on_scene:
        idx_boScene = grp.drop_duplicates('Scene', keep='first')
        idx_eoScene = grp.drop_duplicates('Scene', keep='last')
        grp.loc[grp.index.isin(idx_boScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_boScene), 'Element'].apply(
            lambda x: '<s>'+str(x))
        grp.loc[grp.index.isin(idx_eoScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_eoScene), 'Element'].apply(
            lambda x: str(x) + '</s>' )
      
      # calculate breakup n if decoder-input exceeds 1024
      str_masked = grp.loc[grp.index.isin(idx_w_mask), 'mgold'].apply(
        lambda x: x + ' ').sum()
      tc_masked = len(str(str_masked).split(' '))
      if tc_masked > 1024:
        n = tc_masked//1024 + 1
      else:
        n = 1

      for i in range(n):
        agrp = grp.copy()
        agrp['mgold'] = np.NaN
        idx_agrp = grp['mgold'].dropna().sample(
          frac=1/n, random_state=random_state+i).index
        agrp.loc[agrp.index.isin(idx_agrp), 'mgold'] = \
          grp.loc[grp.index.isin(idx_agrp), 'mgold']
        idx_agrp = sorted(idx_agrp)

        agrp['Element_masked'] = np.NaN
        agrp.loc[agrp.index.isin(idx_agrp), 'Element_masked'] = \
         grp.loc[grp.index.isin(idx_agrp), 'Element_masked']

        d = {}

        agrp_fillna = agrp['Element_masked'].fillna(agrp['Element'])
        if shuffle_element:
          d['input_str'] = agrp_fillna.sample(
              frac=1,
              random_state=random_state+i
          ).sum()
        elif shuffle_scene:
          d['input_str'] = agrp.groupby('Scene').sample(
            frac=1,
            random_state=random_state+i
          )['Element_masked'].fillna(
            agrp['Element']).sum()
        else:
          d['input_str'] = agrp['Element_masked'].fillna(agrp['Element']).sum()

        # if Global Attention on Scene
        if GA_on_scene:
          d['input_str'] = '<s>' + d['input_str'] + '</s>'

        if not agrp.loc[agrp.index.isin(idx_agrp), 'mgold'].empty:
          d['decoder_input_str'] = '<s>' + agrp['Element'].sum()
          d['labels_str'] = agrp['Element'].sum() + '</s>'
        else:
          d['decoder_input_str'] = '<s>'
          d['labels_str'] = '</s>'
        
        cases.append(d)
    return pd.DataFrame(cases)  


  
  @staticmethod
  def create_pretraining_RD(dfsc, 
                            split_size=8192,
                            frac=0.4,
                            shuffle_element=False,
                            shuffle_scene=False,
                            random_state=42,
                            GA_on_scene=True,
                            ) -> pd.DataFrame:
    """
    Recover Dialog
    """
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for grp in grps:
      if GA_on_scene:
        idx_boScene = grp.drop_duplicates('Scene', keep='first')
        idx_eoScene = grp.drop_duplicates('Scene', keep='last')
        grp.loc[grp.index.isin(idx_boScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_boScene), 'Element'].apply(
            lambda x: '<s>'+str(x))
        grp.loc[grp.index.isin(idx_eoScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_eoScene), 'Element'].apply(
            lambda x: str(x) + '</s>' )
      
      for j in range(10):
        local_grp = grp.copy()
        agrp = local_grp[local_grp.Type == 'Dialogue']
        if agrp.empty:
          continue
        agrp = agrp.groupby('Scene').sample(
          frac=frac,
          random_state=random_state+j
        )

        d = {}
        idx_to_mask = agrp.index
        local_grp['Element_masked'] = local_grp['Element'].copy()
        local_grp.loc[local_grp.index.isin(idx_to_mask), 'Element_masked'] = '<mask>\n'
        if shuffle_element:
          d['input_str'] = local_grp['Element_masked'].sample(
            frac=1, 
            random_state=random_state+j).sum()
        elif shuffle_scene:
          d['input_str'] = local_grp.groupby('Scene').sample(
            frac=1,
            random_state=random_state+j
          )['Element_masked'].sum()
        else:
          d['input_str'] = local_grp['Element_masked'].sum()

        if GA_on_scene:
          d['input_str'] = '<s>' + d['input_str'] + '</s>'


        if not local_grp.loc[local_grp.index.isin(idx_to_mask), 'Element'].empty:
         d['decoder_input_str'] = '<s>' + \
          local_grp.loc[local_grp.index.isin(idx_to_mask), 'Element'].apply(
           lambda x: '<s>' + x + '</s>').sum()
         d['labels_str'] = local_grp.loc[local_grp.index.isin(idx_to_mask), 'Element'].apply(
           lambda x: '<s>' + x + '</s>').sum() + '</s>'
        else:
          d['decoder_input_str'] = '<s>'
          d['labels_str'] = '</s>'

        
        cases.append(d)

    return pd.DataFrame(cases)

  @staticmethod
  def create_pretraining_RL(dfsc, 
                            split_size=8192,
                            frac=1,
                            shuffle_element=False,
                            shuffle_scene=False,
                            random_state=42,
                            GA_on_scene=True,
                            ):
  
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for  grp in grps:
      if GA_on_scene:
        idx_boScene = grp.drop_duplicates('Scene', keep='first')
        idx_eoScene = grp.drop_duplicates('Scene', keep='last')
        grp.loc[grp.index.isin(idx_boScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_boScene), 'Element'].apply(
            lambda x: '<s>'+str(x))
        grp.loc[grp.index.isin(idx_eoScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_eoScene), 'Element'].apply(
            lambda x: str(x) + '</s>' )
      frac_local = frac
      repeat = 0
      i = 0
      sh_to_mask = grp.loc[grp.Type == 'Scene Heading', :]
      while sh_to_mask.shape[0] > 320:
        frac_local -= 0.1
        sh_to_mask = grp.loc[grp.Type == 'Scene Heading', :].sample(
        frac=frac_local, 
        random_state=random_state+i
        )
        i += 1
      
      while i > -1:
        d = {}
        idx_to_mask = sh_to_mask.index
        grp['Element_masked'] = grp['Element'].copy()
        grp.loc[grp.index.isin(idx_to_mask), 'Element_masked'] = '<mask>\n'
        if shuffle_element:
          d['input_str'] = grp['Element_masked'].sample(
            frac=1, 
            random_state=random_state).sum()
        elif shuffle_scene:
          d['input_str'] = grp.groupby('Scene').sample(
            frac=1,
            random_state=random_state
          )['Element_masked'].sum()
        else:
          d['input_str'] = grp['Element_masked'].sum()

        if GA_on_scene:
          d['input_str'] = '<s>' + d['input_str'] + '</s>'

        #remove '\n'
        str_masked_dchars = str(grp.loc[grp.index.isin(idx_to_mask), 'Element'].apply(
          lambda x: re.sub('\n', ' ', x)).sum())

        d['decoder_input_str'] = '<s>' + str_masked_dchars
        d['labels_str'] = str_masked_dchars + '</s>'
        cases.append(d)
        sh_to_mask = grp.loc[grp.Type == 'Scene Heading', :].sample(
          frac=frac_local, 
          random_state=random_state+i
          )
        i -= 1

    return pd.DataFrame(cases)

  @staticmethod
  def create_pretraining_RCID(dfsc, 
                              split_size=1024,
                              frac=1,
                              shuffle_element=False,
                              shuffle_scene=False,
                              random_state=42,
                              GA_on_scene=True,
                              ):
   
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for  grp in grps:
      if GA_on_scene:
        idx_boScene = grp.drop_duplicates('Scene', keep='first')
        idx_eoScene = grp.drop_duplicates('Scene', keep='last')
        grp.loc[grp.index.isin(idx_boScene), 'Element'] = \
         grp.loc[grp.index.isin(idx_boScene), 'Element'].apply(
           lambda x: '<s>'+str(x))
        grp.loc[grp.index.isin(idx_eoScene), 'Element'] = \
         grp.loc[grp.index.isin(idx_eoScene), 'Element'].apply(
           lambda x: str(x) + '</s>' )
      frac_local = frac
      repeat = 0
      i = 0
      dchar_to_mask = grp.loc[grp.Type == 'Character', :]
      while dchar_to_mask.shape[0] > 320:
        frac_local -= 0.1
        dchar_to_mask = grp.loc[grp.Type == 'Character', :].sample(
        frac=frac_local, 
        random_state=random_state+i
        )
        i += 1
      
      while i > -1:
        d = {}
        idx_to_mask = dchar_to_mask.index
        grp['Element_masked'] = grp['Element'].copy()
        grp.loc[grp.index.isin(idx_to_mask), 'Element_masked'] = '<mask>\n'
        if shuffle_element:
          d['input_str'] = grp['Element_masked'].sample(
            frac=1, 
            random_state=random_state).sum()
        elif shuffle_scene:
          d['input_str'] = grp.groupby('Scene').sample(
            frac=1,
            random_state=random_state
          )['Element_masked'].sum()
        else:
          d['input_str'] = grp['Element_masked'].sum()

        if GA_on_scene:
          d['input_str'] = '<s>' + d['input_str'] + '</s>'

        #remove '\n'
        str_masked_dchars = str(grp['Element'].apply(
          lambda x: re.sub('\n', ' ', x)).sum())

        d['decoder_input_str'] = '<s>' + str_masked_dchars
        d['labels_str'] = str_masked_dchars + '</s>'
        cases.append(d)
        dchar_to_mask = grp.loc[grp.Type == 'Character', :].sample(
          frac=frac_local, 
          random_state=random_state+i
          )
        i -= 1

    return pd.DataFrame(cases)

  @staticmethod
  def create_pretraining_TI(
                            dfsc, 
                            split_size=4096,
                            frac=0.2,
                            shuffle_element=False,
                            random_state=42,
                            ):
    """This funtion create pretraining cases with optional element shuffle
    dfsc: pd.DataFrame
    split_size=4096
    frac=0.2,
    shuffle_element=False,
    random_state=42
    """
    
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for  grp in grps:
      for j in range(10):
        agrp = grp.groupby('Scene').sample(frac=frac, random_state=j).copy()
        if (grp['tc'].sum() < int(split_size*1.5) and 
            400 <= agrp['tc'].sum() < 960):
          d = {}
          idx_to_mask = agrp.index
          grp['Element_masked'] = grp['Element'].copy()
          grp.loc[grp.index.isin(idx_to_mask), 'Element_masked'] = '<mask>'
          if shuffle_element:
            d['input_str'] = grp['Element_masked'].sample(
              frac=1, 
              random_state=random_state).sum()
          else:
            d['input_str'] = grp['Element_masked'].sum()
          d['decoder_input_str'] = '<s>' + str(grp.loc[grp.index.isin(idx_to_mask), 'Element'].sum())
          d['labels_str'] = str(grp.loc[grp.index.isin(idx_to_mask), 'Element'].sum()) + '</s>'
          cases.append(d)

    return pd.DataFrame(cases)
  
  @staticmethod
  def extract_str_by_method(dfsc, method='AF', return_type='str', n_main=1):
    """
    method = 'AF', 'AO', 'DO', 'FMC', 'FMCAO', 
    'EVERY2at0', 'EVERY2at1',
    'EVERY3at0', 'EVERY3at1', 'EVERY3at2',
    'EVERY4at0', 'EVERY4at1', 'EVERY4at2', 'EVERY4at3'
    F: All Flat
    AO: Action Only
    DO: Dialog Only
    FMC: Follow Main Character
    FMCAO: Follow Main Character + Action Only

    return_type= 'str', 'df', 'df_full'
    """
    dfsc['Element'] = dfsc['Element'].apply(lambda x: str(x).lstrip('\w+').rstrip('\w+').rstrip('\n') + '\n')
    dfsc['Element'] = dfsc['Element'].apply(lambda x: re.sub(':SC:', '', str(x)) )
    dfsc['Scene'] = dfsc['Scene'].fillna(-1)

    if not method:
      return dfsc

    ########
    if method == 'PSentAO':
     dfsc_w_method = dfsc.loc[dfsc['Type'] == 'Action', :].copy()
     dfsent = dfsc_w_method.copy()
     dfsent['Element_sent'] = dfsent['Element'].apply(lambda x: sent_tokenize(x))
     dfsent = dfsent.explode('Element_sent').reset_index()
     dfsent['Element_sent_word'] = dfsent['Element_sent'].astype('str').apply(word_tokenize)
     # Counter of sentences
     dfsent['counter_sent_word'] = dfsent['Element_sent_word'].apply(collections.Counter)
     dfsent['wc_sent'] = dfsent['counter_sent_word'].apply(lambda x: sum(x.values()))
     # Counter of the entire screenplay
     Counter_SC = dfsent['counter_sent_word'].sum()
     dfsent['counter_rest'] = dfsent['counter_sent_word'].apply(lambda x: Counter_SC - x)
     dfsent['wc_rest'] = dfsent['counter_rest'].apply(lambda x: sum(x.values()))
     # Calculate overlap between sent and the rest
     dfsent['counter_diff'] = dfsent['counter_rest'] - dfsent['counter_sent_word']
     # Filter 0 denominator
     dfsent = dfsent[(dfsent.wc_rest > 0) | dfsent.wc_rest > 0]
     # Calculate Rouge 1
     dfsent['wc_overlap_sent_to_rest'] = dfsent['wc_rest'] - dfsent['counter_diff'].apply(lambda x: sum(x.values()))
     #dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest']
     dfsent['recall'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_sent']

     #dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall'])
     dfsent = dfsent.sort_values('recall', ascending=False)
     dfsent['wc_cum'] = dfsent['wc_sent'].cumsum()
     df_selection = dfsent[dfsent['wc_cum'] < 12000].sort_index()
     df_selection['txt'] = df_selection['Element_sent'].fillna('').astype('str')

     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()
    ########
    if method == 'PSentAO_F1':
     dfsc_w_method = dfsc.loc[dfsc['Type'] == 'Action', :].copy()
     dfsent = dfsc_w_method.copy()
     dfsent['Element_sent'] = dfsent['Element'].apply(lambda x: sent_tokenize(x))
     dfsent = dfsent.explode('Element_sent').reset_index()
     dfsent['Element_sent_word'] = dfsent['Element_sent'].astype('str').apply(word_tokenize)
     # Counter of sentences
     dfsent['counter_sent_word'] = dfsent['Element_sent_word'].apply(collections.Counter)
     dfsent['wc_sent'] = dfsent['counter_sent_word'].apply(lambda x: sum(x.values()))
     # Counter of the entire screenplay
     Counter_SC = dfsent['counter_sent_word'].sum()
     dfsent['counter_rest'] = dfsent['counter_sent_word'].apply(lambda x: Counter_SC - x)
     dfsent['wc_rest'] = dfsent['counter_rest'].apply(lambda x: sum(x.values()))
     # Calculate overlap between sent and the rest
     dfsent['counter_diff'] = dfsent['counter_rest'] - dfsent['counter_sent_word']
     # Filter 0 denominator
     dfsent = dfsent[(dfsent.wc_rest > 0) | dfsent.wc_rest > 0]
     # Calculate Rouge 1
     dfsent['wc_overlap_sent_to_rest'] = dfsent['wc_rest'] - dfsent['counter_diff'].apply(lambda x: sum(x.values()))
     dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest']
     dfsent['recall'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_sent']

     dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall'])
     dfsent = dfsent.sort_values('R1-F1', ascending=False)
     dfsent['wc_cum'] = dfsent['wc_sent'].cumsum()
     df_selection = dfsent[dfsent['wc_cum'] < 12000].sort_index()
     df_selection['txt'] = df_selection['Element_sent'].fillna('').astype('str')

     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()

    ########  
    if method == 'PSent_F1':
     dfsc_w_method = dfsc.copy()
     dfsent = dfsc_w_method.copy()
     dfsent['Element_sent'] = dfsent['Element'].apply(lambda x: sent_tokenize(x))
     dfsent = dfsent.explode('Element_sent').reset_index()
     dfsent['Element_sent_word'] = dfsent['Element_sent'].astype('str').apply(word_tokenize)
     # Counter of sentences
     dfsent['counter_sent_word'] = dfsent['Element_sent_word'].apply(collections.Counter)
     dfsent['wc_sent'] = dfsent['counter_sent_word'].apply(lambda x: sum(x.values()))
     # Counter of the entire screenplay
     Counter_SC = dfsent['counter_sent_word'].sum()
     dfsent['counter_rest'] = dfsent['counter_sent_word'].apply(lambda x: Counter_SC - x)
     dfsent['wc_rest'] = dfsent['counter_rest'].apply(lambda x: sum(x.values()))
     # Calculate overlap between sent and the rest
     dfsent['counter_diff'] = dfsent['counter_rest'] - dfsent['counter_sent_word']
     # Filter 0 denominator
     dfsent = dfsent[(dfsent.wc_rest > 0) | dfsent.wc_rest > 0]
     # Calculate Rouge 1
     dfsent['wc_overlap_sent_to_rest'] = dfsent['wc_rest'] - dfsent['counter_diff'].apply(lambda x: sum(x.values()))
     dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest']
     dfsent['recall'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_sent']

     dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall'])
     dfsent = dfsent.sort_values('R1-F1', ascending=False)
     dfsent['wc_cum'] = dfsent['wc_sent'].cumsum()
     df_selection = dfsent[dfsent['wc_cum'] < 12000].sort_index()
     df_selection['txt'] = df_selection['Element_sent'].fillna('').astype('str')

     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()
    ########  
    if method == 'PSent':
     dfsc_w_method = dfsc.copy()
     dfsent = dfsc_w_method.copy()
     dfsent['Element_sent'] = dfsent['Element'].apply(lambda x: sent_tokenize(x))
     dfsent = dfsent.explode('Element_sent').reset_index()
     dfsent['Element_sent_word'] = dfsent['Element_sent'].astype('str').apply(word_tokenize)
     # Counter of sentences
     dfsent['counter_sent_word'] = dfsent['Element_sent_word'].apply(collections.Counter)
     dfsent['wc_sent'] = dfsent['counter_sent_word'].apply(lambda x: sum(x.values()))
     # Counter of the entire screenplay
     Counter_SC = dfsent['counter_sent_word'].sum()
     dfsent['counter_rest'] = dfsent['counter_sent_word'].apply(lambda x: Counter_SC - x)
     dfsent['wc_rest'] = dfsent['counter_rest'].apply(lambda x: sum(x.values()))
     # Calculate overlap between sent and the rest
     dfsent['counter_diff'] = dfsent['counter_rest'] - dfsent['counter_sent_word']
     # Filter 0 denominator
     dfsent = dfsent[(dfsent.wc_rest > 0) | dfsent.wc_rest > 0]
     # Calculate Rouge 1
     dfsent['wc_overlap_sent_to_rest'] = dfsent['wc_rest'] - dfsent['counter_diff'].apply(lambda x: sum(x.values()))
     #dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest']
     dfsent['recall'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_sent']

     #dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall'])
     dfsent = dfsent.sort_values('recall', ascending=False)
     dfsent['wc_cum'] = dfsent['wc_sent'].cumsum()
     df_selection = dfsent[dfsent['wc_cum'] < 12000].sort_index()
     df_selection['txt'] = df_selection['Element_sent'].fillna('').astype('str')

     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()
    ########  
    if method == 'PSceneH5T5AO':      
     dfsc_w_method = dfsc.loc[dfsc['Type'] == 'Action', :].copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     #dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     #dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])  
     
     # split dataframe to first5, center and last5
     idx_first5 = dfscene.index[:5]
     idx_last5 = dfscene.index[-5:]
     idx_center = dfscene.index[5:-5]
     first5_wc = dfscene.loc[idx_first5, 'wc_scene'].sum()
     last5_wc = dfscene.loc[idx_last5, 'wc_scene'].sum()
     # Sort center df by R1-F1
     df_center = dfscene.loc[idx_center, :].sort_values(
      'recall', ascending=False)
     df_center['cum_wc']  = df_center['wc_scene'].cumsum()
     idx_ros = df_center[df_center.cum_wc < (12000-first5_wc-last5_wc)].index
     df_selection= dfscene.loc[idx_first5.union(idx_ros).union(idx_last5),:]
     df_selection = df_selection.sort_values('Scene')

     if return_type == 'str':
       ret = df_selection['txt'].sum()
     if return_type == 'df':
       ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
       ret = df_selection.copy()
    ########  
    if method == 'PSceneH5T5AO_F1':      
     dfsc_w_method = dfsc.loc[dfsc['Type'] == 'Action', :].copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])  
     
     # split dataframe to first5, center and last5
     idx_first5 = dfscene.index[:5]
     idx_last5 = dfscene.index[-5:]
     idx_center = dfscene.index[5:-5]
     first5_wc = dfscene.loc[idx_first5, 'wc_scene'].sum()
     last5_wc = dfscene.loc[idx_last5, 'wc_scene'].sum()
     # Sort center df by R1-F1
     df_center = dfscene.loc[idx_center, :].sort_values(
      'R1-F1', ascending=False)
     df_center['cum_wc']  = df_center['wc_scene'].cumsum()
     idx_ros = df_center[df_center.cum_wc < (12000-first5_wc-last5_wc)].index
     df_selection= dfscene.loc[idx_first5.union(idx_ros).union(idx_last5),:]
     df_selection = df_selection.sort_values('Scene')

     if return_type == 'str':
       ret = df_selection['txt'].sum()
     if return_type == 'df':
       ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
       ret = df_selection.copy()
    ##########################
    if method == 'PSceneH5T5':
      
     dfsc_w_method = dfsc.copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     #dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     #dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])  
     
     # split dataframe to first5, center and last5
     idx_first5 = dfscene.index[:5]
     idx_last5 = dfscene.index[-5:]
     idx_center = dfscene.index[5:-5]
     first5_wc = dfscene.loc[idx_first5, 'wc_scene'].sum()
     last5_wc = dfscene.loc[idx_last5, 'wc_scene'].sum()
     # Sort center df by R1-F1
     df_center = dfscene.loc[idx_center, :].sort_values('recall', ascending=False)
     df_center['cum_wc']  = df_center['wc_scene'].cumsum()
     idx_ros = df_center[df_center.cum_wc < (12000-first5_wc-last5_wc)].index
     df_selection= dfscene.loc[idx_first5.union(idx_ros).union(idx_last5),:]
     df_selection = df_selection.sort_values('Scene')

     if return_type == 'str':
       ret = df_selection['txt'].sum()
     if return_type == 'df':
       ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
       ret = df_selection.copy()

    ##########################
    if method == 'PSceneH5T5_F1':
      
     dfsc_w_method = dfsc.copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])  
     
     # split dataframe to first5, center and last5
     idx_first5 = dfscene.index[:5]
     idx_last5 = dfscene.index[-5:]
     idx_center = dfscene.index[5:-5]
     first5_wc = dfscene.loc[idx_first5, 'wc_scene'].sum()
     last5_wc = dfscene.loc[idx_last5, 'wc_scene'].sum()
     # Sort center df by R1-F1
     df_center = dfscene.loc[idx_center, :].sort_values('R1-F1', ascending=False)
     df_center['cum_wc']  = df_center['wc_scene'].cumsum()
     idx_ros = df_center[df_center.cum_wc < (12000-first5_wc-last5_wc)].index
     df_selection= dfscene.loc[idx_first5.union(idx_ros).union(idx_last5),:]
     df_selection = df_selection.sort_values('Scene')

     if return_type == 'str':
       ret = df_selection['txt'].sum()
     if return_type == 'df':
       ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
       ret = df_selection.copy()
    ##########################
    if method == 'PScene':

     dfsc_w_method = dfsc.copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     #dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     #dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])
     dfscene = dfscene.sort_values('recall', ascending=False)
     dfscene['cum_wc'] = dfscene['wc_scene'].cumsum()
     # Make Selection by R1-F1 rank
     df_selection = dfscene[dfscene.cum_wc < 12000].sort_values(
     'Scene', ascending=True)
     
     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()
    ##########################
    if method == 'PScene_F1':

     dfsc_w_method = dfsc.copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])
     dfscene = dfscene.sort_values('R1-F1', ascending=False)
     dfscene['cum_wc'] = dfscene['wc_scene'].cumsum()
     # Make Selection by R1-F1 rank
     df_selection = dfscene[dfscene.cum_wc < 12000].sort_values(
     'Scene', ascending=True)
     
     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()
    ##########################
    if method == 'EVERY2at0':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 2 == 1, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY2at1':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 2 == 0, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY3at0':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 3 == 0, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY3at1':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 3 == 1, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY3at2':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 3 == 2, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY4at0':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 4 == 0,: ].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY4at1':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 4 == 1, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY4at2':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 4 == 2, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY4at3':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 4 == 3, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'AF':
      dfsc_w_method = dfsc.copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'AO':
      dfsc_w_method = dfsc.loc[dfsc['Type'] == 'Action', :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'DO':
      #dfsc.loc[dfsc['Type'].str.contains('Character'), 'Element'] = \
      #  dfsc.loc[dfsc['Type'].str.contains('Character'), 'Element'].apply(lambda x: x.rstrip('\n') + ':')
      dfsc_w_method = dfsc.loc[dfsc['Type'].str.contains('Character|Dialogue'), :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'FMC':
      dfsc['Element'] = dfsc['Element'].astype('str')
      df_chars = dfsc.loc[dfsc['Type']=='Character', ['Element', 'Type']].groupby(
        'Element').count().reset_index()
      df_chars.columns = ['Char', 'Count']
      df_chars = df_chars.sort_values(by='Count', ascending=False)
      ls_Scene = dfsc.loc[dfsc['Element'].isin(df_chars['Char'][:n_main]), 'Scene'].drop_duplicates().tolist()
      dfsc_w_method = dfsc.loc[dfsc['Scene'].isin(ls_Scene), :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()     
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'FMCAO':
      dfsc['Element'] = dfsc['Element'].astype('str')
      df_chars = dfsc.loc[dfsc['Type']=='Character', ['Element', 'Type']].groupby(
        'Element').count().reset_index()
      df_chars.columns = ['Char', 'Count']
      df_chars = df_chars.sort_values(by='Count', ascending=False)
      ls_Scene = dfsc.loc[dfsc['Element'].isin(df_chars['Char'][:n_main]), 'Scene'].drop_duplicates().tolist()
      dfsc_w_method = dfsc.loc[(dfsc['Scene'].isin(ls_Scene)) & (dfsc['Type'] == 'Action'), :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'TPS':
      dfsc_w_method = dfsc.loc[dfsc['tps'] > 0, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()      
      if return_type =='df_full':
        ret = dfsc_w_method

    if return_type == 'str':
      return ret
    if return_type == 'df':
      return ret.copy()
    if return_type == 'df_full':
      return ret.copy()
import pandas as pd
import numpy as np
import re
import copy
import nltk
from nltk.tokenize import TreebankWordTokenizer
from nltk.tokenize.treebank import TreebankWordDetokenizer
from nltk.tokenize import line_tokenize, sent_tokenize, word_tokenize
nltk.download('punkt')
from rouge_score import rouge_scorer
import collections

class SC_Elements(object):
 
  def __init__(self):
    """
    """
    super(SC_Elements, self).__init__()

  @staticmethod
  def create_pretraining_RSceneO(dfsc)->pd.DataFrame:
    dfscene = dfsc.copy()
    dfscene = dfscene.groupby('Scene')['Element'].sum().reset_index()
    dfscene['wc'] = dfscene['Element'].apply(lambda x: len(word_tokenize(x)))

    cases = [] 
    for i in range(int(dfsc['wc'].sum()//15)):
      d = {}
      scenes = dfscene.sample(n=np.random.poisson(lam=4))
      if 0 < scenes['wc'].sum() < 802:
        d['input_str'] = scenes.sample(frac=1, random_state=42+i)['Element'].apply(
        lambda x: x + '</s><s>').sum()
        d['input_str'] = '<s>' + d['input_str'] + '</s>'
        label = scenes.sort_values('Scene')['Element'].apply(
          lambda x: x + '</s><s>').sum()
        d['decoder_input_str'] = '<s><s>' + label + '</s>'
        d['labels_str'] = '<s>' + label + '</s></s>'
        cases.append(d)

    return pd.DataFrame(cases).drop_duplicates().copy()

  @staticmethod
  def create_pretraining_GSentG(dfsc,
                                percent_top=0.3,
                                use_f1=False,
                                ) -> pd.DataFrame:
    dfsent = dfsc.copy()
    dfsent['Element_sent'] = dfsent['Element'].apply(lambda x: sent_tokenize(x))
    dfsent = dfsent.explode('Element_sent').reset_index()
    dfsent['Element_sent_word'] = dfsent['Element_sent'].astype('str').apply(word_tokenize)
    # Counter of sentences
    dfsent['counter_sent_word'] = dfsent['Element_sent_word'].apply(collections.Counter)
    dfsent['wc_sent'] = dfsent['counter_sent_word'].apply(lambda x: sum(x.values()))
    # Counter of the entire screenplay
    Counter_SC = dfsent['counter_sent_word'].sum()
    dfsent['counter_rest'] = dfsent['counter_sent_word'].apply(lambda x: Counter_SC - x)
    dfsent['wc_rest'] = dfsent['counter_rest'].apply(lambda x: sum(x.values()))
    # Calculate overlap between sent and the rest
    dfsent['counter_diff'] = dfsent['counter_rest'] - dfsent['counter_sent_word']
    # Filter 0 denominator
    dfsent = dfsent[(dfsent.wc_rest > 0) | dfsent.wc_rest > 0]
    # Calculate Rouge 1
    dfsent['wc_overlap_sent_to_rest'] = dfsent['wc_rest'] - dfsent['counter_diff'].apply(lambda x: sum(x.values()))
    dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest']
    dfsent['recall'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_sent']
    dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall'])

    if use_f1:
      dfsent = dfsent.sort_values('R1-F1', ascending=False)
    else:
      dfsent = dfsent.sort_values('recall', ascending=False)
    n_top_percent = int(dfsent.shape[0]*percent_top)
    idx_topn = dfsent[:n_top_percent].index
    dfsent_sorted = dfsent.sort_index()
    
    idx_all = dfsent.index
    cases = []
    for j in sorted(idx_topn):
      sent_prev, sent_next = '', ''
      d = {}
      if j-1 in idx_all:
        sent_prev = str(dfsent.loc[j-1, 'Element_sent'])
      if j+1 in idx_all:
        sent_next = str(dfsent.loc[j+1, 'Element_sent'])
      d['input_str'] = '<s>'+ '</s><s>'.join([sent_prev, '<mask>', sent_next]) +'</s>'
      d['decoder_input_str'] = '<s><s>'+ '</s><s>'.join(
        [sent_prev, dfsent.loc[j, 'Element_sent'], sent_next]) + '</s>'
      d['labels_str'] = '<s>'+ '</s><s>'.join(
        [sent_prev, dfsent.loc[j, 'Element_sent'], sent_next]) + '</s></s>'
      if d:
        cases.append(d)  
    dftrain = pd.DataFrame(cases)
    return dftrain

  @staticmethod
  def create_pretraining_GSceneG(dfsc,
                                percent_top=0.3,
                                use_f1=False,
                                ) -> pd.DataFrame:

    # GSceneG TOP 30%
    dfscene = dfsc.groupby('Scene')['Element'].sum().reset_index()
    dfscene.columns = ['Scene', 'txt']
    dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
    # Counter of the scene and wc of the scene
    dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
    dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
    # Counter of the entire screeplay text
    Counter_SC = dfscene['counter_scene'].sum()
    # Counter of the rest and wc of the rest
    dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
    dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
    # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
    dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
    # Calculate overlap between scene and the rest
    dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
    dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
    # Calculate Rouge-1 Precision, Recall and F1
    dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
    dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
    dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])
    if use_f1:
      dfscene = dfscene.sort_values('R1-F1', ascending=False)
    else:
      dfscene = dfscene.sort_values('recall', ascending=False)
    # Make Selection by R1-F1 rank with top 30%
    n_top30p = int(dfscene.shape[0] *percent_top)
    idx_top30p = dfscene.iloc[:n_top30p].index

    cases = []
    idx_all = dfscene.index
    for i in idx_top30p:
      d = {}
      prev1, prev2 = '', ''
      next1, next2 = '', ''
      curr = dfscene.loc[i, 'txt']
      if i-1 in idx_all:
        prev1 = dfscene.loc[i-1, 'txt']
      if i-2 in idx_all:
        prev2 = dfscene.loc[i-2, 'txt']
      if i+1 in idx_all:
        nex1 = dfscene.loc[i+1, 'txt']
      if i+2 in idx_all:
        nex1 = dfscene.loc[i+2, 'txt']

      if len(word_tokenize(
          ' '.join([prev2, prev1, curr, next1, next2]))) <  768:
          d['input_str'] = '<s>' + \
          '<s></s>'.join([prev2, prev1, curr, next1, next2]) + '</s>'
          d['decoder_input_str'] =  '<s><s>' + \
            '</s><s>'.join([prev2, prev1, curr, next1, next2]) + '</s>'
          d['labels_str'] = '<s>' + \
            '</s><s>'.join([prev2, prev1, curr, next1, next2]) + '</s></s>'
      elif len(word_tokenize(
          ' '.join([prev1, curr, next1]))) <  768:
          d['input_str'] = '<s>' + \
          '<s></s>'.join([ prev1, curr, next1]) + '</s>'
          d['decoder_input_str'] =  '<s><s>' + \
          '<s></s>'.join([prev1, curr, next1]) + '</s>'
          d['labels_str'] =  '<s>' + \
          '<s></s>'.join([prev1, curr, next1]) + '</s></s>'
      else: 
          d['input_str'] = '<s>' + \
          '<s></s>'.join([ prev1, curr, next1]) + '</s>'
          d['decoder_input_str'] =  '<s><s>' + curr + '</s>'
          d['labels_str'] =  '<s>' + curr + '</s></s>'
      if d:
        cases.append(d)

    return pd.DataFrame(cases)

##########

  @staticmethod
  def generate_repattern_for_topchars(
    dfsc,
    top_n_char=10,
    ):
    """
    returns 
      pat_topchars, 
    """
    # identify top appearance characters by Character
    df_chars_sorted = dfsc[dfsc.Type == 'Character'].groupby(
      'Element')['Element'].count().sort_values(
        ascending=False).rename('Characters').reset_index()
    df_chars_sorted['Element'] = df_chars_sorted['Element'].apply(
      lambda x: re.sub('\n', '', x))
    df_chars_sorted.columns = ['Character', 'count_type']
    df_chars_sorted['Character']= df_chars_sorted['Character'].apply(
      lambda x: re.sub('\(.*\)', '', x).lstrip(' ').rstrip(' ')).apply(
        lambda x: re.sub('\(|\)', '', x))
    df_chars_sorted = df_chars_sorted.groupby('Character').sum().reset_index()
    df_chars_sorted['count_AO'] = 0

    # calcualte appearances by name in Actions
    dfsc_AO = dfsc[dfsc.Type == 'Action'].copy()
    for i, row in df_chars_sorted.iterrows():
      res = re.findall(str(row['Character']),
                      dfsc_AO['Element'].dropna().sum(),
                      flags=re.IGNORECASE
                      )
      df_chars_sorted.loc[i, 'count_AO'] = len(res)

    # combine and re-sort
    df_chars_sorted['count_total'] = \
    df_chars_sorted['count_AO'] + df_chars_sorted['count_AO']
    df_chars_sorted = df_chars_sorted.sort_values(
      'count_total', ascending=False).reset_index(drop=True)

    # create pairs from Characters
    ls_chars = df_chars_sorted['Character'].tolist()[:top_n_char]
    ls_chars += ['he', 'she', 'him', 'her', 'they', 'them']
    pairs = []
    L = len(ls_chars)
    for i in range(L):
      for j in range(i+1, L):
        pairs.append((ls_chars[i], ls_chars[j]))

    # Create pat for re
    ls_pat = []

    pat_topchars_for_word_search = '$|^'.join(ls_chars)
    pat_topchars_for_word_search = '(^' + pat_topchars_for_word_search + '$)'

    for pair in pairs:
      left, right = pair
      ls_pat.append( left + '(?=[\s\.\:\,])' +'.*'+ '(?<=[\s])' + right)
      ls_pat.append(right + '(?=[\s\.\:\,])' + '.*'+ '(?<=[\s])' + left)
      #ls_pat.append(right+'.*'+left)
    pat_pair = '({})'.format('|'.join(ls_pat))

    pat_topchars_for_str_search = '{}'.format('|(?<=[\s\.\,\!\:])'.join(ls_chars))
    pat_topchars_for_str_search = '((?<=[\s\.\,\!\:])' + pat_topchars_for_str_search + ')'
    tmp = '{}'.format('(?=[\s\.\,\!\:])|'.join(ls_chars))
    tmp = '(' + tmp + '(?=[\s\.\,\!\:]))'
    pat_topchars_for_str_search = pat_topchars_for_str_search + '|' + tmp


    return pat_pair, pat_topchars_for_word_search, pat_topchars_for_str_search

  @staticmethod
  def create_pretrainig_RCAA(
    dfsc,
    split_size=1024,
    frac=0.5,
    top_n_char=10,
    shuffle_element=False,
    shuffle_scene=False,
    random_state=42,
    GA_on_scene=True
    )-> pd.DataFrame:
    """
    RECOVER CHARACTER ACROSS ACTION
    (1) Identify top n characters -> List[str]=['Tom', 'Jerry']
    
    """

    # Instantiate Detokenizer
    detokenizer = TreebankWordDetokenizer()
    
    # Get reg patterns
    pat_pair, pat_topchars_for_word_search, pat_topchars_for_str_search = \
      SC_Elements.generate_repattern_for_topchars(dfsc=dfsc, top_n_char=top_n_char)
    
    # Identify elements to mask
    def to_mask(ls_sent) -> bool:
      local_ls_sent = copy.deepcopy(ls_sent)
      num_char_in_sent = []
      for sent in local_ls_sent:
        char_match_in_sent = []
        for word in sent:
          if re.match(pat_topchars_for_word_search, word, flags=re.IGNORECASE):
            char_match_in_sent.append(1)
          else:
            char_match_in_sent.append(0)
        num_char_in_sent.append(sum(char_match_in_sent))
        if max(num_char_in_sent) > 1:
          return True
      return False
    ##
    dfsc_to_mask = dfsc[dfsc['Element_word'].apply(to_mask)].copy()
    ##

  # Apply Mask and Get Gold
    def apply_mask_to_Element_word(ls_ls_word):
      local_ls_ls_word = copy.deepcopy(ls_ls_word)
      num_char_in_sent = []
      gold = []
      for i, sent in enumerate(local_ls_ls_word):
        char_match_in_sent = []
        for j, word in enumerate(sent):
          if re.match(pat_topchars_for_word_search, word, flags=re.IGNORECASE):
            char_match_in_sent.append(1)
          else:
            char_match_in_sent.append(0)
        
        if sum(char_match_in_sent) > 1:
          for j, word in enumerate(sent):
            if re.match(pat_topchars_for_word_search, word, flags=re.IGNORECASE):
              gold.append(word)
              local_ls_ls_word[i][j] = '<mask>'

          return [local_ls_ls_word, gold]
    ##
    dfsc_to_mask['Element_word_masked'] = dfsc_to_mask['Element_word'].apply(
      lambda x: apply_mask_to_Element_word(x)[0])
    dfsc_to_mask['gold'] = dfsc_to_mask['Element_word'].apply(
      lambda x: apply_mask_to_Element_word(x)[1])
    ##

    def word_to_sent(ls_ls_word):
      ls_sent = []
      for sent in ls_ls_word:
        ls_sent.append(detokenizer.detokenize(sent))
      return ls_sent

    def sent_to_element(ls_sent):
      return ' '.join(ls_sent) + '\n'
    ##
    dfsc_to_mask['Element_masked'] = dfsc_to_mask['Element_word_masked'].apply(
        word_to_sent).apply(sent_to_element)
    ##

    idx_w_mask = dfsc_to_mask.index
    dfsc['Element_masked'] = np.NaN
    dfsc.loc[dfsc.index.isin(idx_w_mask), 'Element_masked'] = \
    dfsc_to_mask['Element_masked'].copy()

    dfsc['mgold'] = np.NaN
    dfsc.loc[dfsc.index.isin(idx_w_mask), 'mgold'] = \
    dfsc_to_mask['gold'].copy()

    # Split dfsc with split_size
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for grp in grps:

      if GA_on_scene:
        idx_boScene = grp.drop_duplicates('Scene', keep='first')
        idx_eoScene = grp.drop_duplicates('Scene', keep='last')
        grp.loc[grp.index.isin(idx_boScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_boScene), 'Element'].apply(
            lambda x: '<s>'+str(x))
        grp.loc[grp.index.isin(idx_eoScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_eoScene), 'Element'].apply(
            lambda x: str(x) + '</s>' )
      
      # calculate breakup n if decoder-input exceeds 1024
      str_masked = grp.loc[grp.index.isin(idx_w_mask), 'mgold'].apply(
        lambda x: ' '.join(str(x))).sum()
      tc_masked = len(str(str_masked).split(' '))
      if tc_masked > 1024:
        n = tc_masked//1024 + 1
      else:
        n = 1

      for i in range(n):
        agrp = grp.copy()
        agrp['mgold'] = np.NaN
        idx_agrp = grp['mgold'].dropna().sample(
          frac=1/n, random_state=random_state+i).index
        agrp.loc[agrp.index.isin(idx_agrp), 'mgold'] = \
          grp.loc[grp.index.isin(idx_agrp), 'mgold']

        d = {}
        agrp_fillna = agrp['Element_masked'].fillna(agrp['Element'])

        if shuffle_element:
          d['input_str'] = agrp_fillna.sample(
              frac=1,
              random_state=random_state+i
          ).sum()
        elif shuffle_scene:
          d['input_str'] = agrp.groupby('Scene').sample(
            frac=1,
            random_state=random_state+i
          )['Element_masked'].fillna(
            agrp['Element']).sum()
        else:
          d['input_str'] = agrp_fillna.sum()

        # if Global Attention on Scene
        if GA_on_scene:
          d['input_str'] = '<s>' + d['input_str'] + '</s>'

        if not agrp.loc[agrp.index.isin(idx_agrp), 'mgold'].empty:
          d['decoder_input_str'] = '<s>' + agrp['Element'].sum()
          d['labels_str'] = agrp['Element'].sum() + '</s>'
        else:
          d['decoder_input_str'] = '<s>'
          d['labels_str'] = '</s>'
        
        cases.append(d)
    return pd.DataFrame(cases)  

 
  @staticmethod
  def create_pretrainig_RAAC(
    dfsc,
    split_size=1024,
    frac=0.5,
    top_n_char=10,
    shuffle_element=False,
    shuffle_scene=False,
    random_state=42,
    GA_on_scene=True
    )-> pd.DataFrame:

    # Instantiate Detokenizer
    detokenizer = TreebankWordDetokenizer()
    
    # Get reg patterns
    pat_pair, pat_topchars_for_word_search, pat_topchars_for_str_search = \
      SC_Elements.generate_repattern_for_topchars(dfsc=dfsc, top_n_char=top_n_char)
    
    # Identify elements to mask
    def to_mask(ls_sent) -> bool:
      local_ls_sent = copy.deepcopy(ls_sent)
      num_char_in_sent = []
      for sent in local_ls_sent:
        char_match_in_sent = []
        for word in sent:
          if re.match(pat_topchars_for_word_search, word, flags=re.IGNORECASE):
            char_match_in_sent.append(1)
          else:
            char_match_in_sent.append(0)
        num_char_in_sent.append(sum(char_match_in_sent))
        if max(num_char_in_sent) > 1:
          return True
      return False
    ##
    dfsc_to_mask = dfsc[dfsc['Element_word'].apply(to_mask)].copy()
    ##

  # Apply Mask and Get Gold
    def apply_mask_to_Element_word(ls_ls_word):
      local_ls_ls_word = copy.deepcopy(ls_ls_word)
      num_char_in_sent = []
      gold = []
      ls_sent = []
      for i, sent in enumerate(local_ls_ls_word):
        char_match_in_sent = []
        for j, word in enumerate(sent):
          if re.match(pat_topchars_for_word_search, word, flags=re.IGNORECASE):
            char_match_in_sent.append(1)
          else:
            char_match_in_sent.append(0)
        
        if sum(char_match_in_sent) > 1:
          sent_gold = []
          sent_mask = []
          j = 0
          while j < len(sent):
            if not re.match(pat_topchars_for_word_search, sent[j], flags=re.IGNORECASE):
              sent_gold.append(sent[j])
              if not sent_mask:
                sent_mask.append('<mask>')
              elif sent_mask and sent_mask[-1] != '<mask>':
                sent_mask.append('<mask>')
            else:
              sent_mask.append(sent[j])
            j += 1
          gold.append(' '.join(sent_gold))
          ls_sent.append(' '.join(sent_mask))
        return [' '.join(ls_sent), ' '.join(gold)]
    ##
    dfsc_to_mask['Element_masked'] = dfsc_to_mask['Element_word'].apply(
      lambda x: apply_mask_to_Element_word(x)[0])
    dfsc_to_mask['gold'] = dfsc_to_mask['Element_word'].apply(
      lambda x: apply_mask_to_Element_word(x)[1])

    idx_w_mask = dfsc_to_mask.index
    dfsc['Element_masked'] = np.NaN
    dfsc.loc[dfsc.index.isin(idx_w_mask), 'Element_masked'] = \
     dfsc_to_mask['Element_masked'].copy()

    dfsc['mgold'] = np.NaN
    dfsc.loc[dfsc.index.isin(idx_w_mask), 'mgold'] = \
     dfsc_to_mask['gold'].copy()

    # Split dfsc with split_size
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for grp in grps:

      if GA_on_scene:
        idx_boScene = grp.drop_duplicates('Scene', keep='first')
        idx_eoScene = grp.drop_duplicates('Scene', keep='last')
        grp.loc[grp.index.isin(idx_boScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_boScene), 'Element'].apply(
            lambda x: '<s>'+str(x))
        grp.loc[grp.index.isin(idx_eoScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_eoScene), 'Element'].apply(
            lambda x: str(x) + '</s>' )
      
      # calculate breakup n if decoder-input exceeds 1024
      str_masked = grp.loc[grp.index.isin(idx_w_mask), 'mgold'].apply(
        lambda x: x + ' ').sum()
      tc_masked = len(str(str_masked).split(' '))
      if tc_masked > 1024:
        n = tc_masked//1024 + 1
      else:
        n = 1

      for i in range(n):
        agrp = grp.copy()
        agrp['mgold'] = np.NaN
        idx_agrp = grp['mgold'].dropna().sample(
          frac=1/n, random_state=random_state+i).index
        agrp.loc[agrp.index.isin(idx_agrp), 'mgold'] = \
          grp.loc[grp.index.isin(idx_agrp), 'mgold']
        idx_agrp = sorted(idx_agrp)

        agrp['Element_masked'] = np.NaN
        agrp.loc[agrp.index.isin(idx_agrp), 'Element_masked'] = \
         grp.loc[grp.index.isin(idx_agrp), 'Element_masked']

        d = {}

        agrp_fillna = agrp['Element_masked'].fillna(agrp['Element'])
        if shuffle_element:
          d['input_str'] = agrp_fillna.sample(
              frac=1,
              random_state=random_state+i
          ).sum()
        elif shuffle_scene:
          d['input_str'] = agrp.groupby('Scene').sample(
            frac=1,
            random_state=random_state+i
          )['Element_masked'].fillna(
            agrp['Element']).sum()
        else:
          d['input_str'] = agrp['Element_masked'].fillna(agrp['Element']).sum()

        # if Global Attention on Scene
        if GA_on_scene:
          d['input_str'] = '<s>' + d['input_str'] + '</s>'

        if not agrp.loc[agrp.index.isin(idx_agrp), 'mgold'].empty:
          d['decoder_input_str'] = '<s>' + agrp['Element'].sum()
          d['labels_str'] = agrp['Element'].sum() + '</s>'
        else:
          d['decoder_input_str'] = '<s>'
          d['labels_str'] = '</s>'
        
        cases.append(d)
    return pd.DataFrame(cases)  


  
  @staticmethod
  def create_pretraining_RD(dfsc, 
                            split_size=8192,
                            frac=0.4,
                            shuffle_element=False,
                            shuffle_scene=False,
                            random_state=42,
                            GA_on_scene=True,
                            ) -> pd.DataFrame:
    """
    Recover Dialog
    """
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for grp in grps:
      if GA_on_scene:
        idx_boScene = grp.drop_duplicates('Scene', keep='first')
        idx_eoScene = grp.drop_duplicates('Scene', keep='last')
        grp.loc[grp.index.isin(idx_boScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_boScene), 'Element'].apply(
            lambda x: '<s>'+str(x))
        grp.loc[grp.index.isin(idx_eoScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_eoScene), 'Element'].apply(
            lambda x: str(x) + '</s>' )
      
      for j in range(10):
        local_grp = grp.copy()
        agrp = local_grp[local_grp.Type == 'Dialogue']
        if agrp.empty:
          continue
        agrp = agrp.groupby('Scene').sample(
          frac=frac,
          random_state=random_state+j
        )

        d = {}
        idx_to_mask = agrp.index
        local_grp['Element_masked'] = local_grp['Element'].copy()
        local_grp.loc[local_grp.index.isin(idx_to_mask), 'Element_masked'] = '<mask>\n'
        if shuffle_element:
          d['input_str'] = local_grp['Element_masked'].sample(
            frac=1, 
            random_state=random_state+j).sum()
        elif shuffle_scene:
          d['input_str'] = local_grp.groupby('Scene').sample(
            frac=1,
            random_state=random_state+j
          )['Element_masked'].sum()
        else:
          d['input_str'] = local_grp['Element_masked'].sum()

        if GA_on_scene:
          d['input_str'] = '<s>' + d['input_str'] + '</s>'


        if not local_grp.loc[local_grp.index.isin(idx_to_mask), 'Element'].empty:
         d['decoder_input_str'] = '<s>' + \
          local_grp.loc[local_grp.index.isin(idx_to_mask), 'Element'].apply(
           lambda x: '<s>' + x + '</s>').sum()
         d['labels_str'] = local_grp.loc[local_grp.index.isin(idx_to_mask), 'Element'].apply(
           lambda x: '<s>' + x + '</s>').sum() + '</s>'
        else:
          d['decoder_input_str'] = '<s>'
          d['labels_str'] = '</s>'

        
        cases.append(d)

    return pd.DataFrame(cases)

  @staticmethod
  def create_pretraining_RL(dfsc, 
                            split_size=8192,
                            frac=1,
                            shuffle_element=False,
                            shuffle_scene=False,
                            random_state=42,
                            GA_on_scene=True,
                            ):
  
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for  grp in grps:
      if GA_on_scene:
        idx_boScene = grp.drop_duplicates('Scene', keep='first')
        idx_eoScene = grp.drop_duplicates('Scene', keep='last')
        grp.loc[grp.index.isin(idx_boScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_boScene), 'Element'].apply(
            lambda x: '<s>'+str(x))
        grp.loc[grp.index.isin(idx_eoScene), 'Element'] = \
          grp.loc[grp.index.isin(idx_eoScene), 'Element'].apply(
            lambda x: str(x) + '</s>' )
      frac_local = frac
      repeat = 0
      i = 0
      sh_to_mask = grp.loc[grp.Type == 'Scene Heading', :]
      while sh_to_mask.shape[0] > 320:
        frac_local -= 0.1
        sh_to_mask = grp.loc[grp.Type == 'Scene Heading', :].sample(
        frac=frac_local, 
        random_state=random_state+i
        )
        i += 1
      
      while i > -1:
        d = {}
        idx_to_mask = sh_to_mask.index
        grp['Element_masked'] = grp['Element'].copy()
        grp.loc[grp.index.isin(idx_to_mask), 'Element_masked'] = '<mask>\n'
        if shuffle_element:
          d['input_str'] = grp['Element_masked'].sample(
            frac=1, 
            random_state=random_state).sum()
        elif shuffle_scene:
          d['input_str'] = grp.groupby('Scene').sample(
            frac=1,
            random_state=random_state
          )['Element_masked'].sum()
        else:
          d['input_str'] = grp['Element_masked'].sum()

        if GA_on_scene:
          d['input_str'] = '<s>' + d['input_str'] + '</s>'

        #remove '\n'
        str_masked_dchars = str(grp.loc[grp.index.isin(idx_to_mask), 'Element'].apply(
          lambda x: re.sub('\n', ' ', x)).sum())

        d['decoder_input_str'] = '<s>' + str_masked_dchars
        d['labels_str'] = str_masked_dchars + '</s>'
        cases.append(d)
        sh_to_mask = grp.loc[grp.Type == 'Scene Heading', :].sample(
          frac=frac_local, 
          random_state=random_state+i
          )
        i -= 1

    return pd.DataFrame(cases)

  @staticmethod
  def create_pretraining_RCID(dfsc, 
                              split_size=1024,
                              frac=1,
                              shuffle_element=False,
                              shuffle_scene=False,
                              random_state=42,
                              GA_on_scene=True,
                              ):
   
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for  grp in grps:
      if GA_on_scene:
        idx_boScene = grp.drop_duplicates('Scene', keep='first')
        idx_eoScene = grp.drop_duplicates('Scene', keep='last')
        grp.loc[grp.index.isin(idx_boScene), 'Element'] = \
         grp.loc[grp.index.isin(idx_boScene), 'Element'].apply(
           lambda x: '<s>'+str(x))
        grp.loc[grp.index.isin(idx_eoScene), 'Element'] = \
         grp.loc[grp.index.isin(idx_eoScene), 'Element'].apply(
           lambda x: str(x) + '</s>' )
      frac_local = frac
      repeat = 0
      i = 0
      dchar_to_mask = grp.loc[grp.Type == 'Character', :]
      while dchar_to_mask.shape[0] > 320:
        frac_local -= 0.1
        dchar_to_mask = grp.loc[grp.Type == 'Character', :].sample(
        frac=frac_local, 
        random_state=random_state+i
        )
        i += 1
      
      while i > -1:
        d = {}
        idx_to_mask = dchar_to_mask.index
        grp['Element_masked'] = grp['Element'].copy()
        grp.loc[grp.index.isin(idx_to_mask), 'Element_masked'] = '<mask>\n'
        if shuffle_element:
          d['input_str'] = grp['Element_masked'].sample(
            frac=1, 
            random_state=random_state).sum()
        elif shuffle_scene:
          d['input_str'] = grp.groupby('Scene').sample(
            frac=1,
            random_state=random_state
          )['Element_masked'].sum()
        else:
          d['input_str'] = grp['Element_masked'].sum()

        if GA_on_scene:
          d['input_str'] = '<s>' + d['input_str'] + '</s>'

        #remove '\n'
        str_masked_dchars = str(grp['Element'].apply(
          lambda x: re.sub('\n', ' ', x)).sum())

        d['decoder_input_str'] = '<s>' + str_masked_dchars
        d['labels_str'] = str_masked_dchars + '</s>'
        cases.append(d)
        dchar_to_mask = grp.loc[grp.Type == 'Character', :].sample(
          frac=frac_local, 
          random_state=random_state+i
          )
        i -= 1

    return pd.DataFrame(cases)

  @staticmethod
  def create_pretraining_TI(
                            dfsc, 
                            split_size=4096,
                            frac=0.2,
                            shuffle_element=False,
                            random_state=42,
                            ):
    """This funtion create pretraining cases with optional element shuffle
    dfsc: pd.DataFrame
    split_size=4096
    frac=0.2,
    shuffle_element=False,
    random_state=42
    """
    
    i = 0
    off_set = split_size//16
    grps = []
    while i < dfsc['tc_cumsum'].max():
      grp = dfsc[(dfsc.tc_cumsum >= i) & (dfsc.tc_cumsum < i+split_size)].copy()
      grps.append(grp)
      i += off_set

    cases = []
    for  grp in grps:
      for j in range(10):
        agrp = grp.groupby('Scene').sample(frac=frac, random_state=j).copy()
        if (grp['tc'].sum() < int(split_size*1.5) and 
            400 <= agrp['tc'].sum() < 960):
          d = {}
          idx_to_mask = agrp.index
          grp['Element_masked'] = grp['Element'].copy()
          grp.loc[grp.index.isin(idx_to_mask), 'Element_masked'] = '<mask>'
          if shuffle_element:
            d['input_str'] = grp['Element_masked'].sample(
              frac=1, 
              random_state=random_state).sum()
          else:
            d['input_str'] = grp['Element_masked'].sum()
          d['decoder_input_str'] = '<s>' + str(grp.loc[grp.index.isin(idx_to_mask), 'Element'].sum())
          d['labels_str'] = str(grp.loc[grp.index.isin(idx_to_mask), 'Element'].sum()) + '</s>'
          cases.append(d)

    return pd.DataFrame(cases)
  
  @staticmethod
  def extract_str_by_method(dfsc, method='AF', return_type='str', n_main=1):
    """
    method = 'AF', 'AO', 'DO', 'FMC', 'FMCAO', 
    'EVERY2at0', 'EVERY2at1',
    'EVERY3at0', 'EVERY3at1', 'EVERY3at2',
    'EVERY4at0', 'EVERY4at1', 'EVERY4at2', 'EVERY4at3'
    F: All Flat
    AO: Action Only
    DO: Dialog Only
    FMC: Follow Main Character
    FMCAO: Follow Main Character + Action Only

    return_type= 'str', 'df', 'df_full'
    """
    dfsc['Element'] = dfsc['Element'].apply(lambda x: str(x).lstrip('\w+').rstrip('\w+').rstrip('\n') + '\n')
    dfsc['Element'] = dfsc['Element'].apply(lambda x: re.sub(':SC:', '', str(x)) )
    dfsc['Scene'] = dfsc['Scene'].fillna(-1)

    if not method:
      return dfsc

    ########
    if method == 'PSentAO':
     dfsc_w_method = dfsc.loc[dfsc['Type'] == 'Action', :].copy()
     dfsent = dfsc_w_method.copy()
     dfsent['Element_sent'] = dfsent['Element'].apply(lambda x: sent_tokenize(x))
     dfsent = dfsent.explode('Element_sent').reset_index()
     dfsent['Element_sent_word'] = dfsent['Element_sent'].astype('str').apply(word_tokenize)
     # Counter of sentences
     dfsent['counter_sent_word'] = dfsent['Element_sent_word'].apply(collections.Counter)
     dfsent['wc_sent'] = dfsent['counter_sent_word'].apply(lambda x: sum(x.values()))
     # Counter of the entire screenplay
     Counter_SC = dfsent['counter_sent_word'].sum()
     dfsent['counter_rest'] = dfsent['counter_sent_word'].apply(lambda x: Counter_SC - x)
     dfsent['wc_rest'] = dfsent['counter_rest'].apply(lambda x: sum(x.values()))
     # Calculate overlap between sent and the rest
     dfsent['counter_diff'] = dfsent['counter_rest'] - dfsent['counter_sent_word']
     # Filter 0 denominator
     dfsent = dfsent[(dfsent.wc_rest > 0) | dfsent.wc_rest > 0]
     # Calculate Rouge 1
     dfsent['wc_overlap_sent_to_rest'] = dfsent['wc_rest'] - dfsent['counter_diff'].apply(lambda x: sum(x.values()))
     #dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest']
     dfsent['recall'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_sent']

     #dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall'])
     dfsent = dfsent.sort_values('recall', ascending=False)
     dfsent['wc_cum'] = dfsent['wc_sent'].cumsum()
     df_selection = dfsent[dfsent['wc_cum'] < 12000].sort_index()
     df_selection['txt'] = df_selection['Element_sent'].fillna('').astype('str')

     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()
    ########
    if method == 'PSentAO_F1':
     dfsc_w_method = dfsc.loc[dfsc['Type'] == 'Action', :].copy()
     dfsent = dfsc_w_method.copy()
     dfsent['Element_sent'] = dfsent['Element'].apply(lambda x: sent_tokenize(x))
     dfsent = dfsent.explode('Element_sent').reset_index()
     dfsent['Element_sent_word'] = dfsent['Element_sent'].astype('str').apply(word_tokenize)
     # Counter of sentences
     dfsent['counter_sent_word'] = dfsent['Element_sent_word'].apply(collections.Counter)
     dfsent['wc_sent'] = dfsent['counter_sent_word'].apply(lambda x: sum(x.values()))
     # Counter of the entire screenplay
     Counter_SC = dfsent['counter_sent_word'].sum()
     dfsent['counter_rest'] = dfsent['counter_sent_word'].apply(lambda x: Counter_SC - x)
     dfsent['wc_rest'] = dfsent['counter_rest'].apply(lambda x: sum(x.values()))
     # Calculate overlap between sent and the rest
     dfsent['counter_diff'] = dfsent['counter_rest'] - dfsent['counter_sent_word']
     # Filter 0 denominator
     dfsent = dfsent[(dfsent.wc_rest > 0) | dfsent.wc_rest > 0]
     # Calculate Rouge 1
     dfsent['wc_overlap_sent_to_rest'] = dfsent['wc_rest'] - dfsent['counter_diff'].apply(lambda x: sum(x.values()))
     dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest']
     dfsent['recall'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_sent']

     dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall'])
     dfsent = dfsent.sort_values('R1-F1', ascending=False)
     dfsent['wc_cum'] = dfsent['wc_sent'].cumsum()
     df_selection = dfsent[dfsent['wc_cum'] < 12000].sort_index()
     df_selection['txt'] = df_selection['Element_sent'].fillna('').astype('str')

     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()

    ########  
    if method == 'PSent_F1':
     dfsc_w_method = dfsc.copy()
     dfsent = dfsc_w_method.copy()
     dfsent['Element_sent'] = dfsent['Element'].apply(lambda x: sent_tokenize(x))
     dfsent = dfsent.explode('Element_sent').reset_index()
     dfsent['Element_sent_word'] = dfsent['Element_sent'].astype('str').apply(word_tokenize)
     # Counter of sentences
     dfsent['counter_sent_word'] = dfsent['Element_sent_word'].apply(collections.Counter)
     dfsent['wc_sent'] = dfsent['counter_sent_word'].apply(lambda x: sum(x.values()))
     # Counter of the entire screenplay
     Counter_SC = dfsent['counter_sent_word'].sum()
     dfsent['counter_rest'] = dfsent['counter_sent_word'].apply(lambda x: Counter_SC - x)
     dfsent['wc_rest'] = dfsent['counter_rest'].apply(lambda x: sum(x.values()))
     # Calculate overlap between sent and the rest
     dfsent['counter_diff'] = dfsent['counter_rest'] - dfsent['counter_sent_word']
     # Filter 0 denominator
     dfsent = dfsent[(dfsent.wc_rest > 0) | dfsent.wc_rest > 0]
     # Calculate Rouge 1
     dfsent['wc_overlap_sent_to_rest'] = dfsent['wc_rest'] - dfsent['counter_diff'].apply(lambda x: sum(x.values()))
     dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest']
     dfsent['recall'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_sent']

     dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall'])
     dfsent = dfsent.sort_values('R1-F1', ascending=False)
     dfsent['wc_cum'] = dfsent['wc_sent'].cumsum()
     df_selection = dfsent[dfsent['wc_cum'] < 12000].sort_index()
     df_selection['txt'] = df_selection['Element_sent'].fillna('').astype('str')

     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()
    ########  
    if method == 'PSent':
     dfsc_w_method = dfsc.copy()
     dfsent = dfsc_w_method.copy()
     dfsent['Element_sent'] = dfsent['Element'].apply(lambda x: sent_tokenize(x))
     dfsent = dfsent.explode('Element_sent').reset_index()
     dfsent['Element_sent_word'] = dfsent['Element_sent'].astype('str').apply(word_tokenize)
     # Counter of sentences
     dfsent['counter_sent_word'] = dfsent['Element_sent_word'].apply(collections.Counter)
     dfsent['wc_sent'] = dfsent['counter_sent_word'].apply(lambda x: sum(x.values()))
     # Counter of the entire screenplay
     Counter_SC = dfsent['counter_sent_word'].sum()
     dfsent['counter_rest'] = dfsent['counter_sent_word'].apply(lambda x: Counter_SC - x)
     dfsent['wc_rest'] = dfsent['counter_rest'].apply(lambda x: sum(x.values()))
     # Calculate overlap between sent and the rest
     dfsent['counter_diff'] = dfsent['counter_rest'] - dfsent['counter_sent_word']
     # Filter 0 denominator
     dfsent = dfsent[(dfsent.wc_rest > 0) | dfsent.wc_rest > 0]
     # Calculate Rouge 1
     dfsent['wc_overlap_sent_to_rest'] = dfsent['wc_rest'] - dfsent['counter_diff'].apply(lambda x: sum(x.values()))
     #dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest']
     dfsent['recall'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_sent']

     #dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall'])
     dfsent = dfsent.sort_values('recall', ascending=False)
     dfsent['wc_cum'] = dfsent['wc_sent'].cumsum()
     df_selection = dfsent[dfsent['wc_cum'] < 12000].sort_index()
     df_selection['txt'] = df_selection['Element_sent'].fillna('').astype('str')

     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()
    ########  
    if method == 'PSceneH5T5AO':      
     dfsc_w_method = dfsc.loc[dfsc['Type'] == 'Action', :].copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     #dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     #dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])  
     
     # split dataframe to first5, center and last5
     idx_first5 = dfscene.index[:5]
     idx_last5 = dfscene.index[-5:]
     idx_center = dfscene.index[5:-5]
     first5_wc = dfscene.loc[idx_first5, 'wc_scene'].sum()
     last5_wc = dfscene.loc[idx_last5, 'wc_scene'].sum()
     # Sort center df by R1-F1
     df_center = dfscene.loc[idx_center, :].sort_values(
      'recall', ascending=False)
     df_center['cum_wc']  = df_center['wc_scene'].cumsum()
     idx_ros = df_center[df_center.cum_wc < (12000-first5_wc-last5_wc)].index
     df_selection= dfscene.loc[idx_first5.union(idx_ros).union(idx_last5),:]
     df_selection = df_selection.sort_values('Scene')

     if return_type == 'str':
       ret = df_selection['txt'].sum()
     if return_type == 'df':
       ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
       ret = df_selection.copy()
    ########  
    if method == 'PSceneH5T5AO_F1':      
     dfsc_w_method = dfsc.loc[dfsc['Type'] == 'Action', :].copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])  
     
     # split dataframe to first5, center and last5
     idx_first5 = dfscene.index[:5]
     idx_last5 = dfscene.index[-5:]
     idx_center = dfscene.index[5:-5]
     first5_wc = dfscene.loc[idx_first5, 'wc_scene'].sum()
     last5_wc = dfscene.loc[idx_last5, 'wc_scene'].sum()
     # Sort center df by R1-F1
     df_center = dfscene.loc[idx_center, :].sort_values(
      'R1-F1', ascending=False)
     df_center['cum_wc']  = df_center['wc_scene'].cumsum()
     idx_ros = df_center[df_center.cum_wc < (12000-first5_wc-last5_wc)].index
     df_selection= dfscene.loc[idx_first5.union(idx_ros).union(idx_last5),:]
     df_selection = df_selection.sort_values('Scene')

     if return_type == 'str':
       ret = df_selection['txt'].sum()
     if return_type == 'df':
       ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
       ret = df_selection.copy()
    ##########################
    if method == 'PSceneH5T5':
      
     dfsc_w_method = dfsc.copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     #dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     #dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])  
     
     # split dataframe to first5, center and last5
     idx_first5 = dfscene.index[:5]
     idx_last5 = dfscene.index[-5:]
     idx_center = dfscene.index[5:-5]
     first5_wc = dfscene.loc[idx_first5, 'wc_scene'].sum()
     last5_wc = dfscene.loc[idx_last5, 'wc_scene'].sum()
     # Sort center df by R1-F1
     df_center = dfscene.loc[idx_center, :].sort_values('recall', ascending=False)
     df_center['cum_wc']  = df_center['wc_scene'].cumsum()
     idx_ros = df_center[df_center.cum_wc < (12000-first5_wc-last5_wc)].index
     df_selection= dfscene.loc[idx_first5.union(idx_ros).union(idx_last5),:]
     df_selection = df_selection.sort_values('Scene')

     if return_type == 'str':
       ret = df_selection['txt'].sum()
     if return_type == 'df':
       ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
       ret = df_selection.copy()

    ##########################
    if method == 'PSceneH5T5_F1':
      
     dfsc_w_method = dfsc.copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])  
     
     # split dataframe to first5, center and last5
     idx_first5 = dfscene.index[:5]
     idx_last5 = dfscene.index[-5:]
     idx_center = dfscene.index[5:-5]
     first5_wc = dfscene.loc[idx_first5, 'wc_scene'].sum()
     last5_wc = dfscene.loc[idx_last5, 'wc_scene'].sum()
     # Sort center df by R1-F1
     df_center = dfscene.loc[idx_center, :].sort_values('R1-F1', ascending=False)
     df_center['cum_wc']  = df_center['wc_scene'].cumsum()
     idx_ros = df_center[df_center.cum_wc < (12000-first5_wc-last5_wc)].index
     df_selection= dfscene.loc[idx_first5.union(idx_ros).union(idx_last5),:]
     df_selection = df_selection.sort_values('Scene')

     if return_type == 'str':
       ret = df_selection['txt'].sum()
     if return_type == 'df':
       ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
       ret = df_selection.copy()
    ##########################
    if method == 'PScene':

     dfsc_w_method = dfsc.copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     #dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     #dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])
     dfscene = dfscene.sort_values('recall', ascending=False)
     dfscene['cum_wc'] = dfscene['wc_scene'].cumsum()
     # Make Selection by R1-F1 rank
     df_selection = dfscene[dfscene.cum_wc < 12000].sort_values(
     'Scene', ascending=True)
     
     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()
    ##########################
    if method == 'PScene_F1':

     dfsc_w_method = dfsc.copy()
     dfscene = dfsc_w_method.groupby('Scene')['Element'].sum().rename('txt').reset_index()
     dfscene['tokens'] = dfscene['txt'].apply(word_tokenize)
     # Counter of the scene and wc of the scene
     dfscene['counter_scene'] = dfscene['tokens'].apply(collections.Counter)
     dfscene['wc_scene'] = dfscene['counter_scene'].apply(lambda x: sum(x.values()))
     # Counter of the entire screeplay text
     Counter_SC = dfscene['counter_scene'].sum()
     # Counter of the rest and wc of the rest
     dfscene['counter_rest'] = dfscene['counter_scene'].apply(lambda x: Counter_SC - x)
     dfscene['wc_rest'] = dfscene['counter_rest'].apply(lambda x: sum(x.values()))
     # Drop wc_scene=0 and wc_rest=0 for 0 denominator division
     dfscene = dfscene[(dfscene.wc_scene > 0) | (dfscene.wc_rest > 0)]
     # Calculate overlap between scene and the rest
     dfscene['counter_diff'] = dfscene['counter_rest'] - dfscene['counter_scene']
     dfscene['wc_overlap_scene_to_rest'] = dfscene['wc_rest'] - dfscene['counter_diff'].apply(lambda x: sum(x.values()))
     # Calculate Rouge-1 Precision, Recall and F1
     dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest']
     dfscene['recall'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_scene']
     dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall'])
     dfscene = dfscene.sort_values('R1-F1', ascending=False)
     dfscene['cum_wc'] = dfscene['wc_scene'].cumsum()
     # Make Selection by R1-F1 rank
     df_selection = dfscene[dfscene.cum_wc < 12000].sort_values(
     'Scene', ascending=True)
     
     if return_type == 'str':
        ret = df_selection['txt'].sum()
     if return_type == 'df':
        ret = df_selection[['Scene', 'txt']].groupby('Scene')['txt'].sum()
     if return_type =='df_full':
        ret = df_selection.copy()
    ##########################
    if method == 'EVERY2at0':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 2 == 1, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY2at1':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 2 == 0, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY3at0':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 3 == 0, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY3at1':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 3 == 1, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY3at2':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 3 == 2, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY4at0':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 4 == 0,: ].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY4at1':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 4 == 1, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY4at2':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 4 == 2, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'EVERY4at3':
      dfsc_w_method = dfsc.loc[dfsc['Scene'] % 4 == 3, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'AF':
      dfsc_w_method = dfsc.copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'AO':
      dfsc_w_method = dfsc.loc[dfsc['Type'] == 'Action', :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'DO':
      #dfsc.loc[dfsc['Type'].str.contains('Character'), 'Element'] = \
      #  dfsc.loc[dfsc['Type'].str.contains('Character'), 'Element'].apply(lambda x: x.rstrip('\n') + ':')
      dfsc_w_method = dfsc.loc[dfsc['Type'].str.contains('Character|Dialogue'), :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'FMC':
      dfsc['Element'] = dfsc['Element'].astype('str')
      df_chars = dfsc.loc[dfsc['Type']=='Character', ['Element', 'Type']].groupby(
        'Element').count().reset_index()
      df_chars.columns = ['Char', 'Count']
      df_chars = df_chars.sort_values(by='Count', ascending=False)
      ls_Scene = dfsc.loc[dfsc['Element'].isin(df_chars['Char'][:n_main]), 'Scene'].drop_duplicates().tolist()
      dfsc_w_method = dfsc.loc[dfsc['Scene'].isin(ls_Scene), :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()     
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'FMCAO':
      dfsc['Element'] = dfsc['Element'].astype('str')
      df_chars = dfsc.loc[dfsc['Type']=='Character', ['Element', 'Type']].groupby(
        'Element').count().reset_index()
      df_chars.columns = ['Char', 'Count']
      df_chars = df_chars.sort_values(by='Count', ascending=False)
      ls_Scene = dfsc.loc[dfsc['Element'].isin(df_chars['Char'][:n_main]), 'Scene'].drop_duplicates().tolist()
      dfsc_w_method = dfsc.loc[(dfsc['Scene'].isin(ls_Scene)) & (dfsc['Type'] == 'Action'), :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()
      if return_type =='df_full':
        ret = dfsc_w_method

    if method == 'TPS':
      dfsc_w_method = dfsc.loc[dfsc['tps'] > 0, :].copy()
      if return_type == 'str':
        ret = dfsc_w_method['Element'].sum()
      if return_type == 'df':
        ret = dfsc_w_method.groupby('Scene')['Element'].sum()      
      if return_type =='df_full':
        ret = dfsc_w_method

    if return_type == 'str':
      return ret
    if return_type == 'df':
      return ret.copy()
    if return_type == 'df_full':
      return ret.copy()
# Counter of sentences # Counter of the entire screenplay # Calculate overlap between sent and the rest # Filter 0 denominator # Calculate Rouge 1 # GSceneG TOP 30% # Counter of the scene and wc of the scene # Counter of the entire screeplay text # Counter of the rest and wc of the rest # Drop wc_scene=0 and wc_rest=0 for 0 denominator division # Calculate overlap between scene and the rest # Calculate Rouge-1 Precision, Recall and F1 # Make Selection by R1-F1 rank with top 30% ########## returns        pat_topchars, # identify top appearance characters by Character # calcualte appearances by name in Actions # combine and re-sort # create pairs from Characters # Create pat for re #ls_pat.append(right+'.*'+left) RECOVER CHARACTER ACROSS ACTION     (1) Identify top n characters -> List[str]=['Tom', 'Jerry'] # Instantiate Detokenizer # Get reg patterns # Identify elements to mask ## ## # Apply Mask and Get Gold ## ## ## ## # Split dfsc with split_size # calculate breakup n if decoder-input exceeds 1024 # if Global Attention on Scene # Instantiate Detokenizer # Get reg patterns # Identify elements to mask ## ## # Apply Mask and Get Gold ## # Split dfsc with split_size # calculate breakup n if decoder-input exceeds 1024 # if Global Attention on Scene Recover Dialog #remove '\n' #remove '\n' This funtion create pretraining cases with optional element shuffle     dfsc: pd.DataFrame     split_size=4096     frac=0.2,     shuffle_element=False,     random_state=42 method = 'AF', 'AO', 'DO', 'FMC', 'FMCAO',      'EVERY2at0', 'EVERY2at1',     'EVERY3at0', 'EVERY3at1', 'EVERY3at2',     'EVERY4at0', 'EVERY4at1', 'EVERY4at2', 'EVERY4at3'     F: All Flat     AO: Action Only     DO: Dialog Only     FMC: Follow Main Character     FMCAO: Follow Main Character + Action Only      return_type= 'str', 'df', 'df_full' ######## # Counter of sentences # Counter of the entire screenplay # Calculate overlap between sent and the rest # Filter 0 denominator # Calculate Rouge 1 #dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest'] #dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall']) ######## # Counter of sentences # Counter of the entire screenplay # Calculate overlap between sent and the rest # Filter 0 denominator # Calculate Rouge 1 ######## # Counter of sentences # Counter of the entire screenplay # Calculate overlap between sent and the rest # Filter 0 denominator # Calculate Rouge 1 ######## # Counter of sentences # Counter of the entire screenplay # Calculate overlap between sent and the rest # Filter 0 denominator # Calculate Rouge 1 #dfsent['precision'] = dfsent['wc_overlap_sent_to_rest'] / dfsent['wc_rest'] #dfsent['R1-F1'] = (2*dfsent['precision']*dfsent['recall']) / (dfsent['precision']+dfsent['recall']) ######## # Counter of the scene and wc of the scene # Counter of the entire screeplay text # Counter of the rest and wc of the rest # Drop wc_scene=0 and wc_rest=0 for 0 denominator division # Calculate overlap between scene and the rest # Calculate Rouge-1 Precision, Recall and F1 #dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest'] #dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall']) # split dataframe to first5, center and last5 # Sort center df by R1-F1 ######## # Counter of the scene and wc of the scene # Counter of the entire screeplay text # Counter of the rest and wc of the rest # Drop wc_scene=0 and wc_rest=0 for 0 denominator division # Calculate overlap between scene and the rest # Calculate Rouge-1 Precision, Recall and F1 # split dataframe to first5, center and last5 # Sort center df by R1-F1 ########################## # Counter of the scene and wc of the scene # Counter of the entire screeplay text # Counter of the rest and wc of the rest # Drop wc_scene=0 and wc_rest=0 for 0 denominator division # Calculate overlap between scene and the rest # Calculate Rouge-1 Precision, Recall and F1 #dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest'] #dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall']) # split dataframe to first5, center and last5 # Sort center df by R1-F1 ########################## # Counter of the scene and wc of the scene # Counter of the entire screeplay text # Counter of the rest and wc of the rest # Drop wc_scene=0 and wc_rest=0 for 0 denominator division # Calculate overlap between scene and the rest # Calculate Rouge-1 Precision, Recall and F1 # split dataframe to first5, center and last5 # Sort center df by R1-F1 ########################## # Counter of the scene and wc of the scene # Counter of the entire screeplay text # Counter of the rest and wc of the rest # Drop wc_scene=0 and wc_rest=0 for 0 denominator division # Calculate overlap between scene and the rest # Calculate Rouge-1 Precision, Recall and F1 #dfscene['precision'] = dfscene['wc_overlap_scene_to_rest'] / dfscene['wc_rest'] #dfscene['R1-F1'] = (2*dfscene['precision']*dfscene['recall']) / (dfscene['precision']+dfscene['recall']) # Make Selection by R1-F1 rank ########################## # Counter of the scene and wc of the scene # Counter of the entire screeplay text # Counter of the rest and wc of the rest # Drop wc_scene=0 and wc_rest=0 for 0 denominator division # Calculate overlap between scene and the rest # Calculate Rouge-1 Precision, Recall and F1 # Make Selection by R1-F1 rank ########################## #dfsc.loc[dfsc['Type'].str.contains('Character'), 'Element'] = \ #  dfsc.loc[dfsc['Type'].str.contains('Character'), 'Element'].apply(lambda x: x.rstrip('\n') + ':')