align_utils.py

import spacy
from simalign import SentenceAligner
from pathlib import Path
HERE = Path(__file__).parent


SIMALIGNER = SentenceAligner(model="xlmr", layer=8, token_type="bpe", matching_methods="i")
 

STEM_CACHE = {}

POS_CACHE = {}


_LEMMATIZER_MODELS = {
    'es': "es_core_news_lg",
    'it': "it_core_news_lg",
    'fr': "fr_core_news_lg",
    'en': "en_core_web_lg",
    'ro': "ro_core_news_lg",
    'zh': "zh_core_web_lg",
    'xx': "xx_ent_wiki_sm",
}

_LEMMATIZER_CACHE = {}

TAGDICT = {'NP': 'n', 'VBD': 'v', 'SPACE': 'x', 'DET': 'x', 'JJ': 'a', 'NN': 'n', 'NNS': 'n', 'ADV': 'r', 'ADJ': 'a', 'PRON': 'n', 'AUX': 'v', 'CCONJ': 'x', 'SCONJ': 'x', 'X': 'x',
           'VB': 'v', 'PP$': 'x', 'RB': 'r', 'VBN': 'v', 'IN': 'x', ',': 'x', 'SENT': 'x', 'PP': 'x', 'VERB': 'v', 'PROPN': 'n', 'NOUN': 'n', 'SYM': 'x', 'INTJ': 'x',
           'CD': 'x', 'VBZ': 'v', 'CC': 'x', 'MD': 'x', 'TO': 'x', 'PDT': 'x', 'WRB': 'x', 'POS': 'x', 'ADP': 'x', 'PUNCT': 'x', 'NUM': 'x', 'PART': 'x'}

PUNCTUATION = {'¿', '`', '.', ',', '!', '?', '.', ':', ';', '*', ' ', " ", '-', '«', 
               '.', ',', '!', '?', ':', ';', '*', '-', '—', '(', ')', '[', ']', '{', '}', 
    '"', "'", '«', '»', '/', '\\', '|', '@', '#', '$', '%', '^', '&', '_', '+', '«', '»', '¿', '?', 
    '°', '¡', '.', ',', '!', '?', '.', ':', ';', '*', ' ', " ", '-', '«', 
               '.', ',', '!', '?', ':', ';', '*', '-', '—', '(', ')', '[', ']', '{', '}', 
    '"', "'", '«', '»', '/', '\\', '|', '@', '#', '$', '%', '^', '&', '_', '+', 
    '=', '<', '>', '~', '`', '’', '、', '。', '。', '、', '·'
    '=', '<', '>', '~', '`'
}

DO_NOT_CALL_BABELNET = {'«', '»', '¿', '?', '°', '¡', '.', ',', '!', '?', '.', ':', ';', '*', ' ', " ", '-', '«', 
               '.', ',', '!', '?', ':', ';', '*', '-', '—', '(', ')', '[', ']', '{', '}', 
    '"', "'", '«', '»', '/', '\\', '|', '@', '#', '$', '%', '^', '&', '_', '+', "''"
    '=', '<', '>', '~', '`', '’', '、', '。', '。', '、', '·', '©', '·', '¢', '``', '``'} 



ALIGNMENT_CACHE = {}
SYNSET_CACHED_DICT = {}


def safe_split(large_string, split_by):
    if split_by == "":
        return list(large_string)
    return large_string.split(split_by)

    

def convert_quad_to_tuple(ours):
    ans = []
    for align in ours.split():
        sb, se, tb, te = align.split('-')
        sb, se, tb, te = int(sb), int(se), int(tb), int(te)
        for x in range(sb, se + 1):

            for y in range(tb, te + 1):
                ans.append((x, y))
    return ans
    
