text: from collections import OrderedDict
  import os 
  
  import pandas as pd 
  import numpy as np
  np.random.seed(896)
  
  from os.path import join, dirname, abspath
  import sys
  
  # this monstrosity produces the module directory in an environment where this is unpacked
  sys.path.insert(0, abspath(join(dirname(abspath(__file__)), '..', '..')))
  
  from sklearn.feature_extraction.text import CountVectorizer
  from evidence_inference.preprocess.sentence_split import find_span_location, split_into_sentences, gen_exact_evid_array
  from evidence_inference.preprocess.article_reader import TableHTMLParser
  import evidence_inference.preprocess.article_reader as article_reader
  parser = TableHTMLParser()
  
  # this monstrosity points us to a root directory relative to this file
  annotation_root = os.path.abspath(os.path.join(__file__, "..", "..", "..", "annotations"))
  anno_csv_path = os.path.join(annotation_root, "annotations_merged.csv")  # "pilot_run_data/annotations.csv"
  prompts_csv_path = os.path.join(annotation_root, "prompts_merged.csv")  # "pilot_run_data/prompts.csv"
  base_XML_path = os.path.join(annotation_root, "xml_files")
  _train_id_file, _validation_id_file, _test_id_file = [os.path.join(annotation_root, 'splits', d) for d in ['train_article_ids.txt', 'validation_article_ids.txt', 'test_article_ids.txt']]
  
  if not all(os.path.exists(x) for x in [anno_csv_path, prompts_csv_path, base_XML_path]):
      raise RuntimeError("One of {} does not exist".format([anno_csv_path, prompts_csv_path, base_XML_path]))
  
  # base_XML_path = os.path.join("pilot_run_data", "xml_files")
  
  PROMPT_ID_COL_NAME = "PromptID"
  LBL_COL_NAME = "Label Code"
  LABEL = 'Label'
  EVIDENCE_COL_NAME = "Annotations"
  EVIDENCE_START = "Evidence Start"
  EVIDENCE_END = "Evidence End"
  STUDY_ID_COL = "PMCID"
  VALID_LABEL = "Valid Label"
  VALID_REASONING = "Valid Reasoning"
  
  
  def get_article(article_id):
      xml_str = "PMC{}.nxml".format(article_id)
      xml_path = os.path.join(base_XML_path, xml_str)
      return article_reader.Article(xml_path)
  
  
  def read_in_articles(article_ids=None):
  
      anno_df = pd.read_csv(anno_csv_path)
      unique_article_ids = anno_df[STUDY_ID_COL].unique()
  
      articles = []
  
      for article_id in unique_article_ids:
          #import pdb; pdb.set_trace()
  
          if article_ids is None or article_id in article_ids:
              # 2376383
              articles.append(get_article(article_id))
  
      return articles
  
  
  def extract_raw_text(article, sections_of_interest=None):
      if sections_of_interest is None:
          #sections_of_interest = ["results", ""]
          sections_of_interest = article.article_dict.keys()
     
      ti_ab = article.ti_ab_str()
  
      article_sections = [sec for sec in article.article_dict.keys() if any(
                              [s in sec for s in sections_of_interest])]
      article_body = article.to_raw_str(fields=article_sections)
      
      raw_text = ti_ab + "  " + article_body
          
      return raw_text.replace("<p>", "")
  
  def extract_section_text(article):
      """
      Takes in an article, and maps section headers to the text that would be produced.
      """
      sections_of_interest = article.article_dict.keys()
     
      ti_ab = article.ti_ab_str()
  
      article_sections = [sec for sec in article.article_dict.keys() if any(
                              [s in sec for s in sections_of_interest])]
          
      section_mapping = {'Title': (ti_ab + "  ").replace("<p>", "")}
      for section_header in article_sections:
          section_mapping[section_header] = article.to_raw_str(fields=[section_header]).replace("<p>", "")           
              
          ## if not the last one
          if article_sections[-1] != section_header:
              section_mapping[section_header] += "\n\n"
          
      return section_mapping
  
  
  def extract_text_from_prompts(prompts_df):
      I, C, O = prompts_df['Intervention'].values, prompts_df['Comparator'].values, prompts_df['Outcome'].values
      all_prompt_text = [s.lower() for s in np.concatenate([I, C, O])]
      return all_prompt_text
  
  def get_inference_vectorizer(article_ids=None, sections_of_interest=None, vocabulary_file=None):
  
      # if article_ids is None, will use all articles
      # in the CSV passed to the read_in_articles method.
      articles = read_in_articles(article_ids=article_ids)
      raw_texts = [extract_raw_text(article, sections_of_interest) for article in articles]
  
      # we also use the prompts text to construct our vectorizer
      prompts = read_prompts()
      raw_prompt_text = " ".join(extract_text_from_prompts(prompts))
  
      raw_texts.append(raw_prompt_text)
  
      # there is at least one prompt with tokens short enough that CountVectorizer's default destroys it, so we allow any single character through.
      if vocabulary_file is not None:
          with open(vocabulary_file, 'r') as vf:
              vocab = [line.strip() for line in vf]
          vectorizer = CountVectorizer(vocabulary=vocab, token_pattern=r"\b\w+\b")
          print("Loaded {} words from vocab file {}".format(len(vocab), vocabulary_file))
      else:
          vectorizer = CountVectorizer(max_features=20000, token_pattern=r"\b\w+\b")
      vectorizer.fit(raw_texts)
      tokenizer = vectorizer.build_tokenizer() 
  
      str_to_idx = vectorizer.vocabulary_
      str_to_idx[SimpleInferenceVectorizer.PAD] = max(vectorizer.vocabulary_.values())
      str_to_idx[SimpleInferenceVectorizer.UNK] = str_to_idx[SimpleInferenceVectorizer.PAD]+1
      
      # note that for now the vectorizer is fit using only the
      # article texts (i.e., the vocab is based on words in full-texts,
      # not in prompts necessarily).
      return SimpleInferenceVectorizer(str_to_idx, tokenizer)
  
  
  def read_annotations():
      anno_df = pd.read_csv(anno_csv_path)
      # we need to force EVIDENCE_COL_NAME to be strings in all cases; pandas occasionally reads some values as floats.
      anno_df = anno_df[anno_df.apply(lambda row: bool(row[VALID_LABEL]) and bool(row[VALID_REASONING]) and len(str(row[EVIDENCE_COL_NAME])) > 0 and row[LABEL] != 'invalid prompt', axis=1)]
      #annos[~annos["Answer_Val"].isin([-1, 0, 1])]
      # TODO revisit this; right now just overwriting for convienence
      # anno_df["Answer_Val"].replace({3:0}, inplace=True)
      return anno_df
  
  
  def read_prompts():
      prompts_df = pd.read_csv(prompts_csv_path)
      prompts_df = prompts_df[prompts_df.apply(lambda row: all(map(lambda x: type(x) == str and x is not None and bool(x.strip()), [row['Comparator'], row['Intervention'], row['Outcome']])), axis=1)]
      return prompts_df 
  
  
  def assemble_Xy_for_prompts(training_prompts, inference_vectorizer, lbls_too=False, annotations=None, sections_of_interest=None, include_sentence_span_splits = False): 
      Xy = []
  
      for prompt_id in training_prompts[PROMPT_ID_COL_NAME].values:
          if lbls_too:
              Xy_dict = inference_vectorizer.vectorize(training_prompts, prompt_id, 
                                  include_lbls=True, annotations_df=annotations, sections_of_interest=sections_of_interest, include_sentence_span_splits = include_sentence_span_splits)
          else:
              Xy_dict = inference_vectorizer.vectorize(training_prompts, 
                                                       prompt_id, 
                                                       sections_of_interest=sections_of_interest, 
                                                       include_sentence_span_splits = include_sentence_span_splits)
          Xy.append(Xy_dict)
      return Xy
  
  
  def train_document_ids():
      """ Returns the set of document ids for a fixed training set """
      with open(_train_id_file, 'r') as tf:
          ids = list(int(x.strip()) for x in tf.readlines())
          ids_dict = OrderedDict()
          for x in ids:
              ids_dict[x] = x
          return ids_dict.keys()
  
  def validation_document_ids():
      """ Returns the set of document ids for a fixed validation set """
      with open(_validation_id_file, 'r') as vf:
          ids = list(int(x.strip()) for x in vf.readlines())
          ids_dict = OrderedDict()
          for x in ids:
              ids_dict[x] = x
          return ids_dict.keys()
  
  def test_document_ids():
      """ Returns the set of documents for a fixed test set """
      with open(_test_id_file, 'r') as tf:
          ids = list(int(x.strip()) for x in tf.readlines())
          ids_dict = OrderedDict()
          for x in ids:
              ids_dict[x] = x
          return ids_dict.keys()
  
  def get_train_Xy(train_doc_ids, sections_of_interest=None, vocabulary_file=None, include_sentence_span_splits = False):
      """ Loads the relevant documents, builds a vectorizer, and returns a list of training instances"""
      prompts = read_prompts()
      annotations = read_annotations()
  
      # filter out prompts for which we do not have annotations for whatever reason
      # this was actually just one case; not sure what was going on there.
      def have_annotations_for_prompt(prompt_id):
          return len(annotations[annotations[PROMPT_ID_COL_NAME] == prompt_id]) > 0
  
      prompts = [prompt for row_idx, prompt in prompts.iterrows() if 
                              have_annotations_for_prompt(prompt[PROMPT_ID_COL_NAME])]
      prompts = pd.DataFrame(prompts)
  
      inference_vectorizer = get_inference_vectorizer(article_ids=train_doc_ids, sections_of_interest=sections_of_interest, vocabulary_file=vocabulary_file)
  
      training_prompts = prompts[prompts[STUDY_ID_COL].isin(train_doc_ids)]
  
      training_prompts = pd.DataFrame(training_prompts)
      train_Xy = assemble_Xy_for_prompts(training_prompts, 
                                         inference_vectorizer, 
                                         lbls_too=True, 
                                         annotations=annotations, 
                                         include_sentence_span_splits = include_sentence_span_splits)
  
      return train_Xy, inference_vectorizer
  
  
  def get_Xy(docids, inference_vectorizer: 'SimpleInferenceVectorizer', sections_of_interest=None, include_sentence_span_splits = False):
      prompts = read_prompts()
      annotations = read_annotations()
  
      # filter out prompts for which we do not have annotations for whatever reason
      # this was actually just one case; not sure what was going on there.
      def have_annotations_for_prompt(prompt_id):
          return len(annotations[annotations[PROMPT_ID_COL_NAME] == prompt_id]) > 0
  
      prompts = [prompt for row_idx, prompt in prompts.iterrows() if
                 have_annotations_for_prompt(prompt[PROMPT_ID_COL_NAME])]
      prompts = pd.DataFrame(prompts)
  
      prompts = prompts[prompts[STUDY_ID_COL].isin(docids)]
      Xy = assemble_Xy_for_prompts(prompts, inference_vectorizer, lbls_too=True, 
                                   annotations=annotations, 
                                   sections_of_interest=sections_of_interest, 
                                   include_sentence_span_splits = include_sentence_span_splits)
      return Xy
  
  
  class SimpleInferenceVectorizer:
      UNK = "<unk>"
      PAD = "<pad>"
  
      def __init__(self, str_to_idx, tokenizer):
          self.str_to_idx = str_to_idx
          self.idx_to_str = [None]*(len(self.str_to_idx))
          self.sentence_splits = {} # map of article ids to array of sentence splits
          self.token_evidence  = {}
          
          for w, idx in self.str_to_idx.items():
              try:
                  self.idx_to_str[idx] = w 
              except:
                  import pdb; pdb.set_trace()
  
          self.tokenizer = tokenizer
  
      def string_to_seq(self, s):
          tokenized = self.tokenizer(s)
          unk_idx = self.str_to_idx[SimpleInferenceVectorizer.UNK]
          vectorized = [self.str_to_idx.get(token, unk_idx) for token in tokenized]
          return np.array(vectorized)
  
      def vectorize(self, prompts_df, prompt_id, include_lbls=False, annotations_df=None, sections_of_interest=None, include_sentence_span_splits = False):
          """
          Vectorize the prompt specified by the ID.
          """
          if include_lbls and annotations_df is None:
              raise ValueError("When including annotations, they must already be defined")
          
          prompt = prompts_df[prompts_df[PROMPT_ID_COL_NAME]==prompt_id]
  
          ###
          # vectorize the article itself.
          article_id = str(prompt[STUDY_ID_COL].values[0])
          article = get_article(article_id)
          article_text = extract_raw_text(article, sections_of_interest)
          article_text = article_text.lower()
          vectorized_article = self.string_to_seq(article_text)
          
          ###
          # and now vectorize the prompt (I/C/O)
          I_v = self.string_to_seq(prompt["Intervention"].values[0].lower())
          C_v = self.string_to_seq(prompt["Comparator"].values[0].lower())
          O_v = self.string_to_seq(prompt["Outcome"].values[0].lower())
          
          ###
          # if we want section splits, we need to encode individual sections and count
          sections = extract_section_text(article)
          lengths = []
          section_titles = []
          for k in sections.keys():
              lengths.append(len(self.string_to_seq(sections[k])))
              section_titles.append(k)
              
          if not(sum(lengths) == len(vectorized_article)):
              import pdb; pdb.set_trace()
          
          return_dict = {"article":vectorized_article, 
                         "I":I_v, 
                         "C":C_v, 
                         "O":O_v, 
                         "a_id": article_id, 
                         "p_id": prompt_id, 
                         "section_splits":lengths,
                         "section_titles": section_titles}
  
          if include_lbls:
              # then also read out the labels.
              assert (annotations_df is not None)
              annotations_for_prompt = annotations_df[annotations_df[PROMPT_ID_COL_NAME] == prompt_id]
              labels = annotations_for_prompt[[LBL_COL_NAME,EVIDENCE_COL_NAME]].values
              return_dict["y"] = labels
              # remove html tags
              for l in labels:
                  parser.feed(str(l[1]))
                  l[1] = parser.get_data()
              
              spans = annotations_for_prompt[[EVIDENCE_START,EVIDENCE_END]].values
              if len(spans) > 0 and sections_of_interest is None:
                  # split into sentences, find which are evidence, and also encode all.
                  sentence_spans = []
                  if include_sentence_span_splits:
                      sen = split_into_sentences(article_id, article_text, self.sentence_splits)
                      tmp = find_span_location(sen, [s[0] for s in spans], [e[1] for e in spans])
                      for t in tmp:
                          sentence_spans.append([self.string_to_seq(t[0]), t[1]])
                   
                  # encode the evidence spans 
                  evidence_spans = set()
                  for start, end in spans:
                      article_before_span = article_text[:int(start)]
                      # +1 because the slice gets every character (and therefore token) *before* the evidence, so we want to offset the token count by 1 to actually start in the evidence
                      span_start_idx = len(self.tokenizer(article_before_span)) + 1
                      article_at_end_of_span = article_text[:int(end)]
                      span_end_idx = len(self.tokenizer(article_at_end_of_span))
                      evidence_spans.add((span_start_idx, span_end_idx))
                     
                  return_dict['sentence_span'] = sentence_spans  
                  return_dict['evidence_spans'] = evidence_spans
                  if include_sentence_span_splits:
                      return_dict['token_ev_labels'] = gen_exact_evid_array(sentence_spans, evidence_spans, return_dict, self.idx_to_str)
                  
          return return_dict
  
      def decode(self, v):
          return [self.idx_to_str[idx] for idx in v]