Chronofleks / korba-ameba

# -*- coding: utf-8 -*-
import os, io, argparse
from lxml import etree

# Lista frekwencyjna w trzech kolumnach rozdzielonych tabulacją: liczba wystąpień, lemat, pos
FREQLIST = '/home/dorota/projekty/korba-ameba/freq-nkjp1m.txt'

# Lista typów (elementów tekstu TEI), które zostaną pominięte w eksporcie do dag -
# wszystkie, które nie są interpretowane w anotacji bez 'foreign'.
OMITLIST = ['quote']

namespaces = {
  'tei': 'http://www.tei-c.org/ns/1.0',
  'xi': 'http://www.w3.org/2001/XInclude',
  'nkjp': 'http://www.nkjp.pl/ns/1.0',
  'xml': 'http://www.w3.org/XML/1998/namespace'
}

def new_segment(line_list):
  seg = {'start': int(line_list[0]), 'end': int(line_list[1]), 'orth': line_list[2], 'disamb': True if line_list[-1] == 'disamb' else False}
  if seg['disamb']:
    seg['interpretation'] = line_list[3] + ':' + line_list[4]
  return seg

def check_freq(last_lex, last_pos, new_lex, new_pos):
  found = 0
  last_f = 0
  new_f = 0
  for line in io.open(FREQLIST, 'r', encoding='utf-8'):
    f, lex, pos = line.strip().split('\t')
    if (last_lex == lex) and (last_pos == pos):
      last_f = int(f)
      found += 1
    if (new_lex == lex) and (new_pos == pos):
      new_f = int(f)
      found += 1
    if found == 2:
      break
  return last_f, new_f

# Przetwarza tekst po tagerze filtrując tylko interpretacje oznaczone disamb
# Uwaga na sytuacje, gdy jest kilka takich interpretacji dla segmentu - zdarza się
# to jeśli interpretacje różnią się tylko lematem a znaczniki są takie same.
def read_text(filename):
  file_in = io.open(filename, 'r', encoding='utf-8')
  seg_list = []
  newpara = False
  
  for line in file_in:
    line_list = line.split()
    if len(line_list) > 0:
      if line_list[-1] == 'disamb':
        if len(seg_list) and not newpara:
          last_seg = seg_list[-1]
          # Jeśli aktualnie czytany segment ma takie same 'start' i 'end', to sprawdź, który wybrać
          # Na razie - zignoruj
          if last_seg['start'] == int(line_list[0]) and last_seg['end'] == int(line_list[1]):
#            pass
#            print u"Kilka wybranych interpretacji: {}".format(line.strip()).encode('utf-8')
            last_f, new_f = check_freq(last_seg['interpretation'].split(':',2)[0], last_seg['interpretation'].split(':',2)[1], line_list[3], line_list[4].split(':',1)[0])
            if new_f > last_f:
              print u"zmiana dla segmentu {}: {} {}:{} -> {} {}:{}".format(line_list[2], last_f, last_seg['interpretation'].split(':',2)[0], last_seg['interpretation'].split(':',2)[1], new_f, line_list[3], line_list[4].split(':',1)[0]).encode('utf-8')
              seg_list.pop()
              seg_list.append(new_segment(line_list))
          else:
            seg_list.append(new_segment(line_list))
        # Jeszcze pusta lista segmentów lub był nowy akapit
        else:
          seg_list.append(new_segment(line_list))
          newpara = False
    # pusta linia w pliku - przestaw newpara
    else:
      newpara = True
      
  file_in.close()
  return seg_list

def print_dag_simple(seg_list):

  for seg in seg_list:
    print u"{}\t{}\t{}".format(seg['start'], seg['end'], seg['orth']).encode('utf-8')