class DBAligner:
    def __init__(self, src_l, tar_l, dictionary='BN', path_to_dict=None, src_join_char='_', tgt_join_char='_'):
        self.source_language = src_l
        self.target_language = tar_l
        self.src_join_char = src_join_char
        self.tgt_join_char = tgt_join_char
        if dictionary.lower() == 'BN'.lower():
            self.dict_in_use = 'bn'
            self.dict = {}
            self.load_cache()
            global LANGS
            import babelnet as bn
            from babelnet import Language
            from babelnet.pos import POS
            LANGS = {"tr": Language.TR, "es": Language.ES, "it": Language.IT, "ar": Language.AR, "de": Language.DE,
           "fr": Language.FR, "en": Language.EN, "ja": Language.JA, "th": Language.TH, "zh": Language.ZH,
          " ko": Language.KO, "ro": Language.RO }
            global POS_TAGS
            POS_TAGS = {"VERB": POS.VERB, "NOUN": POS.NOUN, "ADV": POS.ADV, "ADJ": POS.ADJ, "PROPN": POS.NOUN}
            self.max_subseq_length = 10
            
        elif dictionary.lower() == 'custom':
            self.dict_in_use = 'custom'
            self.dict = {}
            self.load_dict(path_to_dict)
            
        else:
            assert False, "please pass third argument as either 'BN' or 'Custom'"
            
    def are_synonyms_by_dictionary(self, word_one, word_two, lemma_one, lemma_two, lang_two, lang_one):
        if self.dict_in_use == 'custom':
            total_ans =  self.are_synonyms_by_custom(word_one, word_two, lemma_one, lemma_two)
        elif self.dict_in_use == 'bn':
            total_ans = are_synonyms_by_bn(word_one, word_two, lang_two, lang_one)
        else:
            assert False
        if total_ans:
            return total_ans

        return False
            
    def are_synonyms_by_custom(self, first_word, second_word, first_lemma, second_lemma):
        
      
      
      first_word_first_form = first_word.replace(' ', self.src_join_char)
      
      dict_first_form = self.dict.get(first_word_first_form)
      if dict_first_form and second_word in dict_first_form:
            return 'strict'
      
      first_word_second_form = first_word.replace(self.src_join_char, ' ')
      dict_second_form = self.dict.get(first_word_second_form)
      
      if dict_second_form and second_word in dict_second_form:
          print("Unexpected match found, check your dictionary matches your choice of join characters")
          return 'strict'
      
      if self.tgt_join_char in second_word or self.src_join_char in first_word:
          
          split_first = safe_split(first_word, self.src_join_char)
          split_second = safe_split(second_word, self.tgt_join_char)
          
          for guy in split_first:
            if guy in PUNCTUATION: return False
          for guy in split_second:
            if guy in PUNCTUATION: return False
        
      dict_lemma1 = self.dict.get(first_lemma)
      
      
      if dict_lemma1 and (second_lemma in dict_lemma1 or second_word in dict_lemma1):
            return 'loose'
      
      dict_first_word = self.dict.get(first_word)
      if dict_first_word and second_lemma in dict_first_word:
            return 'loose'
    
      return False
  
    def using_custom(self):
        if self.dict_in_use == 'custom':
            return True
        return False
      
            
    
    def load_dict(self, addr):
        self.max_subseq_length = 0
        # Load a dictionary file
        found_words = set()
        with open(addr, 'r', encoding='utf-8') as infile:
            for line in infile:
                stripped_line = line.strip()
                assert len(stripped_line.split('\t')) == 2 or len(stripped_line) == 0, "Expected each line in dictionary file to have one tab character. The offending line is: " + str(line)
                input_word, translations = stripped_line.split('\t')
                input_word = input_word.replace(' ', self.src_join_char).replace('_', self.src_join_char)
                
                input_sub_length = input_word.count(self.src_join_char)
                
                if input_sub_length > self.max_subseq_length:
                    self.max_subseq_length = input_sub_length
                                        
                found_words.add(input_word)
                for translation in translations.split():
                    
                    if input_word in self.dict:
                        self.dict[input_word].append(translation.replace(' ', self.tgt_join_char).replace('_', self.tgt_join_char))
                    else:
                        self.dict[input_word] = [translation.replace(' ', self.tgt_join_char).replace('_', self.tgt_join_char)]
                    
                    trans_subseq_length = translation.replace(' ', self.tgt_join_char).replace('_', self.tgt_join_char).count(self.tgt_join_char)
                    if trans_subseq_length > self.max_subseq_length:
                        self.max_subseq_length = trans_subseq_length
                        
    def new_align(self, src, tgt, srcl, tgtl, steps=3):
       
        # Use tokenization if the input is a list, otherwise tokenize ourselves using spaces
        def to_list(data):
            if isinstance(data, str): return data.split()
            if isinstance(data, list): return data.copy()
            raise TypeError(f"Expected string or list, got {type(data)}")

        source_words = to_list(src)
        source_lemmas = to_list(srcl)
        target_words = to_list(tgt)
        target_lemmas = to_list(tgtl)
    
     
        aligns = AlignlistObj([], source_words, target_words)
    
        unal_source_indices = list(range(len(source_words)))
        unal_target_indices = list(range(len(target_words)))
    
        if steps == 4:
            intersectmwe_pass(self.source_language, self.target_language, 
                                  source_words, target_words, 
                                  source_lemmas, target_lemmas,
                                  aligns, 
                                  unal_source_indices, unal_target_indices, 
                                  False, self, self.src_join_char, self.tgt_join_char, 
                                  self.max_subseq_length)
              
        aligns.assert_no_conflicts()
        if steps > 0: 
            intersection_pass(self.source_language, self.target_language, 
                              source_words, target_words, 
                              source_lemmas, target_lemmas,
                              aligns, 
                              unal_source_indices, unal_target_indices, 
                              False, self, self.src_join_char, self.tgt_join_char)
        
        aligns.assert_no_conflicts()
        
        if steps > 1:
            babelmwe_pass(self.source_language, self.target_language, 
                          source_words, target_words, 
                          source_lemmas, target_lemmas,
                          aligns, 
                          unal_source_indices, unal_target_indices, 
                          False, self, self.src_join_char, self.tgt_join_char, 
                          self.max_subseq_length, True)
            
            babelnet_pass(self.source_language, self.target_language, 
                          source_words, target_words, 
                          source_lemmas, target_lemmas,
                          aligns, 
                          unal_source_indices, unal_target_indices, 
                          False, self, self.src_join_char, self.tgt_join_char, 
                          False)
        
        aligns.assert_no_conflicts()
        
        if steps > 2:
            simalign_pass(self.source_language, self.target_language, 
                          source_words, target_words, 
                          aligns, unal_source_indices, unal_target_indices, self)   

        aligns.assert_no_conflicts()
       
        answer = aligns.string_version()
        
        return answer
                
    
    
def sort_by_first_number(s: str) -> str:
    groups = s.split()  # split by spaces
    # sort by the first number in each group
    groups.sort(key=lambda g: int(g.split('-')[0]))
    return ' '.join(groups) 
    
