package de.dfki.lt.weka;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;

/**
 * {@link ArffWriter} provides methods to create Weka arff files from plain text
 * files containing a single sentence per line. Each sentence is regarded as a
 * document.
 *
 * @author Joerg Steffen, DFKI
 * @version $Id: NgramExtractor.java,v 1.1 2010/11/17 11:18:50 steffen Exp $
 */
public final class ArffWriter {

  /**
   * Would create a new instance of {@link ArffWriter}. Not to be used.
   */
  private ArffWriter() {

    // private constructor to enforce noninstantiability
  }


  /**
   * Creates arff files from the training and testing sentences in the given
   * data folder.
   *
   * @param dataFolderName
   *          the data folder
   * @throws IOException
   *           if there is an error when reading or writing files
   */
  public static void createArff(String dataFolderName)
      throws IOException {

    // collect languages from the files names in the data folder
    Collection<String> langs = new HashSet<String>();
    File dataFolder = new File(dataFolderName);
    File[] dataFiles = dataFolder.listFiles();
    for (File oneDataFile : dataFiles) {
      langs.add(oneDataFile.getName().substring(0, 3));
    }

    // maps each ngram to its global count
    Map<String, Integer> globalNgram2CountMap =
      new HashMap<String, Integer>();

    // maps each sentence to its language
    Map<String, String> sent2LangMap =
      new HashMap<String, String>();

    // maps each ngram to the set of sentences containing it
    Map<String, Set<String>> ngram2SentsMap =
      new HashMap<String, Set<String>>();

    // extract ngrams from training files
    Map<String, Map<String, Integer>> trainSent2NgramsMap =
      extractTrainNgrams(dataFolderName, langs, globalNgram2CountMap,
        sent2LangMap, ngram2SentsMap);

    // apply feature selection using document frequency threshold
    List<String> features = applyFeatureSelection(globalNgram2CountMap);

    // calculate idf for each ngram
    Map<String, Double> ngram2idf =
      calculateIdf(trainSent2NgramsMap.size(), ngram2SentsMap);

    // write training arff file
    writeArffFile(
      "train.arff", features, langs,
      ngram2idf, trainSent2NgramsMap, sent2LangMap);

    // extract ngrams from testing files
    Map<String, Map<String, Integer>> testSent2NgramsMap =
      extractTestNgrams(dataFolderName, langs, sent2LangMap);

    // write testing arff; ngrams that are NOT part of the features from
    // training are IGNORED;
    // the idf used for a feature is the one calculated in training;
    // only the term frequency is based on the counts in the testing sentences
    writeArffFile(
      "test.arff", features, langs,
      ngram2idf, testSent2NgramsMap, sent2LangMap);
  }


  /**
   * Extracts training ngrams from training documents in the given folder for
   * the given languages. As a side effect, updates the given maps.
   *
   * @param dataFolderName
   *          the data folder
   * @param langs
   *          the languages
   * @param globalNgram2CountMap
   *          maps each ngram to its global count
   * @param sent2LangMap
   *          maps each sentence to its language
   * @param ngram2SentsMap
   *          maps each ngram to the set of sentences containing it
   * @return trainSent2NgramsMap mapping of training sentences to its ngrams
   *         with their local counts
   * @throws IOException
   *           if there is an error when reading the training documents
   */
  private static Map<String, Map<String, Integer>> extractTrainNgrams(
      String dataFolderName,
      Collection<String> langs,
      Map<String, Integer> globalNgram2CountMap,
      Map<String, String> sent2LangMap, Map<String, Set<String>> ngram2SentsMap)
      throws IOException {

    // init result
    Map<String, Map<String, Integer>> trainSent2NgramsMap =
      new HashMap<String, Map<String, Integer>>();

    // get number of training sentences to use
    int trainDocNumber = Config.getInstance().getInt("training.docs");
    // for each language extract the training ngrams
    for (String oneLang : langs) {
      String trainFileName = dataFolderName + "/" + oneLang + "-train.txt";
      System.out.println(
        "extracting training ngrams from " + trainFileName + " ...");
      BufferedReader in = new BufferedReader(
          new InputStreamReader(
            new FileInputStream(trainFileName), "utf-8"));
      String sent;
      int trainCount = 0;
      while ((sent = in.readLine()) != null
          && trainCount < trainDocNumber) {
        trainCount++;
        // extract ngrams from sentence
        Map<String, Integer> localNgram2CountMap = extractNgrams(sent);
        // store ngrams for this sentence
        trainSent2NgramsMap.put(sent, localNgram2CountMap);
        // store language for this sentence
        sent2LangMap.put(sent, oneLang);
        // update global ngram counts and the mapping of ngrams to sentences
        for (Map.Entry<String, Integer> oneEntry
            : localNgram2CountMap.entrySet()) {
          String localNgram = oneEntry.getKey();
          int localNgramCount = oneEntry.getValue();
          // update global ngram counts
          Integer globalNgramCount = globalNgram2CountMap.get(localNgram);
          if (null == globalNgramCount) {
            globalNgram2CountMap.put(
              localNgram, localNgramCount);
          }
          else {
            globalNgram2CountMap.put(
              localNgram, globalNgramCount + localNgramCount);
          }
          // update mapping of ngrams to sentences
          Set<String> sentSet = ngram2SentsMap.get(localNgram);
          if (null == sentSet) {
            sentSet = new HashSet<String>();
            ngram2SentsMap.put(localNgram, sentSet);
          }
          sentSet.add(sent);
        }
      }
      in.close();
    }

    return trainSent2NgramsMap;
  }


