package zemberek.lm.compression;

import com.google.common.base.Joiner;
import com.google.common.base.Splitter;
import com.google.common.io.Closeables;
import com.google.common.io.Files;
import com.google.common.io.LineProcessor;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.List;
import zemberek.core.SpaceTabTokenizer;
import zemberek.core.logging.Log;
import zemberek.core.quantization.DoubleLookup;
import zemberek.core.quantization.Quantizer;
import zemberek.core.quantization.QuantizerType;
import zemberek.lm.LmVocabulary;

/**
 * This is a multiple file representation of an uncompressed backoff language model. Files are:
 * <p/>info <p/>int32 order <p/>int32 1 gram count <p/>int32 2 gram count <p/>... <p/>[n].gram
 * <p/>int32 order <p/>int32 count <p/>int32... id[0,..,n] <p> <p/>[n].prob <p/>int32 count
 * <p/>float32 prob <p/>.... <p> <p/>[n].backoff <p/>int32 count <p/>float32 prob <p/>.... <p>
 * <p/>vocab <p/>int32 count <p/>UTF-8... word0...n
 */
public class MultiFileUncompressedLm {

  public static String INFO_FILE_NAME = "info";
  public static String GRAM_IDS_FILE_SUFFIX = ".gram";
  public static String PROB_FILE_SUFFIX = ".prob";
  public static String BACKOFF_FILE_SUFFIX = ".backoff";
  public static String VOCAB_FILE_NAME = "vocab";

  int[] counts;
  int order;
  File dir;

  public MultiFileUncompressedLm(File dir) throws IOException {
    this.dir = dir;
    try (DataInputStream dis = new DataInputStream(new FileInputStream(getFile(INFO_FILE_NAME)))) {
      order = dis.readInt();
      counts = new int[order + 1];
      for (int i = 0; i < order; i++) {
        counts[i + 1] = dis.readInt();
      }
    }
  }

  public static MultiFileUncompressedLm generate(
      File arpaFile,
      File dir,
      String encoding,
      int fractionDigits) throws IOException {
    if (dir.exists() && !dir.isDirectory()) {
      throw new IllegalArgumentException(dir + " is not a directory!");
    } else {
      java.nio.file.Files.createDirectories(dir.toPath());
    }

    long elapsedTime = Files.asCharSource(arpaFile, Charset.forName(encoding)).readLines(
        new ArpaToBinaryConverter(dir, fractionDigits));
    Log.info("Multi file uncompressed binary model is generated in " + (double) elapsedTime / 1000d
        + " seconds");
    if (!new File(dir, INFO_FILE_NAME).exists()) {
      throw new IllegalStateException("Arpa file " + arpaFile
          + " cannot be parsed. Possibly incorrectly formatted or empty file is supplied.");
    }
    return new MultiFileUncompressedLm(dir);
  }

  public File getLmDir() {
    return dir;
  }

  public int getRankSize(File f) throws IOException {
    DataInputStream dis = new DataInputStream(new FileInputStream(f));
    int count = dis.readInt();
    dis.close();
    return count;
  }

  public int getRankBlockSize(File f) throws IOException {
    DataInputStream dis = new DataInputStream(new FileInputStream(f));
    dis.readInt();
    int blockSize = dis.readInt();
    dis.close();
    return blockSize;
  }

  public File getGramFile(int n) {
    return new File(dir, n + GRAM_IDS_FILE_SUFFIX);
  }

  public File getProbFile(int n) {
    return new File(dir, n + PROB_FILE_SUFFIX);
  }

  public File getProbRankFile(int n) {
    return new File(dir, n + PROB_FILE_SUFFIX + ".rank");
  }

  public File getBackoffRankFile(int n) {
    return new File(dir, n + BACKOFF_FILE_SUFFIX + ".rank");
  }

  public File getBackoffFile(int n) {
    return new File(dir, n + BACKOFF_FILE_SUFFIX);
  }

  public File getProbabilityLookupFile(int n) {
    return new File(dir, getProbFile(n).getName() + ".lookup");
  }

  public File getBackoffLookupFile(int n) {
    return new File(dir, getBackoffFile(n).getName() + ".lookup");
  }

  public int getOrder() {
    return order;
  }

  public int getCount(int n) {
    return counts[n];
  }

  private File getFile(String name) {
    return new File(dir, name);
  }

  public void generateRankFiles(int i, int bit, QuantizerType quantizerType) throws IOException {
    if (bit > 24) {
      throw new IllegalArgumentException(
          "Cannot generate rank file larger than 24 bits but it is:" + bit);
    }
    Log.info("Calculating probabilty rank values for :" + i + " Grams");
    File probFile = getProbFile(i);
    generateRankFile(bit, i, probFile, new File(dir, i + PROB_FILE_SUFFIX + ".rank"),
        quantizerType);
    if (i < counts.length - 1) {
      File backoffFile = getBackoffFile(i);
      Log.info("Calculating back-off rank values for :" + i + " Grams");
      generateRankFile(bit, i, backoffFile, new File(dir, i + BACKOFF_FILE_SUFFIX + ".rank"),
          quantizerType);
    }
  }