class AlignlistObj:
  def __init__(self, list_start, source, target):
      self.source = source
      self.target = target
      self.alignments = list_start
      
  def _get_spans(self, pair):
        if len(pair) == 4:
            return pair # Already a quadrad
        
        p0, p1 = pair[0], pair[1]
        s_start = p0 if isinstance(p0, int) else min(p0)
        s_end   = p0 if isinstance(p0, int) else max(p0)
        t_start = p1 if isinstance(p1, int) else min(p1)
        t_end   = p1 if isinstance(p1, int) else max(p1)
        
        return s_start, s_end, t_start, t_end
    
  def conflict(self, quadrad):
      # Check if a span conflicts with the current alignments
      if len(quadrad) == 4:
        s_start, s_end, t_start, t_end = quadrad
      else:
          assert len(quadrad) == 2
          s_start, s_end, t_start, t_end = quadrad[0], quadrad[0], quadrad[1], quadrad[1]
      for other_quadrad in self.alignments:
         
          if other_quadrad[0] <= s_start <= other_quadrad[1] or other_quadrad[0] <= s_end <= other_quadrad[1]:
              return True
          if s_start <= other_quadrad[0] <= s_end or s_start <= other_quadrad[1] <= s_end:
              return True
          if other_quadrad[2] <= t_start <= other_quadrad[3] or other_quadrad[2] <= t_end <= other_quadrad[3]:
              return True
          if t_start <= other_quadrad[2] <= t_end or t_start <= other_quadrad[3] <= t_end:
              return True
      return False
  
  def assert_no_conflicts(self):
      
      
      for quadrad in self.alignments:
          
          
          s_start, s_end, t_start, t_end = self._get_spans(quadrad)
          for other_quadrad in self.alignments:
              if other_quadrad != quadrad:
                  if other_quadrad[0] <= s_start <= other_quadrad[1] or other_quadrad[0] <= s_end <= other_quadrad[1]:
                    assert False
                  if s_start <= other_quadrad[0] <= s_end or s_start <= other_quadrad[1] <= s_end:
                    assert False
                  if other_quadrad[2] <= t_start <= other_quadrad[3] or other_quadrad[2] <= t_end <= other_quadrad[3]:
                    assert False
                  if t_start <= other_quadrad[2] <= t_end or t_start <= other_quadrad[3] <= t_end:
                    assert False
              
  
  
  def string_of_pairs(self, src_tok, tgt_tok):
      ans = ''
      for guy in convert_quad_to_tuple(self.string_version()):
          ans += src_tok[guy[0]] + ' ' + tgt_tok[guy[1]] + '\n'
      return ans
      
        
  def print_this(self, pair):
      # To print the proper span formatting in a verbose way
      ans = ''
      if isinstance(pair[0], int) and isinstance(pair[1], int):
          s_start, s_end, t_start, t_end = pair[0], pair[0], pair[1], pair[1]
        
      elif isinstance(pair[0], int) and isinstance(pair[1], list):
            s_start, s_end = pair[0], pair[0]
            t_start, t_end = min(pair[1]), max(pair[1])
      elif isinstance(pair[0], list) and isinstance(pair[1], int):
            s_start, s_end = min(pair[0]), max(pair[0])
            t_start, t_end = pair[1], pair[1]
      elif isinstance(pair[0], list) and isinstance(pair[1], list):
          s_start, s_end = min(pair[0]), max(pair[0])
          t_start, t_end = min(pair[1]), max(pair[1])
      else:
          
          assert False
      for i in range(s_start, s_end + 1):
        ans += self.source[i] + ' '
        
      ans += ' ALIGNED WITH'
      for j in range(t_start, t_end + 1):
        ans += ' ' + self.target[j] 
      
    
  def add(self, pair):
      self.alignments.append(self._get_spans(pair))
      
  
  def string_version(self):
      # Get string representaiton
      ans = ''
      for guy in self.alignments:
          assert len(guy) == 4
          ans += str(guy[0]) + '-' + str(guy[1]) + '-' + str(guy[2]) + '-' + str(guy[3]) + ' '
      return sort_by_first_number(ans)
          
        
      
    
def read_function_words(file):
        engans = set()
        with open(file, mode='r', encoding='utf-8') as reader:
            for a, line in enumerate(reader):
                if a > 3:
                    if len(line.split()) < 1:
                        break
                    word = line.split()[0].strip()
                    engans.add(word.replace('"', ''))
        engans.add("'s")
        return engans
    
ENGLISH_FUNCTION_WORDS = read_function_words('res/WList.pm')

def get_lemmatizer(lang: str):
    if lang not in _LEMMATIZER_MODELS:
        raise KeyError(f"No lemmatizer configured for language '{lang}'")

    if lang not in _LEMMATIZER_CACHE:
        model_name = _LEMMATIZER_MODELS[lang]
        try:
            _LEMMATIZER_CACHE[lang] = spacy.load(model_name)
        except OSError as e:
            raise RuntimeError(
                f"spaCy model '{model_name}' is not installed.\n"
                f"Install it with:\n"
                f"  python -m spacy download {model_name}"
            ) from e

    return _LEMMATIZER_CACHE[lang]

        
def get_lemma(word, lang):
    
    cachekey = word + '<>' + lang
    if cachekey in STEM_CACHE:
        return STEM_CACHE[cachekey]
    
    if not word.strip():
        return word
    if lang == 'zh':
        return word
    
    if lang in _LEMMATIZER_MODELS:
        proper_nlp = get_lemmatizer(lang)
    else:
        proper_nlp = get_lemmatizer('xx')
    
    doc1 = proper_nlp(word)

    STEM_CACHE[cachekey] = doc1[0].lemma_
    return doc1[0].lemma_

def erase_cache():
    SYNSET_CACHED_DICT.clear()
    ALIGNMENT_CACHE.clear()
    STEM_CACHE.clear()
    POS_CACHE.clear()
    _LEMMATIZER_CACHE.clear()
    
    
   
def get_synsets_cachable(word, lang, pos=None):
    # Obtain synsets from babelnet
    if word in DO_NOT_CALL_BABELNET:
        return set()
    import babelnet as bn
    if word.startswith('bn:'):
        if word in SYNSET_CACHED_DICT:
            return SYNSET_CACHED_DICT[word]
        SYNSET_CACHED_DICT[word] = bn.get_synset(bn.BabelSynsetID(word))
      
        return SYNSET_CACHED_DICT[word]
    global POS_TAGS
    global LANGS
    if pos is not None and pos in POS_TAGS:
        index = word + lang + pos
        if index in SYNSET_CACHED_DICT:
            return SYNSET_CACHED_DICT[index]
        real_pos = POS_TAGS[pos]
        SYNSET_CACHED_DICT[index] = {s.id for s in bn.get_synsets(word, from_langs=[LANGS[lang]], poses=[real_pos])}
    
        return SYNSET_CACHED_DICT[index]
    else:
        index = word + lang + 'nopos'
        if index in SYNSET_CACHED_DICT:
            return SYNSET_CACHED_DICT[index]
      
        SYNSET_CACHED_DICT[index] = {s.id for s in bn.get_synsets(word, from_langs=[LANGS[lang]])}
       
        return SYNSET_CACHED_DICT[index]
    