# Przetwarza tekst po tagerze zbierając również choice'y
def process_text_choice(filename):
  file_in = io.open(filename, 'r', encoding='utf-8')
  seg_list = []
  newpara = False
  
  for line in file_in:
    line_list = line.split()
    if len(line_list) > 0:
      if len(seg_list) and not newpara:
        found = False
        for seg_var in seg_list[-1]:
          # Jeśli w którymś z wariantów na ostatnim miejscu jest segment o takim samym start i end - zaktualizuj go i zakończ
          if seg_var[-1]['start'] == int(line_list[0]) and seg_var[-1]['end'] == int(line_list[1]):
            if line_list[-1] == 'disamb':
              seg_var[-1]['disamb'] = True
              seg_var[-1]['interpretation'] = line_list[3] + ':' + line_list[4]
            found = True
            break
        if not found:
          # Jeśli segment zaczyna się jak inne warianty - dodaj nowy
          if seg_list[-1][0][0]['start'] == int(line_list[0]):
            seg_list[-1].append([new_segment(line_list)])
          # Segment zaczyna się inaczej - jeśli wszystkie warianty mają ten sam end - dołóż nowy segment do listy,
          # jeśli nie - znajdź wariant, do którego należy go dołożyć.
          else:
            endings = [seg_var[-1]['end'] for seg_var in seg_list[-1]]
            if len(set(endings)) == 1:
              if int(line_list[1]) > endings[0]:
                seg_list.append([[new_segment(line_list)]])
              else:
                print u'Kolejny segment kończy się wcześniej niż poprzednie: {}'.format(line.strip())
            else:
#              print 'endings: {}'.format(endings)
              seg_list[-1][endings.index(int(line_list[0]))].append(new_segment(line_list))
            
      # Jeszcze pusta lista segmentów lub był nowy akapit
      else:
        seg_list.append([[new_segment(line_list)]])
        newpara = False
    # pusta linia w pliku - przestaw newpara
    else:
      newpara = True

      
  file_in.close()
  return seg_list


def print_dag_segs(seg_list):
  seg_len = 0  

  for choice in seg_list:
    for paren in choice:
      seg_len += len(paren)

  print u"Liczba segmentów: {}".format(seg_len).encode('utf-8')

  for var in seg_list:
    print var
  
def test_dag(seg_list, filename):
  namespaces = {
    'tei': 'http://www.tei-c.org/ns/1.0',
    'xi': 'http://www.w3.org/2001/XInclude',
    'nkjp': 'http://www.nkjp.pl/ns/1.0',
    'xml': 'http://www.w3.org/XML/1998/namespace'
  }

  directory = os.path.dirname(filename)
  seg_tree = etree.parse(os.path.join(directory, 'ann_segmentation.xml'))
  sents = seg_tree.xpath("//tei:s", namespaces=namespaces)
  tree_segs = []
  for s in sents:
    tree_segs.extend(s.getchildren())
  if len(tree_segs) == len(seg_list):
    print u">>>Folder {} - ta sama liczba elementów zdania".format(directory).encode("utf-8")
  else:
    print u"<<<Folder {} - różna liczba elementów zdania".format(directory).encode("utf-8")

  for ind, tree_s in enumerate(tree_segs, start=0):
    if ind >= len(seg_list):
      print u"Plik DAG ma za mało segmentów - brak odpowiednika dla id={}".format(tree_s.get('{http://www.w3.org/XML/1998/namespace}id')).encode("utf-8")
    else:
      if tree_s.tag == '{http://www.tei-c.org/ns/1.0}seg':
        if len(seg_list[ind]) != 1:
#          print u"Podział na choice w DAG dla pojedynczego segmentu o id={}".format(tree_s.get('{http://www.w3.org/XML/1998/namespace}id')).encode("utf-8")
          print u"{} - {}".format(etree.tostring(tree_s, method='text', encoding='unicode'), seg_list[ind]).encode("utf-8")
        else:
