import corpus2

def get_chan_dict(base_graph, chan_name):
  '''
  Wydobycie chan z tokenow w zdaniu.
  Wydobywany jest numer tokenu, w ktorym znajduje sie 
  wybrany kanal anotacji oraz jego wartosc w tym kanale
      
  Budowa slownika:
    chan_dict[(id. zdania, id. dla tokenu w kanale)] = id. tokenu w zdaniu
  '''
  chan_dict = {}
  for paragraph in base_graph.document.paragraphs(): 
    for sentence in paragraph.sentences():
      asent = corpus2.AnnotatedSentence.wrap_sentence(sentence)
      if asent.has_channel(chan_name):
        chan = asent.get_channel(chan_name)
        ann_vec = chan.make_annotation_vector()
        for ann in ann_vec:
          key = (sentence.id(), ann.seg_number)
          if not chan_dict.has_key(key):
            chan_dict[key] = ann.head_index
          else:
            print 'W tekscie wejsciowym dla kanalu', chan, \
              'w zdaniu', sentence.id(), \
              'powtorzyl sie numer segmentu', \
              ann.seg_number, '!'
  return chan_dict

def get_rel_list_from_file(base_graph, chan_dict, chan_name):
  '''
  Utworzenie listy relacji znajdujacych sie w pliku input_rel:
  chan -> [
      [sentence_id, from_node_id, to_node_id, rel_name, extra_node_id], 
      [], ..., []]
  '''
  rel_list = []
  for index, rel in enumerate(base_graph.rel_root.findall('rel')):
    chan = rel.find('from').get('chan')
    if chan != chan_name:
      continue
    
    rel_name = rel.get('name')
    if rel_name and rel_name == '3args_rel_2nd_part':
      continue
    
    sentence_id = rel.find('from').get('sent')
    extra_node = None
    extra_node_id = None
    
    if rel_name == '3args_rel_1st_part':
      from_node = int(
        base_graph.rel_root.findall('rel')[index].find('from').text)
      to_node = int(
        base_graph.rel_root.findall('rel')[index + 1].find('to').text)
      extra_node = int(
        base_graph.rel_root.findall('rel')[index].find('to').text)
      extra_node_id = chan_dict[(sentence_id, extra_node)]
    else:
      from_node = int(rel.find('from').text)
      to_node = int(rel.find('to').text)
    
    from_node_id = chan_dict[(sentence_id, from_node)]
    to_node_id = chan_dict[(sentence_id, to_node)]
    rel = [sentence_id, \
        from_node_id, \
        to_node_id, \
        rel_name if chan == 'np_sem_roles' else None, \
        extra_node_id if extra_node_id else None]
    rel_list.append(rel)
  return rel_list