package zemberek.core.hash;

import com.google.common.io.Files;
import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.Closeable;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.RandomAccessFile;
import zemberek.core.logging.Log;

/**
 * This is a MPHF implementation suitable for very large key sets.
 */
public class LargeNgramMphf implements Mphf {

  private static final int DEFAULT_CHUNK_SIZE_IN_BITS = 22;
  final int maxBitMask;
  final int bucketMask;
  final int pageShift;

  MultiLevelMphf mphfs[];
  int[] offsets;

  public LargeNgramMphf(int maxBitMask, int bucketMask, int pageShift, MultiLevelMphf[] mphfs,
      int[] offsets) {
    this.maxBitMask = maxBitMask;
    this.bucketMask = bucketMask;
    this.pageShift = pageShift;
    this.mphfs = mphfs;
    this.offsets = offsets;
  }

  /**
   * Same as generate(File file, int chunkBits) but uses DEFAULT_CHUNK_SIZE_IN_BITS for chunk size.
   *
   * @param file binary key file
   * @return generated LargeNgramMphf
   */
  public static LargeNgramMphf generate(File file) throws IOException {
    return generate(file, DEFAULT_CHUNK_SIZE_IN_BITS);
  }

  /**
   * Generates MPHF from a binary integer key file. File needs to be in this structure: int32 order
   * (how many integers each key) int32 amount of keys . Max is 2^31-1 int32... key1 int32... key2
   * <p/>
   * Keys in the file must be unique. During generation of PHF system does not check for uniqueness.
   * System does the following:
   * <p/>
   * it splits the total amount of keys to large chunks (2^23 ~ 8 million keys by default.) during
   * split operation, all keys are divided to a chunk file according to it's bucket index. this is
   * calculated with PHF algorithms generic hash function g() Key counts in each chunk may be
   * different. But they are generally close values. These count values are important Because they
   * are used during global PHF value. Values are stored in an array (offsets)
   * <p/>
   * After files chunks are generated, a PHF is calculated for each chunk. And they are stored in an
   * array
   *
   * @param file binary key file
   * @return LargeNgramMphf fro the keys in the file
   * @throws IOException If an error occurs during file access.
   */
  public static LargeNgramMphf generate(File file, int chunkBits) throws IOException {
    File tmp = Files.createTempDir();
    Splitter splitter = new Splitter(file, tmp, chunkBits);
    Log.info("Gram count: " + splitter.gramCount);
    Log.info("Segment count: " + splitter.pageCount);
    Log.info("Average segment size: " + (1 << splitter.pageBit));
    Log.info("Segmenting File...");
    splitter.split();
    int bucketBits = splitter.pageBit - 2;
    if (bucketBits <= 0) {
      bucketBits = 1;
    }
    MultiLevelMphf[] mphfs = new MultiLevelMphf[splitter.pageCount];
    int[] offsets = new int[splitter.pageCount];
    int total = 0;
    for (int i = 0; i < splitter.pageCount; i++) {
      final ByteGramProvider keySegment = splitter.getKeySegment(i);
      Log.debug("Segment key count: " + keySegment.keyAmount());
      Log.debug("Segment bucket ratio: " + ((double) keySegment.keyAmount() / (1 << bucketBits)));
      total += keySegment.keyAmount();
      MultiLevelMphf mphf = MultiLevelMphf.generate(keySegment);
      Log.info("MPHF is generated for segment %d with %d keys. Average bits per key: %.3f",
          i,
          mphf.size(),
          mphf.averageBitsPerKey());
      mphfs[i] = mphf;
      if (i > 0) {
        offsets[i] = offsets[i - 1] + mphfs[i - 1].size();
      }
    }
    Log.debug("Total processed keys:" + total);
    int maxMask = (1 << splitter.maxBit) - 1;
    int bucketMask = (1 << bucketBits) - 1;
    return new LargeNgramMphf(maxMask, bucketMask, splitter.pageShift, mphfs, offsets);
  }

