text: # -*- coding: utf-8 -*-
  
  __author__ = "Cristian Cardellino"
  
  import argparse
  import numpy as np
  import os
  import pandas as pd
  import re
  import spacy
  import sys
  
  from collections import defaultdict
  
  
  def build_cooccurrence_matrix(corpus, window_size=5, scale_factor="scaled",
                                vocab_size=5000, unkown_vector=True, verbose=False):
      """
      Builds a co-occurrence matrix with `window_size` and `scale_factor` over `corpus`.
  
      Parameters
      ----------
      corpus : iterator
          Should be a generator where each document is returned as a list of words. This
          corpus should already be normalized.
      window_size : int
          The size of the symmetrical window. The words to be considered when doing the count
          will be the ones being `window_size` before and after the center word.
      scale_factor : str, one of {"flat", "scaled"}
          The factor to use in order to weight the importance of a word regarding the center
          word. Can be either "flat" (equal weights to every word), or "scaled" where the
          word is weighted by 1/d where d is the distance to the center word.
      vocab_size : int
          The maximum size of the vocabulary
      unknown_vector : bool
          Whether to use a vector "UNK" for the words outside of the vocabulary.
      verbose : bool
          Activate to print a simple progress indicator
  
      Returns
      -------
      pd.DataFrame
          The (scaled) matrix of cooccurrences for the top `vocab_size` words in the vocabulary.
      """
  
      assert scale_factor in {"flat", "scaled"}
  
      word_count = defaultdict(int)
      word_word = defaultdict(int)
  
      for didx, document in enumerate(corpus, start=1):
          if verbose:
              print(f"\rDocument No. {didx}", end="", file=sys.stderr)
          for idx, word in enumerate(document):
              word_count[word] += 1
              lwindow = reversed(document[max(idx-window_size, 0):idx])
              rwindow = document[idx+1:idx+1+window_size]
  
              for lidx, lword in enumerate(lwindow):
                  word_word[(word, lword)] += 1/(lidx+1) if scale_factor == "scaled" else 1
  
              for ridx, rword in enumerate(rwindow):
                  word_word[(word, rword)] += 1/(ridx+1) if scale_factor == "scaled" else 1
  
      if verbose:
          print(file=sys.stderr)
  
      vocab = [word for word in sorted(word_count, key=lambda w: word_count[w],
                                       reverse=True)[:vocab_size]]
      vocab = {word: idx for idx, word in enumerate(sorted(vocab))}
  
      if unkown_vector:
          vocab["UNK"] = len(vocab)
  
      word_matrix = np.zeros((len(vocab), len(vocab)))
  
      for (w1, w2), count in word_word.items():
          w1_idx = vocab.get(w1, vocab.get("UNK", None))
          w2_idx = vocab.get(w2, vocab.get("UNK", None))
  
          if w1_idx is not None and w2_idx is not None:
              word_matrix[w1_idx, w2_idx] = count
  
      vocab_ = [word for word in sorted(vocab, key=lambda x: vocab[x])]
  
      return pd.DataFrame(word_matrix, index=vocab_, columns=vocab_)
  
  
  def corpus_processor(corpus_directory, language_model="es_core_news_sm",
                       remove_stopwords=True, lowercase=True):
      """
      Generator to retrieve and process the files of a corpus.
  
      Parameters
      ----------
      corpus_directory : str
          Path to where the corpus is stored as text documents.
      language_model : str
          The SpaCy language model to use for the processing.
      remove_stopwords : bool
          Whether to remove stopwords from the processed corpus.
      lowercase : bool
          Whether to lowercase the words.
  
      Returns
      -------
          generator
          A generator object where each entry is a processed document.
      """
  
      nlp = spacy.load(language_model)
  
      for fname in os.listdir(corpus_directory):
          with open(os.path.join(corpus_directory, fname), "r") as fh:
              # Careful with this for very large docs
              document = re.sub(r"\s+", " ", fh.read())
  
          nlp.max_length = max(len(document), nlp.max_length)
  
          tokens = [
              token.text for token in nlp(document, disable=["tagger", "parser", "ner"])
              if not (remove_stopwords and token.is_stop)
          ]
  
          if lowercase:
              tokens = [token.lower() for token in tokens]
  
          yield tokens
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser("preprocessing")
      parser.add_argument("corpus_directory",
                          help="Path to the directory holding the corpus files.")
      parser.add_argument("output_file",
                          help="Path to store the matrix (as csv file).")
      parser.add_argument("--language-model", "-l",
                          default="es_core_news_sm",
                          help="Name of the SpaCy language model to use for tokenization.")
      parser.add_argument("--ignore-unknown", "-u",
                          action="store_true",
                          help="Activate to avoid reserving a `UNK` vector for unkown words.")
      parser.add_argument("--scale-factor", "-f",
                          default="flat",
                          help="The scale factor. Can be either `flat` or `scaled`.")
      parser.add_argument("--use-stopwords", "-s",
                          action="store_true",
                          help="Activate to avoid ignoring the stopwords in the model.")
      parser.add_argument("--use-case", "-c",
                          action="store_true",
                          help="Activate to avoid word normalization via lowercase.")
      parser.add_argument("--vocab-size", "-v",
                          default=5000,
                          help="The maximum amount of words to consider.",
                          type=int)
      parser.add_argument("--verbose",
                          action="store_true",
                          help="Activate to print a simple progress value.")
      parser.add_argument("--window-size", "-w",
                          default=5,
                          help="The size of the co-occurrence window.",
                          type=int)
  
      args = parser.parse_args()
  
      cooccurrence_matrix = build_cooccurrence_matrix(
          corpus=corpus_processor(
              corpus_directory=args.corpus_directory,
              language_model=args.language_model,
              remove_stopwords=not args.use_stopwords,
              lowercase=not args.use_case
          ),
          window_size=args.window_size,
          scale_factor=args.scale_factor,
          vocab_size=args.vocab_size,
          unkown_vector=not args.ignore_unknown,
          verbose=args.verbose
      )
  
      cooccurrence_matrix.to_csv(args.output_file)
  