def are_synonyms_by_bn(word1, word2, lang2, j_c, lang1):
    
    word1 = word1.replace(' ', j_c)
    word2 = word2.replace(' ', j_c)
    
    # certain tokens crash BN
    if word1 in DO_NOT_CALL_BABELNET or word2 in DO_NOT_CALL_BABELNET:
        return False
    
    for element in DO_NOT_CALL_BABELNET:
        if element in word1 or element in word2:
            if element != j_c:
                return False
            
    # Function words shouldn't be lemmatized.
    
    if word1 in ENGLISH_FUNCTION_WORDS and lang1 == 'en':
        lemma1 = word1
        lemma2 = word2
    elif j_c in word2 or j_c in word1:
        for guy in word1.split(j_c) + word2.split(j_c):
            
                if guy in PUNCTUATION:
                    return False
                
       
        lemma1 = j_c.join([get_lemma(a, lang1) for a in word1.split(j_c)])
        lemma2 = j_c.join([get_lemma(a, lang2) for a in word2.split(j_c)])
    else:
        lemma1 = get_lemma(word1, lang1)
        lemma2 = get_lemma(word2, lang2)
    
   
    synsets1 = get_synsets_cachable(word1, lang1)
    synsets2 = get_synsets_cachable(word2, lang2)
    synsets3 = get_synsets_cachable(lemma1, lang1)
    synsets4 = get_synsets_cachable(lemma2, lang2)
    
    
    word_intersection = set.intersection(synsets1, synsets2)
    lem_intersection = set.intersection(synsets3, synsets4)
    
    
    if word_intersection:
        return 'strict' # If the words are synonyms we say they strictly match
    elif lem_intersection:
        return 'loose' # If just their lemmas are we say they loosely match.

    return False        
 

    
def simalign_cachable(source, target):
    key = source + '<and>' + target
    if key in ALIGNMENT_CACHE:
        pass
    else:
        ALIGNMENT_CACHE[key] = SIMALIGNER.get_word_aligns(source, target)
    return ALIGNMENT_CACHE[key]['itermax']

def token_pos_and_morph_tag(sentence, language):
    
    tokens = []
    poses = []
    morphs = []
    
    key = language + '<>' + sentence + '<>' + 'WITHMORPH'
    if key in POS_CACHE:
        return POS_CACHE[key]
    if language in _LEMMATIZER_MODELS:
        doc = get_lemmatizer(language)(sentence)
    else:
        doc = get_lemmatizer('xx')(sentence)
   
    for token in doc:
        tokens.append(token.text)
        poses.append(token.pos_)
        morphs.append(token.morph.to_dict())
        
    ans = (tokens, poses, morphs)
    
    POS_CACHE[key] = ans
    
    return POS_CACHE[key]

def can_have_same_pos(sl, tl, word1, word2, p1, p2, alignobj):
    if alignobj.using_custom():
        return False
   
    synsets1 = get_synsets_cachable(word1, sl)
    synsets2 = get_synsets_cachable(word2, tl).union(get_synsets_cachable(get_lemma(word2, tl), tl)).union(TAGDICT[p2])
    
    poses1 = {TAGDICT[p1]}
    for synset in synsets1:
        poses1.add(str(synset)[-1])
    for synset in synsets2:
        letter = str(synset)[-1]
        if letter in poses1:
            return True
    
    return False
    
def pos_match(src_lang, tgt_lang, pos1, pos2, w1, w2, alignobj):
    
    if pos1 == pos2:
        return True
    if w1[:5] == w2[:5]:
        return True
    if can_have_same_pos(src_lang, tgt_lang, w1, w2, pos1, pos2, alignobj):
        return True
    if pos1 in ['VERB', 'AUX'] and pos2 in ['VERB', 'AUX']: # because this is basically the same
        return True
    if pos1 in ['PRON', 'DET'] and pos2 in ['PRON', 'DET']: # in the case of 'its' can be the same
        return True
    if pos1 in ['VERB', 'ADJ'] and pos2 in ['VERB', 'ADJ']: # in the case of 'we are surrounded' this can be the same
        return True
    if pos1 in ['PART', 'ADP', 'SCONJ'] and pos2 in ['PART', 'ADP', 'SCONJ']: # basically the same, no?
        return True
    if pos1 in ['PROPN', 'NOUN'] and pos2 in ['PROPN', 'NOUN']: # proper nouns are nouns i think this works:
        return True
    if pos1 in ['ADJ', 'DET'] and pos2 in ['ADJ', 'DET']: # in the context of 'other people'
        return True
    if pos1 in ['NOUN', 'VERB'] and pos2 in ['NOUN', 'VERB']: # really losing the plot of this function making any sense, but this is needed sometimes
        return True
    if alignobj.are_synonyms_by_dictionary(w1, w2, w1, w2, tgt_lang, src_lang):
        return True
    return False

def screen_alignments(a, s, t, len1, len2, slan, tlan, al_obj):
    
    answer = []
    t1, p1, _ = token_pos_and_morph_tag(s, slan)
    t2, p2, _ = token_pos_and_morph_tag(t, tlan)
    if not (len(t1) == len1 and len(t2) == len2):
        return a    
    for pair in a:
        try:
            if pos_match(slan, tlan, p1[pair[0]], p2[pair[1]], t1[pair[0]], t2[pair[1]], al_obj):
                answer.append(pair)
        except IndexError:
            return a
        
    return answer
    
def simalign_cachable_screened(source, target, len1, len2, src_lan, tar_lan, aligner):
    alignments = simalign_cachable(source, target)
    return screen_alignments(alignments, source, target, len1, len2, src_lan, tar_lan, aligner)

def add_nonvetoed_alignment(a, s, sentone, senttwo):
    if not veto(s[0], s[1], sentone, senttwo):
        a.add(s)

def add_alignment(a, s):
    a.add(s)
    
def mean(list_of_vals):
    return sum(list_of_vals) / len(list_of_vals)
    
def choose_diagonally(possible_inds, index, known_index_length, alt_index_length):
    if len(possible_inds) == 1:
        return index, possible_inds[0]
   
   
    flat_possible_indices = []
    list_possibilities = []
    for item in possible_inds:
        if isinstance(item, int):
            flat_possible_indices.append(item)
        elif isinstance(item, list):
            list_possibilities.append(item)
            for guy in item:
                flat_possible_indices.append(guy)
    
  
    if isinstance(index, list):
        position = mean(index)
    else:
        position = index
    src_distance_through = position / known_index_length
    tgt_distances_through = [a / alt_index_length for a in flat_possible_indices]
    
    
    best_closeness = 99999
    best_ans = None
    for i, tgt_dist in enumerate(tgt_distances_through):
        closeness = abs(tgt_dist - src_distance_through)
        if closeness < best_closeness:
            best_closeness = closeness
            best_ans = flat_possible_indices[i]
            
    for guy in list_possibilities:
        if best_ans in guy:
           
            return index, best_ans
    return index, best_ans
    
    