  public void generateRankFiles(int bit, QuantizerType quantizerType) throws IOException {
    if (bit > 24) {
      throw new IllegalArgumentException(
          "Cannot generate rank file larger than 24 bits but it is:" + bit);
    }
    for (int i = 1; i < counts.length; i++) {
      Log.info("Calculating probabilty lookup values for :" + i + " Grams");
      File probFile = getProbFile(i);
      generateRankFile(bit, i, probFile, new File(dir, i + PROB_FILE_SUFFIX + ".rank"),
          quantizerType);
      if (i < counts.length - 1) {
        File backoffFile = getBackoffFile(i);
        Log.info("Calculating lookup values for " + i + " Grams");
        generateRankFile(bit, i, backoffFile, new File(dir, i + BACKOFF_FILE_SUFFIX + ".rank"),
            quantizerType);
      }
    }
  }

  private void generateRankFile(int bit, int currentOrder, File probFile, File rankFile,
      QuantizerType quantizerType) throws IOException {
    try (DataInputStream dis = new DataInputStream(
        new BufferedInputStream(new FileInputStream(probFile)));
        DataOutputStream dos = new DataOutputStream(
            new BufferedOutputStream(new FileOutputStream(rankFile)))) {
      int count = dis.readInt();
      Quantizer quantizer = BinaryFloatFileReader.getQuantizer(probFile, bit, quantizerType);
      dos.writeInt(count);
      Log.info("Writing Rank file for " + currentOrder + " grams");
      int bytecount = (bit % 8 == 0 ? bit / 8 : bit / 8 + 1);
      if (bytecount == 0) {
        bytecount = 1;
      }
      dos.writeInt(bytecount);
      byte[] bytez = new byte[3];
      for (int j = 0; j < count; j++) {
        final int rank = quantizer.getQuantizationIndex(dis.readFloat());
        switch (bytecount) {
          case 1:
            dos.write(rank & 0xff);
            break;
          case 2:
            dos.writeShort(rank & 0xffff);
            break;
          case 3:
            bytez[0] = (byte) ((rank >>> 16) & 0xff);
            bytez[1] = (byte) ((rank >>> 8) & 0xff);
            bytez[2] = (byte) (rank & 0xff);
            dos.write(bytez);
            break;
        }
      }

      DoubleLookup lookup = quantizer.getDequantizer();
      Log.info("Writing lookups for " + currentOrder + " grams. Size= " + lookup.getRange());
      lookup.save(new File(dir, probFile.getName() + ".lookup"));

    }
  }

  public File getVocabularyFile() {
    return new File(dir, VOCAB_FILE_NAME);
  }

  private static class ArpaToBinaryConverter implements LineProcessor<Long> {

    public static final int DEFAULT_UNKNOWN_PROBABILTY = -20;
    int ngramCounter = 0;
    int _n;

    double fractionMultiplier;
    State state = State.BEGIN;
    List<Integer> ngramCounts = new ArrayList<>();
    boolean started = false;
    File dir;
    DataOutputStream gramOs;
    DataOutputStream probOs;
    DataOutputStream backoOffs;
    int order;
    long start;
    SpaceTabTokenizer tokenizer = new SpaceTabTokenizer();
    LmVocabulary.Builder vocabularyBuilder = new LmVocabulary.Builder();
    // This will be generated after reading unigrams.
    LmVocabulary lmVocabulary;
    ArpaToBinaryConverter(File dir, int fractionDigitCount)
        throws FileNotFoundException {
      Log.info("Generating multi file uncompressed language model from Arpa file in directory: %s",
          dir.getAbsolutePath());
      this.dir = dir;
      if (fractionDigitCount >= 0) {
        fractionMultiplier = Math.pow(10, fractionDigitCount);
      } else {
        fractionMultiplier = 0;
      }
      start = System.currentTimeMillis();
    }

    private float reduceFraction(float input) {
      if (fractionMultiplier != 0) {
        return (float) (Math.round(input * fractionMultiplier) / fractionMultiplier);
      } else {
        return input;
      }
    }

    private void newGramStream(int n) throws IOException {
      if (gramOs != null) {
        gramOs.close();
      }
      gramOs = getDos(n + GRAM_IDS_FILE_SUFFIX);
      gramOs.writeInt(n);
      gramOs.writeInt(ngramCounts.get(n - 1));
    }

    private void newProbStream(int n) throws IOException {
      if (probOs != null) {
        probOs.close();
      }
      probOs = getDos(n + PROB_FILE_SUFFIX);
      probOs.writeInt(ngramCounts.get(n - 1));
    }

    private void newBackoffStream(int n) throws IOException {
      if (backoOffs != null) {
        backoOffs.close();
      }
      backoOffs = getDos(n + BACKOFF_FILE_SUFFIX);
      backoOffs.writeInt(ngramCounts.get(n - 1));
    }