  /**
   * A custom deserializer.
   *
   * @param file file that contains serialized data.
   * @param skip amoutn to skip
   * @return a new FastMinimalPerfectHash object.
   * @throws IOException if an error occurs during file access.
   */
  public static LargeNgramMphf deserialize(File file, long skip) throws IOException {
    try (DataInputStream dis = new DataInputStream(
        new BufferedInputStream(new FileInputStream(file), 1000000))) {
      long actualSkip = dis.skip(skip);
      if (actualSkip != skip) {
        throw new IOException("Cannot skip necessary amount of bytes from stream:" + skip);
      }
      return deserialize(dis);
    }
  }

  /**
   * A custom deserializer.
   *
   * @param file file that contains serialized data.
   * @return a new FastMinimalPerfectHash object.
   * @throws IOException if an error occurs during file access.
   */
  public static LargeNgramMphf deserialize(File file) throws IOException {
    return deserialize(file, 0);
  }

  /**
   * A custom deserializer. Look serialization method document for the format.
   *
   * @param dis DataInputStream that contains serialized data.
   * @return a new ChdPerfectHash object.
   * @throws IOException if an error occurs during stream access.
   */
  public static LargeNgramMphf deserialize(DataInputStream dis) throws IOException {
    int maxBitMask = dis.readInt();
    int bucketMask = dis.readInt();
    int pageShift = dis.readInt();
    int phfCount = dis.readInt();

    int[] offsets = new int[phfCount];
    for (int i = 0; i < offsets.length; i++) {
      offsets[i] = dis.readInt();
    }
    MultiLevelMphf[] hashes = new MultiLevelMphf[phfCount];
    for (int i = 0; i < offsets.length; i++) {
      hashes[i] = MultiLevelMphf.deserialize(dis);
    }
    return new LargeNgramMphf(maxBitMask, bucketMask, pageShift, hashes, offsets);
  }

  public int get(int[] ngram) {
    final int hash = MultiLevelMphf.hash(ngram, -1);
    final int pageIndex = (hash & maxBitMask) >>> pageShift;
    return mphfs[pageIndex].get(ngram, hash) + offsets[pageIndex];
  }

  public int get(int[] ngram, int hash) {
    final int pageIndex = (hash & maxBitMask) >>> pageShift;
    return mphfs[pageIndex].get(ngram, hash) + offsets[pageIndex];
  }

  public int get(int g1, int g2, int g3, int hash) {
    final int pageIndex = (hash & maxBitMask) >>> pageShift;
    return mphfs[pageIndex].get(g1, g2, g3, hash) + offsets[pageIndex];
  }

  public int get(int g1, int g2, int hash) {
    final int pageIndex = (hash & maxBitMask) >>> pageShift;
    return mphfs[pageIndex].get(g1, g2, hash) + offsets[pageIndex];
  }

  public int get(String ngram) {
    final int hash = MultiLevelMphf.hash(ngram, -1);
    final int pageIndex = (hash & maxBitMask) >>> pageShift;
    return mphfs[pageIndex].get(ngram, hash) + offsets[pageIndex];
  }

  public int get(String ngram, int hash) {
    final int pageIndex = (hash & maxBitMask) >>> pageShift;
    return mphfs[pageIndex].get(ngram, hash) + offsets[pageIndex];
  }

  public int get(int[] ngram, int begin, int end, int hash) {
    final int pageIndex = (hash & maxBitMask) >>> pageShift;
    return mphfs[pageIndex].get(ngram, begin, end, hash) + offsets[pageIndex];
  }

  /**
   * A custom serializer.
   *
   * @param file file to serialize data.
   * @throws IOException if an error occurs during file access.
   */
  public void serialize(File file) throws IOException {
    serialize(new BufferedOutputStream(new FileOutputStream(file), 1000000));
  }

  public void serialize(OutputStream os) throws IOException {
    try (DataOutputStream dos = new DataOutputStream(os)) {
      dos.writeInt(this.maxBitMask);
      dos.writeInt(this.bucketMask);
      dos.writeInt(this.pageShift);
      dos.writeInt(this.mphfs.length);
      for (int offset : offsets) {
        dos.writeInt(offset);
      }
      for (MultiLevelMphf mphf : mphfs) {
        mphf.serialize(dos);
      }
    }
  }

  @Override
  public double averageBitsPerKey() {
    double total = 0;
    for (MultiLevelMphf mphf : mphfs) {
      total = total + mphf.averageBitsPerKey();
    }
    return mphfs.length > 0 ? total / mphfs.length : 0;
  }