def mwe_intersectalign(index, simals, pis, iwk):
    
    good = True
    for pi in pis + [index]:
        if isinstance(pi, list):
            good = False
    if isinstance(index, int):
        s_side_inds_flat = [index]
    else:
        s_side_inds_flat = index
        
    if good:
        
        return [a for a in simals if a[iwk] == index and a[1 - iwk] in pis]
    else:
        ans = []
        for pi in pis:
            this_pi_valid = True
            if isinstance(pi, int):
                for guy in s_side_inds_flat:
                    if (guy, pi) in simals:
                        pass
                    else:
                        this_pi_valid = False
                
            elif isinstance(pi, list):
                for x in pi:
                    for guy in s_side_inds_flat:
                        if (x, guy) in simals:
                            pass
                        else:
                            this_pi_valid = False
            if this_pi_valid:
                    ans.append((index, pi))
                        
        return ans
     
        
    
def choose_alignment_from_bn_options(poss_indices, index, src_wds, tgt_wds, intersect_pass, choosing):
    
    
    # We have multiple options according to BabelNet, choose them based on SimAlign info
    
    simaligns = simalign_cachable(' '.join(src_wds), ' '.join(tgt_wds))
    
    if choosing == 'source':
        relevant_wds = tgt_wds
        ind_we_know = 1
        ind_we_dont = 0
        other_wds = src_wds
    elif choosing == 'target':
        relevant_wds = src_wds
        other_wds = tgt_wds
        ind_we_know = 0
        ind_we_dont = 1
   
    intersectaligns_of_token = mwe_intersectalign(index, simaligns, poss_indices, ind_we_know) # intersection of alignments
        
    intersectaligns_of_token = [a for a in simaligns if a[ind_we_know] == index and a[ind_we_dont] in poss_indices]
    
    
    if len(intersectaligns_of_token) == 1:
           
        return intersectaligns_of_token[0]
    elif len(intersectaligns_of_token) > 1:
              
        confirmed_poss_indices = [a[ind_we_dont] for a in intersectaligns_of_token]
        
        most_diagonal_alignment = choose_diagonally(confirmed_poss_indices, index, len(relevant_wds), len(other_wds))
     
        return most_diagonal_alignment[ind_we_know], most_diagonal_alignment[ind_we_dont]
    elif intersect_pass:
        return None, None
    elif len(intersectaligns_of_token) == 0 and len(poss_indices) > 0:
        
        # none of the alignments are supported by simalign, but let's choose one by virtue of diagonality
        most_diagonal_alignment = choose_diagonally(poss_indices, index, len(relevant_wds), len(other_wds))
      
        return most_diagonal_alignment[ind_we_know], most_diagonal_alignment[ind_we_dont]
    
    
    return None, None
    
    
def find_diag_score(align, len_one, len_two):
    first_dist = align[0] / len_one
    second_dist = align[1] / len_two
    return abs(first_dist - second_dist)
    
def sorted_by_diagonal_closeness(t, l1, l2):
    return sorted(t, key=lambda p: find_diag_score(p, l1, l2))
    
    
def no_conflict(l, curr):
    for other in l:
        if other[0] == curr[0] or other[1] == curr[1]:
            return False
    return True
    
def diagonally_restricted(t_a, length_one, length_two):
    
    ans = []
    
    for entry in sorted_by_diagonal_closeness(t_a, length_one, length_two):
        if no_conflict(ans, entry):
            ans.append(entry)
        
        
    return ans

def accept_all_alignments(proposed, alignments, usi, uti, toks1, toks2):
    unfinished = False
    to_add = []
    for pair in proposed:
        if not alignments.conflict(pair):
            to_add.append(pair)
            
    to_add = diagonally_restricted(to_add, len(toks1), len(toks2))
    for pair in to_add:
        add_nonvetoed_alignment(alignments, pair, toks1, toks2)
        safe_remove(usi, pair[0])
        safe_remove(uti, pair[1])
          
    
    return unfinished

def safe_remove(s, element):
    if isinstance(element, int):
        
        if element in s:
            s.remove(element)
    elif isinstance(element, list):
      
        for guy in element:
        
            s.remove(guy)
        
       
    
def flatten(mixed_list):
    result = []
    for item in mixed_list:
        if isinstance(item, list):
            result.extend(item)
        else:
            result.append(item)
    return result


    
def get_claimed(mylist):
    ans_s, ans_t = [], []
   
    for guy in mylist:
      
        assert isinstance(guy, tuple) and len(guy) == 2
        if isinstance(guy[0], int) and isinstance(guy[1], int):
            ans_s.append(guy[0])
            ans_t.append(guy[1])
        else:
            if isinstance(guy[0], list):
                for fella in guy[0]:
                    ans_s.append(fella)
            if isinstance(guy[1], list):
                for fella in guy[1]:
                    ans_t.append(fella) 
            if isinstance(guy[0], int):
                ans_s.append(guy[0])
            if isinstance(guy[1], int):
                ans_t.append(guy[1])
    return ans_s, ans_t
  
def claimed_by(arg):
    claimed = []
    if isinstance(arg, int):
        claimed.append(arg)
    elif isinstance(arg, list):
        for guy in arg:
            claimed.append(guy)  
    return claimed
   
def claims(p, s, t):

    if isinstance(p, tuple) and len(p) == 2:
        if isinstance(p[0], int) and isinstance(p[1], int):

            return p[0] in s or p[1] in t
        else:
            claimed_by_first = claimed_by(p[0])
            for guy in claimed_by_first:
                if guy in s:
                    return True
            claimed_by_second = claimed_by(p[1])
            for guy in claimed_by_second:
                if guy in t:
                    return True
            return False
    else:  
        assert False 
    
