package com.github.megallo.markoverator.bigrammer;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.github.megallo.markoverator.utils.Pair;

import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

import static com.github.megallo.markoverator.bigrammer.Bigrammer.DELIM;

public class BigramModelBuilder {

    /**
     * Call me first!
     * Initialize this object with a model based on the provided sentences.
     *
     * @param sentencesList a list of sentences: each sentence is pre-tokenized, usually into words
     */
    public static BigramModel buildModel(List<List<String>> sentencesList) {
        HashMap<Pair, List<String>> forwardCache = new HashMap<>();
        HashMap<Pair, List<String>> backwardCache = new HashMap<>();
        List<String> fullWordList = new ArrayList<>();

        // add sentence delimiters to get more natural sentence starts and ends
        for (List<String> oneSentence : sentencesList) {
            fullWordList.add(DELIM);
            fullWordList.addAll(oneSentence);
        }
        fullWordList.add(DELIM); // don't forget the one at the end

        // for each triplet
        //   map of (<w1, w2> -> w3) = generates forward text
        //   map of (<w2, w3> -> w1) = generates backward text

        for (int i = 0; i < fullWordList.size() - 2; i++) {
            String w1 = fullWordList.get(i);
            String w2 = fullWordList.get(i + 1);
            String w3 = fullWordList.get(i + 2);

            Pair forwardPair = new Pair(w1, w2);
            Pair backwardPair = new Pair(w2, w3);

            if (!forwardCache.containsKey(forwardPair)) {
                forwardCache.put(forwardPair, new ArrayList<String>());
            }
            forwardCache.get(forwardPair).add(w3);

            if (!backwardCache.containsKey(backwardPair)) {
                backwardCache.put(backwardPair, new ArrayList<String>());
            }
            backwardCache.get(backwardPair).add(w1);
        }

        return new BigramModel(fullWordList, forwardCache, backwardCache);
    }

    /**
     * Write out the model object to a stream, e.g. a file on disk
     * @param model a BigramModel object, like the one built by buildModel()
     * @param outputStream a writeable stream, such as java.io.FileOutputStream
     */
    public static void saveModel(BigramModel model, OutputStream outputStream) {
        if (model == null) {
            throw new RuntimeException("Refusing to write empty model.");
        }
        Kryo kryo = new Kryo();
        Output output = new Output(outputStream);
        kryo.writeObject(output, model);
        output.close();
    }

    /**
     * Load a model object from a stream, e.g. a file on disk
     * @param inputStream readable stream we can read a previously built model from
     * @return a BigramModel that's ready to load into Bigrammer
     */
    public static BigramModel loadModel(InputStream inputStream) {
        Kryo kryo = new Kryo();
        Input input = new Input(inputStream);
        BigramModel model = kryo.readObject(input, BigramModel.class);
        input.close();

        return model;
    }
}

@@ -60,231 +60,48 @@
 com.
-esotericsoftware.kryo.Kryo;%0Aimport com.esotericsoftware.kryo.io.Input;%0Aimport com.esotericsoftware.kryo.io.Output;%0Aimport com.github.megallo.markoverator.utils.Pair;%0A%0Aimport java.io.InputStream;%0Aimport java.io.OutputStream;
+github.megallo.markoverator.utils.Pair;%0A
 %0Aimp
@@ -2024,1091 +2024,7 @@
    %7D
-%0A%0A    /**%0A     * Write out the model object to a stream, e.g. a file on disk%0A     * @param model a BigramModel object, like the one built by buildModel()%0A     * @param outputStream a writeable stream, such as java.io.FileOutputStream%0A     */%0A    public static void saveModel(BigramModel model, OutputStream outputStream) %7B%0A        if (model == null) %7B%0A            throw new RuntimeException(%22Refusing to write empty model.%22);%0A        %7D%0A        Kryo kryo = new Kryo();%0A        Output output = new Output(outputStream);%0A        kryo.writeObject(output, model);%0A        output.close();%0A    %7D%0A%0A    /**%0A     * Load a model object from a stream, e.g. a file on disk%0A     * @param inputStream readable stream we can read a previously built model from%0A     * @return a BigramModel that's ready to load into Bigrammer%0A     */%0A    public static BigramModel loadModel(InputStream inputStream) %7B%0A        Kryo kryo = new Kryo();%0A        Input input = new Input(inputStream);%0A        BigramModel model = kryo.readObject(input, BigramModel.class);%0A        input.close();%0A%0A        return model;%0A    %7D
 %0A%7D%0A

Remove kryo serialization from BigramModelBuilder