package zemberek.lm.compression;

import com.google.common.io.Files;
import java.io.BufferedOutputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.BitSet;
import zemberek.core.hash.IntHashKeyProvider;
import zemberek.core.hash.LargeNgramMphf;
import zemberek.core.hash.Mphf;
import zemberek.core.hash.MultiLevelMphf;
import zemberek.core.logging.Log;
import zemberek.core.quantization.QuantizerType;

public class UncompressedToSmoothLmConverter {

  private static final int VERSION = 1;
  File lmFile;
  File tempDir;

  int order;

  public UncompressedToSmoothLmConverter(File lmFile, File tempDir) {
    this.lmFile = lmFile;
    this.tempDir = tempDir;
  }

  public void convertSmall(File binaryUncompressedLmDir, NgramDataBlock block) throws IOException {
    convert(binaryUncompressedLmDir, block, SmoothLm.MphfType.SMALL, null, -1);
  }

  public void convertLarge(File binaryUncompressedLmDir, NgramDataBlock block, int chunkBits)
      throws IOException {
    convert(binaryUncompressedLmDir, block, SmoothLm.MphfType.LARGE, null, chunkBits);
  }

  public void convertLarge(File binaryUncompressedLmDir, NgramDataBlock block,
      File[] oneBasedMphfFiles, int chunkBits) throws IOException {
    convert(binaryUncompressedLmDir, block, SmoothLm.MphfType.LARGE, oneBasedMphfFiles, chunkBits);
  }

  private void convert(File binaryUncompressedLmDir,
      NgramDataBlock block,
      SmoothLm.MphfType type,
      File[] oneBasedMphfFiles,
      int chunkBits) throws IOException {

    Log.info("Generating compressed language model.");

    MultiFileUncompressedLm lm = new MultiFileUncompressedLm(binaryUncompressedLmDir);

    lm.generateRankFiles(block.probabilitySize * 8, QuantizerType.BINNING);

    order = lm.order;

    DataOutputStream dos = new DataOutputStream(
        new BufferedOutputStream(new FileOutputStream(lmFile)));

    // generate Minimal Perfect Hash functions for 2,3...n grams and save them as separate files.

    File[] phfFiles = new File[order + 1];
    if (oneBasedMphfFiles != null) {
      phfFiles = oneBasedMphfFiles;
    } else {
      for (int i = 2; i <= order; i++) {
        Mphf mphf;
        if (type == SmoothLm.MphfType.LARGE) {
          mphf = LargeNgramMphf.generate(lm.getGramFile(i), chunkBits);
        } else {
          mphf = MultiLevelMphf.generate(lm.getGramFile(i));
        }
        Log.info("MPHF is generated for order %d with %d keys. Average bits per key: %.3f",
            i,
            mphf.size(),
            mphf.averageBitsPerKey());
        File mphfFile = new File(tempDir, lmFile.getName() + i + "gram.mphf");
        phfFiles[i] = mphfFile;
        mphf.serialize(mphfFile);
      }
    }
    // generate header.
    Log.info("Writing header");
    // write version and type info
    dos.writeInt(VERSION);

    // write Mphf type
    if (type == SmoothLm.MphfType.SMALL) {
      dos.writeInt(0);
    } else {
      dos.writeInt(1);
    }

    // write log-base
    dos.writeDouble(10d);
    // write n value for grams (3 for trigram model)
    dos.writeInt(order);
    // write counts, generate gramdata

    for (int i = 1; i <= order; i++) {
      dos.writeInt(lm.getCount(i));
    }

    // write rank lookup data (contains size+doubles)
    for (int i = 1; i <= order; i++) {
      Files.copy(lm.getProbabilityLookupFile(i), dos);
    }
    for (int i = 1; i <= order; i++) {
      if (i < order) {
        Files.copy(lm.getBackoffLookupFile(i), dos);
      }
    }

    Log.info("Reordering probability data and saving it together with n-gram fingerprints");
    for (int i = 1; i <= order; i++) {
      InMemoryBigByteArray probData = new InMemoryBigByteArray(lm.getProbRankFile(i));
      InMemoryBigByteArray backoffData = null;

      if (i < order) {
        backoffData = new InMemoryBigByteArray(lm.getBackoffRankFile(i));
      }

      ReorderData reorderData;
      final int gramCount = probData.count;
      if (i == 1) {
        int reorderedIndexes[] = new int[gramCount];
        for (int j = 0; j < gramCount; j++) {
          reorderedIndexes[j] = j;
        }
        reorderData = new ReorderData(reorderedIndexes, new int[0]);
      } else {
        if (type == SmoothLm.MphfType.LARGE) {
          reorderData = reorderIndexes(block, lm, i, LargeNgramMphf.deserialize(phfFiles[i]));
        } else {
          reorderData = reorderIndexes(block, lm, i, MultiLevelMphf.deserialize(phfFiles[i]));
        }
      }
      Log.info("Validating reordered index array for order: %d", i);

      validateIndexArray(reorderData.reorderedKeyIndexes);

      int fingerPrintSize = block.fingerPrintSize;
      if (i == 1) {
        fingerPrintSize = 0;
      }
      int backOffSize = block.backoffSize;
      if (i == order) {
        backOffSize = 0;
      }

      dos.writeInt(gramCount);
      dos.writeInt(fingerPrintSize);
      dos.writeInt(block.probabilitySize);
      dos.writeInt(backOffSize);

      byte[] probBuff = new byte[block.probabilitySize];
      byte[] fpBuff = new byte[fingerPrintSize];
      byte[] backoffBuff = new byte[backOffSize];

      for (int k = 0; k < gramCount; k++) {
        // save fingerprint values for 2,3,.. grams.
        if (i > 1) {
          block.fingerprintAsBytes(reorderData.fingerprints[k], fpBuff);
          dos.write(fpBuff);
        }
        probData.get(reorderData.reorderedKeyIndexes[k], probBuff);
        dos.write(probBuff);
        // write backoff value if exists.
        if (backoffData != null) {
          backoffData.get(reorderData.reorderedKeyIndexes[k], backoffBuff);
          dos.write(backoffBuff);
        }
      }
    }

    // append size of the Perfect hash and its content.
    if (phfFiles.length > 0) {
      Log.info("Saving MPHF values.");
    }

    for (int i = 2; i <= order; i++) {
      Files.copy(phfFiles[i], dos);
    }

    // save vocabulary
    Log.info("Saving vocabulary.");
    Files.copy(lm.getVocabularyFile(), dos);

    dos.close();

  }