def subsumes(other_match, match):
    
    if isinstance(other_match[0], int):
        o_claim_s = [other_match[0]]
    else:
        o_claim_s = other_match[0]
        
    if isinstance(other_match[1], int):
        o_claim_t = [other_match[1]]
    else:
        o_claim_t = other_match[1]
        
    if isinstance(match[1], int):
        claim_t = [match[1]]
    else:
        claim_t = match[1]
    
    if isinstance(match[0], int):
        claim_s = [match[0]]
    else:
        claim_s = match[0]
    for guy in claim_s:
        if guy not in o_claim_s:
         
            return False
    for guy in claim_t:
        if guy not in o_claim_t:
            
            return False
    
    return True
    
    
def generalize_if_possible(proposed):
    '''If one proposed alignment strictly contains another, use that one'''
    for match in proposed:
        for other_match in proposed:
            if other_match != match and subsumes(other_match, match):
                proposed.remove(match)
                break
    return proposed
    
    
def accept_unconflicting_alignments(proposed, alignments, usi, uti, toks1, toks2):
    
    proposed = generalize_if_possible(proposed)
       
    unfinished = False
    claimed_src_indices, claimed_tar_indices = get_claimed(proposed)
    
    conflicted_tgt_indices = [i for i in set(claimed_tar_indices) if claimed_tar_indices.count(i) > 1]
    conflicted_src_indices = [i for i in set(claimed_src_indices) if claimed_src_indices.count(i) > 1]
   
    for pair in proposed:
       
        if not claims(pair, conflicted_src_indices, conflicted_tgt_indices):
            add_nonvetoed_alignment(alignments, pair, toks1, toks2)
            
            safe_remove(usi, pair[0])
            safe_remove(uti, pair[1])
        else:
                       
            unfinished = True
   
    return unfinished
    
def consecutive_subsequences(nums, max_len): 
   

    if not nums:
        return []

    nums = sorted(set(nums))  # Sort and remove duplicates
    subsequences = []
    start = 0

    for i in range(1, len(nums) + 1):
        # If end of list or break in adjacency
        if i == len(nums) or nums[i] != nums[i - 1] + 1:
            segment = nums[start:i]
            # Generate all subsequences of length >= 2 from the segment
            for j in range(len(segment)):
                for k in range(j + 2, len(segment) + 1):
                    
                    subseq = segment[j:k]
                    if max_len is None or len(subseq) <= max_len:
                        subsequences.append(subseq)
            start = i
   
    return subsequences


    
def intersection_pass(src_lang, tar_lang, src_wds, tgt_wds, src_lems, tgt_lems, align_ans, unaligned_source_indices, unaligned_target_indices, strict_lemma, aligner, sjc, tjc):
    babelnet_pass(src_lang, tar_lang, src_wds, tgt_wds, src_lems, tgt_lems, align_ans, unaligned_source_indices, unaligned_target_indices, strict_lemma, aligner, sjc, tjc, True)
    
    
def lemmatize(list_of_words, language):
    return [get_lemma(a, language) for a in list_of_words]
    
    
def babelmwe_pass(sl, tl, src_wds, tgt_wds, src_lms, tgt_lms, align_ans, unaligned_source_indices, unaligned_target_indices, strict_lemma, alignobj, s_j_c, t_j_c, max_subseq_l, strict_intersect):
  unfinished = True
  
  if src_lms is None:
        src_lms = lemmatize(src_wds, sl)
  if tgt_lms is None:
        tgt_lms = lemmatize(tgt_wds, tl)
      
  # Repeat until nothing changes
  while unfinished:
    # first pass using babelnet align
    # Forward pass : align the source words
    proposed_alignments = []
    
    # Loop over all unaligned indices and contiguous subsequences of them
    for i in unaligned_source_indices + consecutive_subsequences(unaligned_source_indices, max_subseq_l):
        if isinstance(i, int):
            word_one = src_wds[i]
            lemma_one = src_lms[i]
        elif isinstance(i, list):
            word_one = s_j_c.join([src_wds[el] for el in i])
            lemma_one = s_j_c.join([src_lms[el] for el in i])

       
        
        possible_alignment_indices_for_this = []
        less_strict_possible_alignment_indices_for_this = []
        for j in unaligned_target_indices + consecutive_subsequences(unaligned_target_indices, max_subseq_l):
            
            # Since this is specifically the MWE path, don't just align lone tokens to each other.
            if isinstance(i, int) and isinstance(j, int):
                continue
            
            # Find the target side word:
            if isinstance(j, int):
                word_two = tgt_wds[j]
                lemma_two = tgt_lms[j]
            elif isinstance(j, list):
                word_two = t_j_c.join([tgt_wds[el] for el in j])
                lemma_two = t_j_c.join([tgt_lms[el] for el in j])
              
          
            # Check for synonyms
            are_syn = alignobj.are_synonyms_by_dictionary(word_one, word_two, lemma_one, lemma_two, tl, sl)
            
            if are_syn == 'strict': # this is true if theyre the same without lemmatizing them
                
                possible_alignment_indices_for_this.append(j)
                if isinstance(j, int):
                    assert j in unaligned_target_indices
                elif isinstance(j, list):
                    for guy in j:
                        assert guy in unaligned_target_indices
            elif are_syn == 'loose': # if they're the same only when lemmatized
                
                less_strict_possible_alignment_indices_for_this.append(j)
           
            
           
        
            
        for x in possible_alignment_indices_for_this:
            if isinstance(x, int):
            
                assert x in unaligned_target_indices
            else:
                for sub_x in x:
                    assert sub_x in unaligned_target_indices
        
        # we prefer to use the strict ones if possible, but if there aren't any, use either
        if len(possible_alignment_indices_for_this) == 0 and not strict_lemma:
            possible_alignment_indices_for_this += less_strict_possible_alignment_indices_for_this
        
            
        
        new_al_src_ind, new_al_tgt_ind = choose_alignment_from_bn_options(possible_alignment_indices_for_this, i, src_wds, tgt_wds, strict_intersect, 'target')   
        
        # if we made a new alignment, remove these from the list of indices that can be aligned still
        if new_al_src_ind is not None and new_al_tgt_ind is not None:
                    
            proposed_alignments.append((new_al_src_ind, new_al_tgt_ind))
             
    unfinished = accept_unconflicting_alignments(proposed_alignments, align_ans, unaligned_source_indices, unaligned_target_indices, src_wds, tgt_wds)
    proposed_alignments = []
    
    # backwards pass: align target words   
    
    for j in unaligned_target_indices + consecutive_subsequences(unaligned_target_indices, max_subseq_l):
        
        if isinstance(j, int):
                
                word_two = tgt_wds[j]
                lemma_two = tgt_lms[j]
        elif isinstance(j, list):
                word_two = t_j_c.join([tgt_wds[el] for el in j])
                lemma_two = t_j_c.join([tgt_lms[el] for el in j])
   
        possible_alignment_indices_for_this = []
        less_strict_possible_alignment_indices_for_this = []
        for i in unaligned_source_indices + consecutive_subsequences(unaligned_source_indices, max_subseq_l):
            if isinstance(i, int) and isinstance(j, int):
                continue
            if isinstance(i, int):
                word_one = src_wds[i]
                lemma_one = src_lms[i]
            elif isinstance(i, list):
                word_one = s_j_c.join([src_wds[el] for el in i])
                lemma_one = s_j_c.join([src_lms[el] for el in i])
           
            are_syn = alignobj.are_synonyms_by_dictionary(word_one, word_two, lemma_one, lemma_two, tl, sl)
            if are_syn == 'strict': # this is true if theyre the same without lemmatizing them
                possible_alignment_indices_for_this.append(i)
                
                if isinstance(i, int):
                    assert i in unaligned_source_indices
                elif isinstance(i, list):
                    for guy in i:
                        assert guy in unaligned_source_indices
            elif are_syn == 'loose': # if they're the same only when lemmatized
                less_strict_possible_alignment_indices_for_this.append(i)
        for x in possible_alignment_indices_for_this:
            if isinstance(x, int):
                assert x in unaligned_source_indices
            elif isinstance(x, list):
                for guy in x:
                    assert guy in unaligned_source_indices
        
        # we prefer to use the strict ones if possible, but if there aren't any, use either
        if len(possible_alignment_indices_for_this) == 0 and not strict_lemma:
            possible_alignment_indices_for_this += less_strict_possible_alignment_indices_for_this
        
            
     
        new_al_src_ind, new_al_tgt_ind = choose_alignment_from_bn_options(possible_alignment_indices_for_this, j, src_wds, tgt_wds, strict_intersect, 'source')   
       
        # if we made a new alignment, remove these from the list of indices that can be aligned still
        if new_al_src_ind is not None and new_al_tgt_ind is not None:
         
            assert new_al_src_ind in unaligned_source_indices or (isinstance(new_al_src_ind, list) and [d in unaligned_source_indices for d in unaligned_source_indices])
            proposed_alignments.append((new_al_src_ind, new_al_tgt_ind))
            
    unfinished = accept_unconflicting_alignments(proposed_alignments, align_ans, unaligned_source_indices, unaligned_target_indices, src_wds, tgt_wds)
    previous_cycle = proposed_alignments.copy()
    
    # If nothing has changed, we're done
    if proposed_alignments == previous_cycle:
        
        unfinished = False
  

   
  
  
                    