#          print u"{}".format(tree_s.xpath('.//tei:w/text()', namespaces = namespaces)).encode("utf-8")
          if len(seg_list[ind][0]) != 1:
            print u"Dla segmentu o id={} jest kilka segmentów w DAG".format(tree_s.get('{http://www.w3.org/XML/1998/namespace}id')).encode("utf-8")
#          else: # w ann_segmentation jest wersja transliterowana - nie można tego sprawdzić
#            if tree_s.xpath('.//tei:w/text()', namespaces = namespaces)[0] != seg_list[ind][0][0]['orth']:
#              print u"Dla segmentu o id={} jest inny napis w DAG".format(tree_s.get('{http://www.w3.org/XML/1998/namespace}id')).encode("utf-8")
      elif tree_s.tag == '{http://www.tei-c.org/ns/1.0}choice':
        tree_choice = [len(elem.xpath('descendant-or-self::tei:seg', namespaces=namespaces)) for elem in tree_s]
        tree_choice.sort()
        seg_list_choice = [len(choice) for choice in seg_list[ind]]
        seg_list_choice.sort()
#        print "ind: {}, {} - {}".format(ind, tree_choice, seg_list_choice)
        if tree_choice != seg_list_choice:
#          print u"Niezgodność w DAG dla segmentu z choice od id={}".format(tree_s.xpath('descendant-or-self::tei:seg', namespaces=namespaces)[0].get('{http://www.w3.org/XML/1998/namespace}id')).encode("utf-8")
          print u"{} - {}".format(etree.tostring(tree_s, method='text', encoding='unicode'), seg_list[ind]).encode("utf-8")
      else:
        print u"Nietypowy element w zdaniu: {}".format(tree_s).encode("utf-8")

def find_morph(seg_id, morph_index, morph_list):
  m_index = morph_index
  while (m_index < len(morph_list)) and (morph_list[m_index].get('corresp','') != "ann_segmentation.xml#{}".format(seg_id)):
    m_index += 1
  if m_index == len(morph_list):
    return (morph_index, []) 
  else:
    return (m_index, [morph_list[m_index]])

def tag_wo_ns(elem):
  return elem.tag.rsplit("}", 1)[-1]

def print_interps(dag_out, start, end, morph_seg, seg_type, disamb):
  orth = morph_seg.xpath('./tei:fs/tei:f[@name="orth"]/tei:string', namespaces = namespaces)[0].text
  if seg_type == 'foreign':
    if disamb:
      dag_out.write(u"{}\t{}\t{}\t{}\t{}\t1.000\tdisamb\n".format(start, end, orth, orth, 'xxx'))
    else:
      dag_out.write(u"{}\t{}\t{}\t{}\t{}\t0.000\n".format(start, end, orth, orth, 'xxx'))
  else:
    if disamb:
      disamb_elems = morph_seg.xpath('./tei:fs/tei:f[@name="disamb"]/tei:fs[@type="tool_report"]/tei:f[@name="choice"]', namespaces = namespaces)
      if len(disamb_elems) == 0:
        print u"Brak ujednoznacznienia dla segmentu: {}".format(orth).encode("utf-8")
        return
      disamb_ids = [elem.get('fVal')[1:] for elem in disamb_elems] # trzeba powinąć początkowy # w fVal
    for lex in morph_seg.xpath('./tei:fs/tei:f[@name="interps"]/tei:fs[@type="lex"]', namespaces = namespaces):
      base = lex.xpath('./tei:f[@name="base"]/tei:string', namespaces = namespaces)[0].text
      ctag = lex.xpath('./tei:f[@name="ctag"]/tei:symbol/@value', namespaces = namespaces)[0]
      msds = lex.xpath('./tei:f[@name="msd"]//tei:symbol', namespaces = namespaces)
      for m in msds:
        msd =  m.get('value')
        msd_id = m.get('{http://www.w3.org/XML/1998/namespace}id')
        msd = ctag + ':' + msd if msd else ctag
        dag_out.write(u"{}\t{}\t{}\t{}\t{}".format(start, end, orth, base, msd))
        if disamb:
          if msd_id in disamb_ids:
            dag_out.write(u"\t{0:.3f}\tdisamb\n".format(1.0/len(disamb_elems)))
          else:
            dag_out.write(u"\t0.000\n")
        else:
          dag_out.write(u"\t0.000\n")