  private void validateIndexArray(int[] arr) {
    BitSet set = new BitSet();
    for (int i : arr) {
      if (i >= arr.length || i < 0) {
        throw new IllegalStateException(
            "array contains a value=" + i + " larger than the array size=" + arr.length);
      }
      set.set(i);
    }
    for (int i = 0; i < arr.length; i++) {
      if (!set.get(i)) {
        throw new IllegalStateException("Not validated.");
      }
    }
  }

  /**
   * This class does the following: suppose we have the keys as: [k0, k1, k2, k3, k4, k5] ->
   * [0,1,2,3,4,5] their mphf values are however: k0=2, k1=5, k2=0, k3=4, k4=1, k5=3 So what we want
   * is to have key indexes in their minimal perfect has index order. reordered keys indexes: [k2,
   * k4, k0, k5, k3, k1] -> [2,4,0,5,3,1]
   *
   * @param lm multifile language language model.
   * @param _order current order of language model
   * @param mphf MPH function
   * @return reordered key indexes and those keys fingerprint values.
   */
  private ReorderData reorderIndexes(NgramDataBlock block, MultiFileUncompressedLm lm, int _order,
      Mphf mphf) throws IOException {
    ChunkingNGramReader reader = new ChunkingNGramReader(lm.getGramFile(_order), _order, 1000000);
    int[] reorderedIndexes = new int[lm.getCount(_order)];
    int[] fingerPrints = new int[lm.getCount(_order)];
    int counter = 0;
    for (IntHashKeyProvider provider : reader) {
      for (int k = 0; k < provider.keyAmount(); k++) {
        final int[] key = provider.getKey(k);
        final int hashVal = mphf.get(key);
        reorderedIndexes[hashVal] = counter;
        fingerPrints[hashVal] = block.fingerprint(key);
        counter++;
      }
    }
    return new ReorderData(reorderedIndexes, fingerPrints);
  }

  public static class NgramDataBlock {

    int fingerPrintSize;
    int probabilitySize;
    int backoffSize;
    int fingerprintMask;

    public NgramDataBlock(int fingerPrintBits, int probabilityBits, int backoffBits) {
      if (fingerPrintBits % 8 != 0) {
        throw new IllegalArgumentException("FingerPrint bit size must be an order of 8");
      }
      if (probabilityBits % 8 != 0) {
        throw new IllegalArgumentException("Probability bit size must be an order of 8");
      }
      if (backoffBits % 8 != 0) {
        throw new IllegalArgumentException("Backoff bit size must be an order of 8");
      }
      this.fingerPrintSize = fingerPrintBits / 8;
      this.probabilitySize = probabilityBits / 8;
      this.backoffSize = backoffBits / 8;
      if (fingerPrintBits == 4) {
        this.fingerprintMask = 0xffffffff;
      } else {
        this.fingerprintMask = (1 << fingerPrintBits) - 1;
      }
    }

    int fingerprint(int[] key) {
      return MultiLevelMphf.hash(key, -1) & fingerprintMask;
    }

    void fingerprintAsBytes(int fingerprint, byte[] res) {
      int k = 0;
      switch (fingerPrintSize) {
        case 1:
          res[k] = (byte) (fingerprint & 0xff);
          break;
        case 2:
          res[k] = (byte) ((fingerprint >>> 8) & 0xff);
          res[k + 1] = (byte) (fingerprint & 0xff);
          break;
        case 3:
          res[k] = (byte) ((fingerprint >>> 16) & 0xff);
          res[k + 1] = (byte) ((fingerprint >>> 8) & 0xff);
          res[k + 2] = (byte) (fingerprint & 0xff);
          break;
        case 4:
          res[k] = (byte) ((fingerprint >>> 24) & 0xff);
          res[k + 1] = (byte) ((fingerprint >>> 16) & 0xff);
          res[k + 2] = (byte) ((fingerprint >>> 8) & 0xff);
          res[k + 3] = (byte) (fingerprint & 0xff);
          break;
      }
    }

    void fingerprintAsBytes(int[] key, byte[] res) {
      fingerprintAsBytes(MultiLevelMphf.hash(key, -1) & fingerprintMask, res);
    }
  }

  private class ReorderData {

    int[] reorderedKeyIndexes;
    int[] fingerprints;

    private ReorderData(int[] reorderedKeyIndexes, int[] fingerprints) {
      this.reorderedKeyIndexes = reorderedKeyIndexes;
      this.fingerprints = fingerprints;
    }
  }
}
