package com.turbolent.aptagger;

import org.msgpack.MessagePack;
import org.msgpack.packer.Packer;
import org.msgpack.template.SetTemplate;
import org.msgpack.unpacker.Unpacker;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Path;
import java.util.*;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static org.msgpack.template.Templates.*;

public class Tagger {
    protected static final String[] START = {"-START-", "-START2-"};
    protected static final String[] END = {"-END-", "-END2-"};
    protected static final Pattern NUMBER = Pattern.compile("[0-9][0-9,.]*");

    private static boolean isNumber(String string) {
        return NUMBER.matcher(string).matches();
    }

    private static String normalize(String word) {
        if (isNumber(word))
            return "!NUMBER";
        return word.toLowerCase();
    }

    protected final AveragedPerceptron perceptron;
    protected final Map<String, String> tags;

    protected Tagger(Map<String, String> tags, AveragedPerceptron perceptron) {
        this.tags = tags;
        this.perceptron = perceptron;
    }

    protected Tagger(Map<String, Map<String, Float>> weights,
                     Map<String, String> tags, Set<String> labels)
    {
        this(tags, new AveragedPerceptron(weights, labels));
    }

    public static Tagger loadFrom(Path inputPath) throws IOException {
        File inputFile = inputPath.toFile();

        try (FileInputStream fileStream = new FileInputStream(inputFile);
             Unpacker unpacker = new MessagePack().createUnpacker(fileStream))
        {
            Map<String, String> tags =
                unpacker.read(tMap(TString, TString));
            Map<String, Map<String, Float>> weights =
                unpacker.read(tMap(TString, tMap(TString, TFloat)));
            Set<String> labels =
                unpacker.read(new SetTemplate<>(TString));

            return new Tagger(weights, tags, labels);
        }
    }

    public void saveTo(Path outputPath) throws IOException {
        File outputFile = outputPath.toFile();

        try (FileOutputStream fileStream = new FileOutputStream(outputFile);
             Packer packer = new MessagePack().createPacker(fileStream))
        {
            packer.write(this.tags);
            packer.write(this.perceptron.weights);
            packer.write(this.perceptron.labels);
        }
    }

    protected static List<String> getContext(List<String> words) {
        return Stream.concat(Stream.concat(Arrays.stream(START),
                                           words.stream().map(Tagger::normalize)),
                             Arrays.stream(END))
                     .collect(Collectors.toList());
    }

    public List<String> tag(List<String> words) {
        List<String> context = getContext(words);
        String prev = START[0];
        String prev2 = START[1];
        int offset = START.length;
        List<String> tags = new ArrayList<>();
        for (int index = 0; index < words.size(); index++) {
            String word = words.get(index);

            String tag = this.tags.get(word);
            if (tag == null) {
                Features features =
                    Features.getFeatures(offset + index, word, context, prev, prev2);
                tag = this.perceptron.predict(features);
            }
            tags.add(tag);
            prev2 = prev;
            prev = tag;
        }
        return tags;
    }
}

@@ -133,16 +133,54 @@
 mplate;%0A
+import org.msgpack.template.Template;%0A
 import o
@@ -205,24 +205,24 @@
 r.Unpacker;%0A
-
 %0Aimport java
@@ -752,16 +752,358 @@
 9,.%5D*%22);
+%0A    protected static final Template%3CMap%3CString, String%3E%3E TAGS_TEMPLATE =%0A        tMap(TString, TString);%0A    protected static final Template%3CMap%3CString, Map%3CString, Float%3E%3E%3E WEIGHTS_TEMPLATE =%0A        tMap(TString, tMap(TString, TFloat));%0A    protected static final Template%3CSet%3CString%3E%3E LABELS_TEMPLATE =%0A        new SetTemplate%3C%3E(TString);
 %0A%0A    pr
@@ -2130,61 +2130,36 @@
 gs =
-%0A                unpacker.read(tMap(TString, TString)
+ unpacker.read(TAGS_TEMPLATE
 );%0A 
@@ -2214,167 +2214,104 @@
 ts =
-%0A                unpacker.read(tMap(TString, tMap(TString, TFloat)));%0A            Set%3CString%3E labels =%0A                unpacker.read(new SetTemplate%3C%3E(TString)
+ unpacker.read(WEIGHTS_TEMPLATE);%0A            Set%3CString%3E labels = unpacker.read(LABELS_TEMPLATE
 );%0A%0A

make templates constants