def make_dag(directory, disamb):
  dag_out = io.open(os.path.join(directory, 'text_transcr.txt.morph'), 'w', encoding='utf-8')
  seg_tree = etree.parse(os.path.join(directory, 'ann_segmentation.xml'))
  morph_tree = etree.parse(os.path.join(directory, 'ann_morphosyntax.xml'))
  morph_seg_list = morph_tree.xpath('/tei:teiCorpus/tei:TEI/tei:text/tei:body/tei:p/tei:s/tei:seg', namespaces = namespaces)

  morph_index = 0
  last_ab = 0
  start = 0

  for seg_entry in seg_tree.xpath("//tei:s/*", namespaces=namespaces):
    tag = tag_wo_ns(seg_entry)

    if tag == 'seg':
      i, morph_seg = find_morph(seg_entry.get('{http://www.w3.org/XML/1998/namespace}id'), morph_index, morph_seg_list)
      ab_corresp = seg_entry.get('corresp')
      seg_type = seg_entry.get('type', '')
      if ab_corresp: # Sprawdź, czy nie trzeba zacząć nowego akapitu
        ab_id = ab_corresp[26:].split('.', 1)[1].split('-', 1)[0]
        if ab_id != last_ab:
          if last_ab != 0:
            dag_out.write(u"\n")
          start = 0
          last_ab = ab_id
      if len(morph_seg) and seg_type not in OMITLIST:
        print_interps(dag_out, start, start + 1, morph_seg[0], seg_type, disamb)
        start += 1
      morph_index = i
    elif tag == 'choice':
      # Stwórz listę wariantów
      variant_list = []
      for var in seg_entry.iterchildren():
        if tag_wo_ns(var) == 'seg':
          variant_list.append([var])
        elif tag_wo_ns(var) == 'paren':
          variant_list.append(list(var))
      # Wylicz przesuniecie end ostatniego segmentu
      last_end = sum([len(v) - 1 for v in variant_list]) + 1
      v_end = start + 1
      for v in variant_list:
        v_start = start
        for seg_ind, seg in enumerate(v):
          i, morph_seg = find_morph(seg.get('{http://www.w3.org/XML/1998/namespace}id'), morph_index, morph_seg_list)
          ab_corresp = seg.get('corresp')
          seg_type = seg.get('type', '')
          if ab_corresp: # Sprawdź, czy nie trzeba zacząć nowego akapitu
            ab_id = ab_corresp[26:].split('.', 1)[1].split('-', 1)[0]
            if ab_id != last_ab:
              if last_ab != 0:
                dag_out.write(u"\n")
              start = 0
              v_start = 0
              v_end = 1
              last_ab = ab_id
          if len(morph_seg) and seg_type not in OMITLIST: 
            if len(v) == seg_ind + 1:
              print_interps(dag_out, v_start, start + last_end, morph_seg[0], seg_type, disamb)
            else:
              print_interps(dag_out, v_start, v_end, morph_seg[0], seg_type, disamb)
              v_start = v_end
              v_end += 1
#          else:
#            print u"Segment w choice bez interpretacji!!! ab_corresp: {}, seg_type: {}".format(ab_corresp, seg_type).encode("utf-8")
          morph_index = i
      start += last_end
  dag_out.write(u"\n")
  dag_out.close()