  /**
   * Applies a document frequency threshold as defined int he config to do
   * feature selection.
   *
   * @param globalNgram2CountMap
   *          maps each ngram to its global count
   * @return the ngrams that have been selected as features
   */
  private static List<String> applyFeatureSelection(
      Map<String, Integer> globalNgram2CountMap) {

    System.out.println("applying feature selection ...");
    int threshold = Config.getInstance().getInt("doc.freq.threshold");
    List<String> features = new ArrayList<String>();
    for (Map.Entry<String, Integer> oneEntry
        : globalNgram2CountMap.entrySet()) {
      if (oneEntry.getValue() < threshold
          // exclude ngrams causing problems in the arff format
          || oneEntry.getKey().contains("'")
          || oneEntry.getKey().contains("\\")) {
        continue;
      }
      features.add(oneEntry.getKey());
    }
    Collections.sort(features);
    System.out.println("number of features: " + globalNgram2CountMap.size());
    System.out.println("filtered: "
      + (globalNgram2CountMap.size() - features.size()));

    return features;
  }


  /**
   * Calculates the inverse document frequency for the ngrams in the given map.
   *
   * @param totalSentCount
   *          the total number of training sentences
   * @param ngram2SentsMap
   *          maps each ngram to the set of sentences containing it
   * @return a mapping of ngrams to their idf weight
   */
  private static Map<String, Double> calculateIdf(
      int totalSentCount,
      Map<String, Set<String>> ngram2SentsMap) {

    System.out.println("calculating idf ...");
    Map<String, Double> ngram2idf = new HashMap<String, Double>();

    for (Map.Entry<String, Set<String>> oneEntry : ngram2SentsMap.entrySet()) {
      double idf =
        Math.log((double)totalSentCount / oneEntry.getValue().size());
      ngram2idf.put(oneEntry.getKey(), idf);
    }
    return ngram2idf;
  }


  /**
   * Extracts testing ngrams from testing documents in the given folder for the
   * given languages. As a side effect, updates the given maps.
   *
   * @param dataFolderName
   *          the data folder
   * @param langs
   *          the languages
   * @param sent2LangMap
   *          maps each sentence to its language
   * @return testSent2NgramsMap mapping of testing sentences to its ngrams with
   *         their local counts
   * @throws IOException
   *           if there is an error when reading the testing documents
   */
  private static Map<String, Map<String, Integer>> extractTestNgrams(
      String dataFolderName,
      Collection<String> langs,
      Map<String, String> sent2LangMap)
      throws IOException {

    // init result
    Map<String, Map<String, Integer>> testSent2NgramsMap =
      new HashMap<String, Map<String, Integer>>();

    // get number of testing sentences to use
    int testDocNumber = Config.getInstance().getInt("testing.docs");
    // for each language extract the testing ngrams
    for (String oneLang : langs) {
      String testFileName = dataFolderName + "/" + oneLang + "-test.txt";
      System.out.println(
        "extracting testing ngrams from " + testFileName + " ...");
      BufferedReader in = new BufferedReader(
          new InputStreamReader(
            new FileInputStream(testFileName), "utf-8"));
      String sent;
      int testCount = 0;
      while ((sent = in.readLine()) != null
          && testCount < testDocNumber) {
        testCount++;
        // extract ngrams from sentence
        Map<String, Integer> localNgram2CountMap = extractNgrams(sent);
        // store ngrams for this sentence
        testSent2NgramsMap.put(sent, localNgram2CountMap);
        // store language for this sentence
        sent2LangMap.put(sent, oneLang);
      }
      in.close();
    }

    return testSent2NgramsMap;
  }