def babelnet_pass(sl, tl, src_wds, tgt_wds, src_lms, tgt_lms, align_ans, unaligned_source_indices, unaligned_target_indices, strict_lemma, alignobj, sjc, tjc, strict_intersect):
  unfinished = True
  
  # Repeat until nothing changes
  while unfinished:
    if src_lms is None:
        src_lms = lemmatize(src_wds, sl)
    if tgt_lms is None:
        tgt_lms = lemmatize(tgt_wds, tl)
    
    
    proposed_alignments = []
    for i in unaligned_source_indices.copy():
        if isinstance(i, int):
            word_one = src_wds[i]
            lemma_one = src_lms[i]
        elif isinstance(i, list):
            word_one = sjc.join([src_wds[el] for el in i])
            lemma_one = sjc.join([src_lms[el] for el in i])

        possible_alignment_indices_for_this = []
        less_strict_possible_alignment_indices_for_this = []
        
        for j in unaligned_target_indices:
            if isinstance(j, int):
                word_two = tgt_wds[j]
                lemma_two = tgt_lms[j]
                
        
            are_syn = alignobj.are_synonyms_by_dictionary(word_one, word_two, lemma_one, lemma_two, tl, sl)
            if are_syn == 'strict': # this is true if theyre the same without lemmatizing them
                possible_alignment_indices_for_this.append(j)
                if isinstance(j, int):
                    assert j in unaligned_target_indices
                elif isinstance(j, list):
                    for guy in j:
                        assert guy in unaligned_target_indices
            elif are_syn == 'loose': # if they're the same only when lemmatized
                less_strict_possible_alignment_indices_for_this.append(j)
        
            
        for x in possible_alignment_indices_for_this:
            if isinstance(x, int):
            
                assert x in unaligned_target_indices
            else:
                for sub_x in x:
                    assert sub_x in unaligned_target_indices
        
        # we prefer to use the strict ones if possible, but if there aren't any, use either
        if len(possible_alignment_indices_for_this) == 0 and not strict_lemma:
            possible_alignment_indices_for_this += less_strict_possible_alignment_indices_for_this
        
            
        
        new_al_src_ind, new_al_tgt_ind = choose_alignment_from_bn_options(possible_alignment_indices_for_this, i, src_wds, tgt_wds, strict_intersect, 'target')   
        
        # if we made a new alignment, remove these from the list of indices that can be aligned still
        if new_al_src_ind is not None and new_al_tgt_ind is not None:
            
            # Ensure formatting:
            if isinstance(new_al_tgt_ind, int):
                assert new_al_tgt_ind in unaligned_target_indices
                proposed_alignments.append((new_al_src_ind, new_al_tgt_ind))

            elif isinstance(new_al_tgt_ind, list):
                for k in new_al_tgt_ind:
                    proposed_alignments.append((new_al_src_ind, k))
                    
    # Accept all alignments that fit
    unfinished = accept_unconflicting_alignments(proposed_alignments, align_ans, unaligned_source_indices, unaligned_target_indices, src_wds, tgt_wds)
    proposed_alignments = []
    
    # backwards pass: align target words   
    
    for j in unaligned_target_indices.copy():
        
        if isinstance(j, int):
                word_two = tgt_wds[j]
                lemma_two = tgt_lms[j]
        elif isinstance(j, list):
                word_two = tjc.join([tgt_wds[el] for el in j])
                lemma_two = tjc.join([tgt_lms[el] for el in j])
   
        possible_alignment_indices_for_this = []
        less_strict_possible_alignment_indices_for_this = []
        for i in unaligned_source_indices:
            if isinstance(i, int):
                word_one = src_wds[i]
                lemma_one = src_lms[i]
            
            are_syn = alignobj.are_synonyms_by_dictionary(word_one, word_two, lemma_one, lemma_two, tl, sl)
            if are_syn == 'strict': # this is true if theyre the same without lemmatizing them
                possible_alignment_indices_for_this.append(i)
          
                if isinstance(i, int):
                    assert i in unaligned_source_indices
                
            elif are_syn == 'loose': # if they're the same only when lemmatized
                less_strict_possible_alignment_indices_for_this.append(i)
        for x in possible_alignment_indices_for_this:
            if isinstance(x, int):
                assert x in unaligned_source_indices
            elif isinstance(x, list):
                for guy in x:
                    assert guy in unaligned_source_indices
        
        # we prefer to use the strict ones if possible, but if there aren't any, use either
        if len(possible_alignment_indices_for_this) == 0 and not strict_lemma:
            possible_alignment_indices_for_this += less_strict_possible_alignment_indices_for_this
        
            
       
        new_al_src_ind, new_al_tgt_ind = choose_alignment_from_bn_options(possible_alignment_indices_for_this, j, src_wds, tgt_wds, strict_intersect, 'source')   
       
        # if we made a new alignment, remove these from the list of indices that can be aligned still
        if new_al_src_ind is not None and new_al_tgt_ind is not None:
            assert new_al_src_ind in unaligned_source_indices
            proposed_alignments.append((new_al_src_ind, new_al_tgt_ind))
      
    unfinished = accept_unconflicting_alignments(proposed_alignments, align_ans, unaligned_source_indices, unaligned_target_indices, src_wds, tgt_wds)
    previous_cycle = proposed_alignments.copy()
    if proposed_alignments == previous_cycle:
        # If nothing has changed, we're done!
        unfinished = False
  
  
    
