package zemberek.embedding;

import java.io.BufferedInputStream;
import java.io.BufferedOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.PrintWriter;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.PriorityQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import zemberek.core.logging.Log;
import zemberek.lm.LmVocabulary;

public class WordDistances {

  public final String source;
  Distance[] distances;

  public WordDistances(String source, Distance[] distList) {
    this.source = source;
    this.distances = distList;
  }

  public static void saveDistanceListBin(
      Path vectorFile,
      Path outFile,
      Path vocabFile,
      int distanceAmount,
      int blockSize,
      int threadSize) throws Exception {
    Log.info("Loading vectors.");
    List<WordVector> wordVectors = WordVector.loadFromBinary(vectorFile);

    Log.info("Writing vocabulary.");
    // write vocabulary.
    List<String> words = new ArrayList<>(wordVectors.size());
    wordVectors.forEach(s -> words.add(s.word));
    LmVocabulary vocabulary = new LmVocabulary(words);
    vocabulary.saveBinary(vocabFile.toFile());

    Log.info("Calculating distances.");
    // create a thread pool executor
    ExecutorService es = Executors.newFixedThreadPool(threadSize);
    CompletionService<List<BlockUnit>> completionService = new ExecutorCompletionService<>(es);

    int blockCounter = 0;
    for (int i = 0; i < wordVectors.size(); i += blockSize) {
      int endIndex = i + blockSize >= wordVectors.size() ? wordVectors.size() : i + blockSize;
      completionService
          .submit(new BlockDistanceTask(wordVectors, distanceAmount, blockSize, i, endIndex));
      blockCounter++;
    }
    es.shutdown();
    List<WordDistances> distancesToWrite = new ArrayList<>(wordVectors.size());
    int i = 0;
    while (i < blockCounter) {
      List<BlockUnit> units = completionService.take().get();
      for (BlockUnit unit : units) {
        String source = unit.vector.word;
        List<Distance> distList = new ArrayList<>(unit.distQueue);
        Collections.sort(distList);
        Collections.reverse(distList);
        distancesToWrite
            .add(new WordDistances(source, distList.toArray(new Distance[distList.size()])));
      }
      i++;
      if ((i * blockSize % 10) == 0) {
        Log.info("%d of %d completed", i * blockSize, wordVectors.size());
      }
    }

    Log.info("Writing.");
    try (DataOutputStream dos = new DataOutputStream(
        new BufferedOutputStream(new FileOutputStream(outFile.toFile())))) {
      dos.writeInt(wordVectors.size());
      dos.writeInt(distanceAmount);
      for (WordDistances w : distancesToWrite) {
        dos.writeInt(vocabulary.indexOf(w.source));
        for (Distance token : w.distances) {
          dos.writeInt(vocabulary.indexOf(token.word));
          dos.writeFloat(token.distance);
        }
      }
    }
  }

  public static List<WordDistances> readFromBinary(Path binFile, Path vocabFile)
      throws IOException {
    LmVocabulary vocabulary = LmVocabulary.loadFromBinary(vocabFile.toFile());
    try (DataInputStream in = new DataInputStream(
        new BufferedInputStream(new FileInputStream(binFile.toFile()), 100000))) {
      int wordSize = in.readInt();
      int vectorSize = in.readInt();
      List<WordDistances> distLists = new ArrayList<>(wordSize);
      for (int i = 0; i < wordSize; i++) {
        String s = vocabulary.getWord(in.readInt());
        Distance[] distances = new Distance[vectorSize];
        for (int j = 0; j < vectorSize; j++) {
          String word = vocabulary.getWord(in.readInt());
          Float f = in.readFloat();
          distances[j] = new Distance(word, f);
        }
        if (i % 10000 == 0) {
          Log.info("%d completed", i);
        }
        distLists.add(new WordDistances(s, distances));
      }
      return distLists;
    }
  }

  public static void writeToTxt(Path in, Path vocab, Path out, int size) throws IOException {
    try (PrintWriter pw = new PrintWriter(out.toFile())) {
      for (WordDistances distList : WordDistances.readFromBinary(in, vocab)) {
        pw.println(distList.source);
        for (int k = 0; k < size; k++) {
          pw.println(distList.distances[k]);
        }
      }
    }
  }

  public static class Distance implements Comparable<Distance> {

    public final String word;
    public final float distance;

    public Distance(String word, float distance) {
      this.word = word;
      this.distance = distance;
    }

    public String toString() {
      return word + " " + distance;
    }

    @Override
    public int compareTo(Distance o) {
      return Float.compare(distance, o.distance);
    }
  }

  private static class BlockUnit {

    public final WordVector vector;
    PriorityQueue<Distance> distQueue = new PriorityQueue<>();

    public BlockUnit(WordVector vector) {
      this.vector = vector;
    }
  }

  private static class BlockDistanceTask implements Callable<List<BlockUnit>> {

    List<WordVector> wordVectors;
    int distanceAmount;
    int blockSize;
    int blockBeginIndex;
    int endIndex;

    private BlockDistanceTask(
        List<WordVector> wordVectors,
        int distanceAmount,
        int blockSize,
        int blockBeginIndex,
        int endIndex) {
      this.wordVectors = wordVectors;
      this.distanceAmount = distanceAmount;
      this.blockSize = blockSize;
      this.blockBeginIndex = blockBeginIndex;
      this.endIndex = endIndex;
    }

    @Override
    public List<BlockUnit> call() throws Exception {
      List<BlockUnit> units = new ArrayList<>();
      for (int j = blockBeginIndex; j < endIndex; j++) {
        WordVector vector = wordVectors.get(j);
        units.add(new BlockUnit(vector));
      }
      for (int k = 0; k < wordVectors.size(); k += blockSize) {
        int end = (k + blockSize) >= wordVectors.size() ? wordVectors.size() : (k + blockSize);
        for (BlockUnit unit : units) {
          for (int j = k; j < end; j++) {
            WordVector v = wordVectors.get(j);
            // skip self.
            if (v == unit.vector) {
              continue;
            }
            float distance = unit.vector.cosDistance(v);
            if (unit.distQueue.size() < distanceAmount) {
              unit.distQueue.add(new Distance(v.word, distance));
            } else {
              Distance weakest = unit.distQueue.peek();
              if (weakest.distance < distance) {
                unit.distQueue.remove();
                unit.distQueue.add(new Distance(v.word, distance));
              }
            }
          }
        }
      }
      return units;
    }
  }

}