  /**
   * Extracts the ngrams from the given sentence.
   *
   * @param sentence
   *          the string where to extract the ngrams from
   * @return a mapping of ngrams to their counts as found in the sentence
   */
  private static Map<String, Integer> extractNgrams(String sentence) {

    Map<String, Integer> ngram2CountMap = new HashMap<String, Integer>();
    int ngramSize = Config.getInstance().getInt("ngram.size");

    // extract ngrams and count them
    for (int i = 0; i < sentence.length() - ngramSize; i++) {
      String ngram = sentence.substring(i, i + ngramSize);
      Integer ngramCount = ngram2CountMap.get(ngram);
      if (null == ngramCount) {
        ngram2CountMap.put(ngram, 1);
      }
      else {
        ngram2CountMap.put(ngram, ngramCount + 1);
      }
    }

    return ngram2CountMap;
  }


  /**
   * Writes an arff file with the given name for the given features and
   * languages. Feature weights are calculated using tf.idf as derived from the
   * given ngram to idf map and the given sentence to ngrams map. The value of
   * the language feature is looked up in the given sentence to language map.
   *
   * @param fileName
   *          the arff file name
   * @param langs
   *          the languages
   * @param features
   *          the features
   * @param ngram2idf
   *          mapping of ngrams to their idf weight
   * @param sent2NgramsMap
   *          mapping of sentences to the ngrams with counts
   * @param sent2LangMap
   *          mapping of sentences to their language
   * @throws IOException
   *           if there is an error when writing the file
   */
  private static void writeArffFile(
      String fileName, List<String> features, Collection<String> langs,
      Map<String, Double> ngram2idf,
      Map<String, Map<String, Integer>> sent2NgramsMap,
      Map<String, String> sent2LangMap)
      throws IOException {

    System.out.println("writing " + fileName + " ...");
    PrintWriter out = new PrintWriter(
      new BufferedWriter(
        new OutputStreamWriter(
          new FileOutputStream(fileName), "utf-8")));

    // write the relation
    out.println("@relation langId");
    out.println();

    // write the attributes for the ngram features, all numeric
    for (String oneFeature : features) {
      out.print("@attribute '");
      out.print(oneFeature);
      out.println("' numeric");
    }

    // write the language attribute, nominal
    out.print("@attribute lang {");
    List<String> sortedLangs = new ArrayList<String>(langs);
    Collections.sort(sortedLangs);
    Iterator<String> langIt = sortedLangs.iterator();
    while (langIt.hasNext()) {
      String oneLang = langIt.next();
      out.print(oneLang);
      if (langIt.hasNext()) {
        out.print(",");
      }
    }
    out.println("}");
    out.println();

    // write the attribute values
    out.println("@data");
    // for each sentence, calculate the tf.idf weight for all its features
    for (Map.Entry<String, Map<String, Integer>> oneEntry
        : sent2NgramsMap.entrySet()) {
      String sent = oneEntry.getKey();
      // sum up the total number of ngrams extracted from this sentence
      int totalSentNgramCount = 0;
      for (Integer oneNgramCount : oneEntry.getValue().values()) {
        totalSentNgramCount = totalSentNgramCount + oneNgramCount;
      }
      Map<String, Integer> localNgram2CountMap = oneEntry.getValue();
      for (String oneFeature : features) {
        // get term frequency
        double termFreq = 0.0;
        Integer localNgramCount = localNgram2CountMap.get(oneFeature);
        if (null != localNgramCount) {
          // normalize ngram count by the total number of ngrams extracted
          // from this sentence
          termFreq = (double)localNgramCount / totalSentNgramCount;
        }
        if (0 == termFreq) {
          out.print("0,");
          continue;
        }
        // get idf
        double idf = 0.0;
        Double idfDouble = ngram2idf.get(oneFeature);
        if (idfDouble != null) {
          idf = idfDouble;
        }
        double tfidf = termFreq * idf;
        out.print(String.format(Locale.ENGLISH, "%.5f", tfidf));
        out.print(",");
      }
      // last attribute is the language
      out.println(sent2LangMap.get(sent));
    }

    out.close();
  }


  /**
   * The main method; requires no arguments.
   *
   * @param args
   *          the arguments; not used here
   */
  public static void main(String[] args) {

    try {
      ArffWriter.createArff("data");
    } catch (Exception e) {
      e.printStackTrace();
    }
  }
}
