package zemberek.lm.backoff;

import com.google.common.base.Charsets;
import com.google.common.base.Splitter;
import com.google.common.io.Files;
import com.google.common.io.LineProcessor;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;
import zemberek.core.logging.Log;
import zemberek.core.math.LogMath;
import zemberek.lm.BaseLanguageModel;
import zemberek.lm.LmVocabulary;
import zemberek.lm.NgramLanguageModel;

/**
 * Simple Back-off N-gram model that loads an ARPA file. This class is not designed for compactness
 * or speed but it can be used for validation and debugging. <p> aaa,cd
 */
public class SimpleBackoffNgramModel extends BaseLanguageModel implements NgramLanguageModel {

  public final double unigramWeight;
  private Map<NgramData, NgramProb> probabilities;
  private List<Integer> counts;

  private SimpleBackoffNgramModel(
      int order,
      LmVocabulary vocabulary,
      Map<NgramData, NgramProb> probabilities,
      List<Integer> counts,
      double unigramWeight) {
    super(order, vocabulary);
    this.probabilities = probabilities;
    this.counts = counts;
    this.unigramWeight = unigramWeight;
  }

  /**
   * Constructs a back-off language model from an Arpa file. Uni-gram weight applies smoothing to
   * uni-gram values. A value of 1 means uni-gram probabilities from model directly used.
   *
   * @param arpaModel arpa model file
   * @param unigramWeigth unigram weight. in the moethod code there is explanation about it.
   * @throws IOException if things go wrong when accessing the arpa file.
   */
  public static SimpleBackoffNgramModel fromArpa(File arpaModel, float unigramWeigth)
      throws IOException {
    if (unigramWeigth < 0 || unigramWeigth > 1) {
      throw new IllegalArgumentException(
          "Unigram weight must be between 0 and 1 but it is:" + unigramWeigth);
    }
    ArpaLoader converter = new ArpaLoader(unigramWeigth);
    return Files.asCharSource(arpaModel, Charsets.UTF_8).readLines(converter);
  }

  /**
   * Constructs a back-off language model from an Arpa file. No unigram weight is applied.
   *
   * @param arpaModel arpa model file
   * @throws IOException if things go wrong when accessing the file.
   */
  public static SimpleBackoffNgramModel fromArpa(File arpaModel) throws IOException {
    return fromArpa(arpaModel, 1.0f);
  }

  @Override
  public float getUnigramProbability(int id) {
    NgramData data = new NgramData(id);
    NgramProb res = probabilities.get(data);
    if (res == null) {
      throw new IllegalArgumentException(
          "Word does not exist!" + vocabulary.getWordsString(id) + " with index:" + id);
    } else {
      return res.prob;
    }
  }

  @Override
  public boolean ngramExists(int... wordIndexes) {
    if (wordIndexes.length < 1 || wordIndexes.length > order - 1) {
      throw new IllegalArgumentException("Amount of tokens must be between 1 and " +
          order + " But it is " + wordIndexes.length);
    }
    NgramData data = new NgramData(wordIndexes);
    NgramProb res = probabilities.get(data);
    return res != null;
  }

  public double getLogBase() {
    return Math.E;
  }

  private float getProbabilityValue(int... ids) {
    NgramData data = new NgramData(ids);
    NgramProb res = probabilities.get(data);
    if (res != null) {
      return res.prob;
    } else {
      return LogMath.LOG_ZERO_FLOAT;
    }
  }

  private float getBackoffValue(int... ids) {
    NgramData data = new NgramData(ids);
    NgramProb res = probabilities.get(data);
    if (res != null) {
      return res.backoff;
    } else {
      return 0;
    }
  }

  public Iterator<NgramData> getAllIndexes() {
    return probabilities.keySet().iterator();
  }