    public boolean processLine(String s) throws IOException {
      String clean = s.trim();
      switch (state) {
        // read n value and ngram counts.
        case BEGIN:
          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) {
            state = State.UNIGRAMS;
            newGramStream(1);
            newProbStream(1);
            newBackoffStream(1);
            Log.info("Gram counts in Arpa file: " + Joiner.on(" ").join(ngramCounts));
            Log.info("Writing unigrams.");
            _n++;
          }
          break;

        // read ngrams. if unigram values, we store the strings and related indexes.
        case UNIGRAMS:
          if (clean.length() == 0 || clean.startsWith("\\")) {
            break;
          }
          String[] tokens = tokenizer.split(clean);
          // parse probabilty
          float logProbability = Float.parseFloat(tokens[0]);

          String word = tokens[1];
          float logBackoff = 0;
          if (tokens.length == 3) {
            logBackoff = Float.parseFloat(tokens[_n + 1]);
          }
          // write unigram id, log-probability and log-backoff value.
          int wordIndex = vocabularyBuilder.add(word);
          gramOs.writeInt(wordIndex);
          probOs.writeFloat(reduceFraction(logProbability));

          // if there are only ngrams, do not write backoff value.
          if (ngramCounts.size() > 1) {
            backoOffs.writeFloat(reduceFraction(logBackoff));
          }

          ngramCounter++;
          if (ngramCounter == ngramCounts.get(0)) {
            handleSpecialToken("<unk>");
            handleSpecialToken("</s>");
            handleSpecialToken("<s>");
            lmVocabulary = vocabularyBuilder.generate();

            ngramCounts.set(0, lmVocabulary.size());

            // we write info file after reading unigrams because we may add special tokens to unigrams
            // so count information may have been changed.
            order = ngramCounts.size();
            try (DataOutputStream infos = getDos(INFO_FILE_NAME)) {
              infos.writeInt(order);
              for (Integer ngramCount : ngramCounts) {
                infos.writeInt(ngramCount);
              }
            }

            ngramCounter = 0;
            state = State.NGRAMS;
            _n++;
            // if there is only unigrams in the arpa file, exit
            if (ngramCounts.size() == 1) {
              state = State.VOCABULARY;
            } else {
              newGramStream(2);
              newProbStream(2);
              if (order > 2) {
                newBackoffStream(2);
              }
              Log.info("Writing 2-grams.");
            }
          }
          break;

        case NGRAMS:
          if (clean.length() == 0 || clean.startsWith("\\")) {
            break;
          }
          tokens = tokenizer.split(clean);
          logProbability = Float.parseFloat(tokens[0]);

          for (int i = 0; i < _n; i++) {
            int id = lmVocabulary.indexOf(tokens[i + 1]);
            gramOs.writeInt(id);
          }

          // probabilities

          probOs.writeFloat(reduceFraction(logProbability));
          if (_n < ngramCounts.size()) {
            logBackoff = 0;
            if (tokens.length == _n + 2) {
              logBackoff = Float.parseFloat(tokens[_n + 1]);
            }
            backoOffs.writeFloat(reduceFraction(logBackoff));
          }

          if (ngramCounter > 0 && ngramCounter % 1000000 == 0) {
            Log.info(ngramCounter + " grams are written so far.");
          }

          ngramCounter++;
          if (ngramCounter == ngramCounts.get(_n - 1)) {
            ngramCounter = 0;
            // if there is no more ngrams, exit
            if (ngramCounts.size() == _n) {
              state = State.VOCABULARY;
            } else {
              _n++;
              newGramStream(_n);
              newProbStream(_n);
              if (order > _n) {
                newBackoffStream(_n);
              }
              Log.info("Writing " + _n + "-grams.");
            }
          }
          break;

        case VOCABULARY:
          Closeables.close(gramOs, true);
          Closeables.close(probOs, true);
          Closeables.close(backoOffs, true);
          Log.info("Writing model vocabulary.");
          lmVocabulary.saveBinary(new File(dir, VOCAB_FILE_NAME));
          return false; // we are done.
      }
      return true;
    }

    // adds undefined specials token with default probability.
    private void handleSpecialToken(String word) throws IOException {
      if (vocabularyBuilder.indexOf(word) == -1
          && vocabularyBuilder.indexOf(word.toUpperCase()) == -1) {
        Log.warn("Special token " + word +
            " does not exist in model. It is added with default [unknown word] probability: " +
            DEFAULT_UNKNOWN_PROBABILTY);
        int index = vocabularyBuilder.add(word);
        gramOs.writeInt(index);
        probOs.writeFloat(DEFAULT_UNKNOWN_PROBABILTY);
        if (ngramCounts.size() > 1) {
          backoOffs.writeFloat(0);
        }
      }
    }

    private DataOutputStream getDos(String name) throws FileNotFoundException {
      return new DataOutputStream(
          new BufferedOutputStream(new FileOutputStream(new File(dir, name)), 100000));
    }

    public Long getResult() {
      // just return the time..
      return System.currentTimeMillis() - start;
    }

    enum State {
      BEGIN, UNIGRAMS, NGRAMS, VOCABULARY
    }
  }
}