  @Override
  public int size() {
    int size = 0;
    for (MultiLevelMphf mphf : mphfs) {
      size += mphf.size();
    }
    return size;
  }

  private static class Splitter {

    final int pageShift;
    final int pageMask;
    File gramFile;
    File[] files;
    int order;
    int gramCount;
    File tmpDir;
    int maxBit;
    int pageCount;
    int pageBit;

    public Splitter(File gramFile, File tmpdir, int pageBit) throws IOException {
      this.gramFile = gramFile;
      RandomAccessFile raf = new RandomAccessFile(gramFile, "r");
      this.order = raf.readInt();
      this.gramCount = raf.readInt();
      raf.close();
      this.tmpDir = tmpdir;
      this.pageBit = pageBit;

      // maxbit ıs the x value where 2^(x-1)<gramCount<2^x
      maxBit = getBitCountHigher(gramCount);
      if (maxBit == 0) {
        maxBit = 1;
      }

      // if the x is smaller than the default page bit, we use the x value for the page bit.
      if (maxBit < this.pageBit) {
        this.pageBit = maxBit;
        pageShift = maxBit;
      } else {
        pageShift = this.pageBit;
      }

      int pageLength = 1 << this.pageBit;
      pageMask = pageLength - 1;
      pageCount = 1 << (maxBit - this.pageBit);
      files = new File[pageCount];
    }

    int getBitCountHigher(int num) {
      int i = 1;
      int c = 0;
      while (i < num) {
        i *= 2;
        c++;
      }
      return c;
    }

    public int keyAmount() {
      return gramCount;
    }

    public void split() throws IOException {

      if (pageCount == 1) {
        files[0] = gramFile;
        return;
      }

      FileKeyWriter[] fileKeyWriters = new FileKeyWriter[pageCount];
      int[] counts = new int[pageCount];

      for (int i = 0; i < pageCount; i++) {
        files[i] = new File(tmpDir, order + "-gramidfile" + i + ".batch");
        fileKeyWriters[i] = new FileKeyWriter(files[i], order);
      }

      byte[] buffer = new byte[(1 << pageBit) * 4 * order];
      RandomAccessFile raf = new RandomAccessFile(gramFile, "r");
      raf.skipBytes(8);
      int actual;
      while ((actual = raf.read(buffer)) > 0) {
        if (actual % (order * 4) != 0) {
          throw new IllegalStateException("Cannot read order*4 aligned bytes from:" + gramFile);
        }

        int[] gramIds = new int[order];
        int p = 0;
        int maxBitMask = (1 << maxBit) - 1;
        for (int i = 0; i < actual; i += 4) {
          gramIds[p++] =
              (buffer[i] & 0xff) << 24 | (buffer[i + 1] & 0xff) << 16 | (buffer[i + 2] & 0xff) << 8
                  | (buffer[i + 3] & 0xff);
          if (p == order) {
            int hash = MultiLevelMphf.hash(gramIds, -1) & maxBitMask;
            int segmentId = hash >>> pageShift;
            fileKeyWriters[segmentId].write(gramIds);
            counts[segmentId]++;
            p = 0;
          }
        }
      }
      raf.close();
      int i = 0;
      for (FileKeyWriter writer : fileKeyWriters) {
        writer.close();
        writer.changeCount(counts[i++]);
      }
    }

    public ByteGramProvider getKeySegment(int segment) throws IOException {
      return new ByteGramProvider(files[segment]);
    }
  }

  private static class FileKeyWriter implements Closeable {

    DataOutputStream dos;
    File file;
    int order;

    private FileKeyWriter(File file, int order) throws IOException {
      this.file = file;
      dos = new DataOutputStream(new BufferedOutputStream(new FileOutputStream(file), 1000000));
      dos.writeInt(order);
      dos.writeInt(0);
      this.order = order;
    }

    void write(int[] key) throws IOException {
      for (int i : key) {
        dos.writeInt(i);
      }
    }

    void changeCount(int count) throws IOException {
      try (RandomAccessFile rafw = new RandomAccessFile(file, "rw")) {
        rafw.writeInt(order);
        rafw.writeInt(count);
      }
    }

    public void close() throws IOException {
      dos.close();
    }
  }
}