  @Override
  public float getProbability(int... ids) {
    for (int id : ids) {
      if (!vocabulary.contains(id)) {
        throw new IllegalArgumentException(
            "Unigram does not exist!" + vocabulary.getWordsString(id) + " with index:" + id);
      }
    }
    float result = 0;
    float probability = getProbabilityValue(ids);
    if (probability == LogMath.LOG_ZERO_FLOAT) { // if probability does not exist.
      if (ids.length == 1) {
        return LogMath.LOG_ZERO_FLOAT;
      }
      float backoffValue = getBackoffValue(head(ids));
      result = result + backoffValue + getProbability(tail(ids));
    } else {
      result = probability;
    }
    return result;
  }

  @Override
  public float getTriGramProbability(int id0, int id1, int id2) {
    return getProbability(id0, id1, id2);
  }

  @Override
  public float getTriGramProbability(int id0, int id1, int id2, int fingerPrint) {
    return getProbability(id0, id1, id2);
  }

  @Override
  public int getOrder() {
    return order;
  }

  public int getGramCount(int order) {
    return counts.get(order - 1);
  }

  @Override
  public LmVocabulary getVocabulary() {
    return vocabulary;
  }

  public static class NgramData {

    int[] indexes;

    private NgramData(int... indexes) {
      this.indexes = indexes;
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) {
        return true;
      }
      if (o == null || getClass() != o.getClass()) {
        return false;
      }
      NgramData ngramData = (NgramData) o;
      return Arrays.equals(indexes, ngramData.indexes);
    }

    @Override
    public int hashCode() {
      return indexes != null ? Arrays.hashCode(indexes) : 0;
    }

    @Override
    public String toString() {
      return Arrays.toString(indexes);
    }