def insert_disamb(morph_seg, fVal, interp):
  fs_elem = morph_seg[0]
  new_f = etree.SubElement(fs_elem, '{http://www.tei-c.org/ns/1.0}f')
  new_f.set("name", "disamb")
  new_fs = etree.SubElement(new_f, '{http://www.tei-c.org/ns/1.0}fs')
  new_fs.set("type", "tool_report")
  new_f2 = etree.SubElement(new_fs, '{http://www.tei-c.org/ns/1.0}f')
  new_f2.set("name", "choice")
  new_f2.set("fVal", fVal)
  new_f3 = etree.SubElement(new_fs, '{http://www.tei-c.org/ns/1.0}f')
  new_f3.set("name", "interpretation")
  new_string = etree.SubElement(new_f3, '{http://www.tei-c.org/ns/1.0}string')
  new_string.text = interp
  
def process_seg(morph_seg, interp):
#  interp = seg_list[seg_list_id]['interpretation']
  if interp[0:2] == '::':  # uwaga na base = ':'
    interp_list = interp[2:].split(':', 1)
    interp_list.insert(0, ':')
  else:
    interp_list = interp.split(':', 2)  
  base = "&quot;" if interp_list[0] == '"' else interp_list[0]
  ctag = interp_list[1]
  msd = interp_list[2] if len(interp_list) > 2 else ''
  lex_list = morph_seg.xpath('./tei:fs/tei:f[@name="interps"]/tei:fs[@type="lex"][tei:f[@name="base"]/tei:string = "'+ base + '" and tei:f[@name="ctag"]/tei:symbol/@value="' + ctag + '"]', namespaces = namespaces)
  if len(lex_list):
    msd_list = lex_list[0].xpath('./tei:f[@name="msd"]//tei:symbol[@value="' + msd + '"]', namespaces=namespaces)
    fVal = "#" + msd_list[0].get('{http://www.w3.org/XML/1998/namespace}id') if len(msd_list) > 0 else ""
    insert_disamb(morph_seg, fVal, interp)
  else:
#    print u"Nie znalazłem base: {}, ctag: {} dla segmentu o id: {}".format(base, ctag, seg_entry.get('{http://www.w3.org/XML/1998/namespace}id')).encode('utf-8')
    insert_disamb(morph_seg, "", interp)

def insert_rejected(seg):
  print u"rejected dla id: {}".format(seg.get('{http://www.w3.org/XML/1998/namespace}id')).encode('utf-8')
  seg.set('{http://www.nkjp.pl/ns/1.0}rejected', "true")
  