def intersectmwe_pass(src_lang, tar_lang, src_wds, tgt_wds, src_lems, tgt_lems, align_ans, unaligned_source_indices, unaligned_target_indices, strict_lemma, aligner, s_j_c, t_j_c, max_s_l):
    babelmwe_pass(src_lang, tar_lang, src_wds, tgt_wds, src_lems, tgt_lems, align_ans, unaligned_source_indices, unaligned_target_indices, strict_lemma, aligner, s_j_c, t_j_c, max_s_l, False, False)
    babelmwe_pass(src_lang, tar_lang, src_wds, tgt_wds, src_lems, tgt_lems, align_ans, unaligned_source_indices, unaligned_target_indices, strict_lemma, aligner, s_j_c, t_j_c, max_s_l, False, True)
      
def are_aligned(ind_one, ind_two, alignments):
    for al in alignments:
        if al[0] == ind_one and al[1] == ind_two:
            return True
    return False
    

def cognates(word_one, word_two):
    return False
    if word_one[:3] == word_two[:3]:
        return True
    return False

def every_word_cognates(words_one, words_two):
    
    for w1 in words_one:
        is_cognate = False
        for w2 in words_two:
            if cognates(w1, w2):
                is_cognate = True
        if not is_cognate:
            return False
        
    for w2 in words_two:
        is_cognate = False
        for w1 in words_one:
            if cognates(w1, w2):
                is_cognate = True
        if not is_cognate:
            return False
    return True
        
def veto(i, j, src_ws, tgt_ws):
    if isinstance(i, int):
        w1 = src_ws[i]
    else:
        w1 = ' '.join([src_ws[k] for k in i])
  
    if isinstance(j, int):
        w2 = tgt_ws[j]
    else:
        w2 = ' '.join([tgt_ws[k] for k in j])
    
  
    if w1 in PUNCTUATION and w2 not in PUNCTUATION:
        return True
    if w2 in PUNCTUATION and w1 not in PUNCTUATION:
        return True
    if w1 in ['(', ')'] and w2 not in ['(', ')']:
        return True
    if w2 in ['(', ')'] and w1 not in ['(', ')']:
        return True
    return False
         
def simalign_pass(sl, tl, src_wds, tgt_wds, align_ans, unaligned_source_indices, unaligned_target_indices, a):
  unfinished = True
  previous_cycle = None
  
  # Loop until nothing changes
  while unfinished:
      
    
    proposed_alignments = []
    simaligns = simalign_cachable_screened(' '.join(src_wds), ' '.join(tgt_wds), len(src_wds), len(tgt_wds), sl, tl, a)
    
    # Iterate through proposed alignments
    for i in unaligned_source_indices.copy():
        possible_alignment_indices_for_this = []
        for j in unaligned_target_indices:

            # Add any that are valid
            if are_aligned(i, j, simaligns) and not veto(i, j, src_wds, tgt_wds):
                possible_alignment_indices_for_this.append(j)
                assert j in unaligned_target_indices
                
        # Ensure correct formatting
        for x in possible_alignment_indices_for_this:
            assert x in unaligned_target_indices
            
        # Use the diagonal heuristic to break ties
        most_diagonal = choose_diagonally(possible_alignment_indices_for_this, i, len(src_wds), len(tgt_wds))
        
        
        new_al_src_ind, new_al_tgt_ind = most_diagonal[0], most_diagonal[1]
        
        # if we made a new alignment, remove these from the list of indices that can be aligned still
        if new_al_src_ind is not None and new_al_tgt_ind is not None:
            
            assert new_al_tgt_ind in unaligned_target_indices
           
            proposed_alignments.append(most_diagonal)
    
    
    unfinished = accept_unconflicting_alignments(proposed_alignments, align_ans, unaligned_source_indices, unaligned_target_indices, src_wds, tgt_wds)
    previous_cycle = proposed_alignments.copy()
    
    # If nothing has changed, we're done.
    if proposed_alignments == previous_cycle:
    
        accept_all_alignments(proposed_alignments, align_ans, unaligned_source_indices, unaligned_target_indices, src_wds, tgt_wds)
        return