    public int[] getIndexes() {
      return indexes;
    }
  }

  private static class NgramProb {

    float prob;
    float backoff;

    private NgramProb(float prob, float backoff) {
      this.prob = prob;
      this.backoff = backoff;
    }

    @Override
    public String toString() {
      return prob + " " + backoff;
    }

  }

  private static class ArpaLoader implements LineProcessor<SimpleBackoffNgramModel> {

    public static final int DEFAULT_UNKNOWN_PROBABILITY = -20;
    private final Pattern SPLIT_PATTERN = Pattern.compile("\\s+");
    LmVocabulary.Builder vocabularyBuilder = LmVocabulary.builder();
    LmVocabulary vocabulary;
    Map<NgramData, NgramProb> probabilities = new HashMap<>();
    int lineCounter = 0;
    int _n;
    int order;
    float logUniformUnigramProbability;
    float logUnigramWeigth;
    float inverseLogUnigramWeigth;
    float unigramWeight;
    State state = State.BEGIN;
    List<Integer> ngramCounts = new ArrayList<>();
    boolean started = false;

    long start;

    ArpaLoader(float unigramWeight) {
      this.unigramWeight = unigramWeight;
      start = System.currentTimeMillis();
      logUnigramWeigth = (float) Math.log(unigramWeight);
      inverseLogUnigramWeigth = (float) Math.log(1 - unigramWeight);
    }

    public boolean processLine(String s) throws IOException {
      String clean = s.trim();
      switch (state) {
        // read n value and ngram counts.
        case BEGIN:
          if (clean.length() == 0) {
            break;
          }
          if (clean.startsWith("\\data\\")) {
            started = true;
          } else if (started && clean.startsWith("ngram")) {
            started = true;
            int count = 0, i = 0;
            for (String str : Splitter.on("=").trimResults().split(clean)) {
              if (i++ == 0) {
                continue;
              }
              count = Integer.parseInt(str);
            }
            ngramCounts.add(count);
          } else if (started) {
            order = ngramCounts.size();
            Log.info("order = " + order);
            state = State.UNIGRAMS;
            int i = 0;
            for (Integer count : ngramCounts) {
              Log.info((++i) + "gram count:" + count);
            }
            Log.info("Processing unigrams.");
            logUniformUnigramProbability = (float) -Math
                .log(ngramCounts.get(0)); // log(1/#count) = -log(#count)
            _n++;
          }
          break;

        // read ngrams. if unigram values, we store the strings and related indexes.
        case UNIGRAMS:
          if (clean.length() == 0 || clean.startsWith("\\")) {
            break;
          }
          Iterator<String> it = Splitter.on(SPLIT_PATTERN).split(clean).iterator();
          // parse probabilty
          float logProbability = (float) LogMath.log10ToLog(Double.parseDouble(it.next()));
          if (unigramWeight < 1) {
            // apply uni-gram weight. This applies smoothing to unigrams. As lowering high probabilities and
            // adding gain to small probabilities.
            // uw = uni-gram weight  , uniformProb = 1/#unigram
            // so in linear domain, we apply this to all probability values as: p(w1)*uw + uniformProb * (1-uw) to
            // maintain the probability total is one while smoothing the values.
            // this converts to log(p(w1)*uw + uniformProb*(1-uw) ) which is calculated with log probabilities
            // a = log(p(w1)) + log(uw) and b = -log(#unigram)+log(1-uw) applying logsum(a,b)
            // approach is taken from Sphinx-4
            float p1 = logProbability + logUnigramWeigth;
            float p2 = logUniformUnigramProbability + inverseLogUnigramWeigth;
            logProbability = LogMath.LOG_SUM_FLOAT.lookup(p1, p2);
          }
          String word = it.next();
          float logBackoff = 0;
          if (it.hasNext()) {
            logBackoff = (float) LogMath.log10ToLog(Double.parseDouble(it.next()));
          }
          int index = vocabularyBuilder.add(word);
          probabilities.put(new NgramData(index), new NgramProb(logProbability, logBackoff));
          lineCounter++;

          if (lineCounter == ngramCounts.get(0)) {
            handleSpecialToken("<unk>");
            handleSpecialToken("</s>");
            handleSpecialToken("<s>");
            vocabulary = vocabularyBuilder.generate();
            // update the ngram counts in case a special token is added.
            ngramCounts.set(0, vocabulary.size());
            lineCounter = 0;
            state = State.NGRAMS;
            _n++;
            // if there is only unigrams in the arpa file, exit
            if (ngramCounts.size() > 1) {
              Log.info("Processing 2-grams.");
            }
          }
          break;

        case NGRAMS:
          if (clean.length() == 0 || clean.startsWith("\\")) {
            break;
          }
          Iterator<String> it2 = Splitter.on(SPLIT_PATTERN).split(clean).iterator();
          logProbability = (float) LogMath.log10ToLog(Double.parseDouble(it2.next()));

          int[] ids = new int[_n];
          for (int i = 0; i < _n; i++) {
            ids[i] = vocabulary.indexOf(it2.next());
          }

          logBackoff = 0;
          if (_n < ngramCounts.size()) {
            if (it2.hasNext()) {
              logBackoff = (float) LogMath.log10ToLog(Double.parseDouble(it2.next()));
            }
          }

          probabilities.put(new NgramData(ids), new NgramProb(logProbability, logBackoff));

          lineCounter++;
          if (lineCounter == ngramCounts.get(_n - 1)) {
            // if there is no more ngrams, exit
            if (ngramCounts.size() == _n) {
              return false;
            } else {
              lineCounter = 0;
              _n++;
              Log.info("Processing " + _n + "-grams.");
            }
          }
          break;
      }
      return true;
    }

    // adds special token with default probability.
    private void handleSpecialToken(String word) throws IOException {
      if (vocabularyBuilder.indexOf(word) == -1) {
        Log.warn("Special token " + word +
            " does not exist in model. It is added with default unknown probability: " +
            DEFAULT_UNKNOWN_PROBABILITY);
        int index = vocabularyBuilder.add(word);
        probabilities.put(new NgramData(index), new NgramProb(DEFAULT_UNKNOWN_PROBABILITY, 0));
      }
    }

    public SimpleBackoffNgramModel getResult() {
      return new SimpleBackoffNgramModel(order, vocabulary, probabilities, ngramCounts,
          unigramWeight);
    }

    enum State {
      BEGIN, UNIGRAMS, NGRAMS
    }
  }
}