def tag_tei(directory):
  print u"{}".format(directory).encode('utf-8')
  seg_list = read_text(os.path.join(directory, 'text_transcr.txt.morph.tagged'))
  parser = etree.XMLParser(remove_blank_text=True)
  seg_tree = etree.parse(os.path.join(directory, 'ann_segmentation.xml'), parser)
  morph_tree = etree.parse(os.path.join(directory, 'ann_morphosyntax.xml'), parser)
  morph_seg_list = morph_tree.xpath('/tei:teiCorpus/tei:TEI/tei:text/tei:body/tei:p/tei:s/tei:seg', namespaces = namespaces)

  morph_index = 0
  seg_list_id = 0

  for seg_entry in seg_tree.xpath("//tei:s/*", namespaces=namespaces):
    tag = tag_wo_ns(seg_entry)

    if tag == 'seg':
      if seg_entry.get('corresp'): # Jeśli segment nie ma "corresp", to nie będzie znakowany, więc go pomiń
        i, morph_seg = find_morph(seg_entry.get('{http://www.w3.org/XML/1998/namespace}id'), morph_index, morph_seg_list)
        
        if len(morph_seg):
          if morph_seg[0].xpath('./tei:fs/tei:f[@name="orth"]/tei:string', namespaces = namespaces)[0].text == seg_list[seg_list_id]['orth']:
            process_seg(morph_seg[0], seg_list[seg_list_id]['interpretation'])
            seg_list_id += 1
            morph_index = i
          else:
            print u"Niezgodność na segmencie o id: {}".format(seg_entry.get('{http://www.w3.org/XML/1998/namespace}id')).encode('utf-8')
            return
        else:
          print u"Brak opisu w morphosyntax dla segmentu o id: {}".format(seg_entry.get('{http://www.w3.org/XML/1998/namespace}id')).encode('utf-8')
      
    elif tag == 'choice':
      variant_found = False
      for variant in seg_entry.iterchildren():
        if tag_wo_ns(variant) == 'seg':
          if variant_found:
            insert_rejected(variant)
          else:
            i, morph_seg = find_morph(variant.get('{http://www.w3.org/XML/1998/namespace}id'), morph_index, morph_seg_list)
            if len(morph_seg):
              if morph_seg[0].xpath('./tei:fs/tei:f[@name="orth"]/tei:string', namespaces = namespaces)[0].text == seg_list[seg_list_id]['orth']:
                process_seg(morph_seg[0], seg_list[seg_list_id]['interpretation'])
                seg_list_id += 1
                morph_index = i
                variant_found = True
              else:
                insert_rejected(variant)
            else:
              print u"Brak opisu w morphosyntax dla segmentu z choice o id: {}".format(variant.get('{http://www.w3.org/XML/1998/namespace}id')).encode('utf-8')
              return
        elif tag_wo_ns(variant) == 'paren':
          if variant_found:
            for s in variant:
              insert_rejected(s)
          else:
            i, morph_seg = find_morph(variant[0].get('{http://www.w3.org/XML/1998/namespace}id'), morph_index, morph_seg_list)
            if len(morph_seg):
              if morph_seg[0].xpath('./tei:fs/tei:f[@name="orth"]/tei:string', namespaces = namespaces)[0].text == seg_list[seg_list_id]['orth']:
                process_seg(morph_seg[0], seg_list[seg_list_id]['interpretation'])
                seg_list_id += 1
                morph_index = i
                variant_found = True
                for s in variant[1:]:
                  j, morph_seg = find_morph(s.get('{http://www.w3.org/XML/1998/namespace}id'), morph_index, morph_seg_list)
                  if len(morph_seg):
                    if morph_seg[0].xpath('./tei:fs/tei:f[@name="orth"]/tei:string', namespaces = namespaces)[0].text == seg_list[seg_list_id]['orth']:
                      process_seg(morph_seg[0], seg_list[seg_list_id]['interpretation'])
                      seg_list_id += 1
                      morph_index = j
                    else:
                      print u"Brak zgodności zapisu orth dla segmentu z choice o id: {}".format(s.get('{http://www.w3.org/XML/1998/namespace}id')).encode('utf-8')
                      return
                  else:
                    print u"Brak opisu w morphosyntax dla segmentu z choice o id: {}".format(s.get('{http://www.w3.org/XML/1998/namespace}id')).encode('utf-8')
                    return
              else:
                for s in variant:
                  insert_rejected(s)
            else:
              print u"Brak opisu w morphosyntax dla segmentu z choice o id: {}".format(variant.get('{http://www.w3.org/XML/1998/namespace}id')).encode('utf-8')
              return
        else:
          print u"Niewłaściwy tag: {} zamiast seg lub paren w choice".format(tag_wo_ns(variant)).encode('utf-8')
          return

    else:
      print u"Niewłaściwy tag: {} zamiast seg lub choice".format(tag).encode('utf-8')
      return

  seg_tree.write(os.path.join(directory, 'ann_segmentation.xml'), pretty_print=True, encoding='UTF-8', xml_declaration=True)
  morph_tree.write(os.path.join(directory, 'ann_morphosyntax.xml'), pretty_print=True, encoding='UTF-8', xml_declaration=True)

#parser = argparse.ArgumentParser(description=u'Czyta pliki tagera w formacie DAG'.encode("utf-8"))
#parser.add_argument('filename', help=u'Ścieżka do pliku'.encode("utf-8"))
#args = parser.parse_args()

#seg_list = read_text(args.filename)

#print_dag_simple(seg_list)

#test_dag(seg_list, args.filename)