import opennlp.tools.namefind.NameFinderME;
import opennlp.tools.namefind.TokenNameFinderModel;
import opennlp.tools.util.Span;

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.InputStream;
import java.util.ArrayList;

/**
 * Created by zen on 11/15/14.
 */
public class NamedEntityFinder {

    private NameFinderME nameFinder;

    public void readModel(String modelPath) {
        InputStream modelIn = null;
        try {
            modelIn = new FileInputStream(modelPath);
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        }

        TokenNameFinderModel model = null;
        try {
            model = new TokenNameFinderModel(modelIn);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            if (modelIn != null) {
                try {
                    modelIn.close();
                }
                catch (IOException e) {
                }
            }
        }
        nameFinder = new NameFinderME(model);
    }

    public int readUniqueNamedEntities(ArrayList<String[]> tokenizedSentences) {

        ArrayList<String> names = new ArrayList<String>();

        for (String[] tokenizedSentence : tokenizedSentences) {
            Span nameSpans[] = nameFinder.find(tokenizedSentence);
            for (Span span : nameSpans) {
                String fullNameAsString = "";
                for (int namePiece = span.getStart(); namePiece < span.getEnd(); namePiece++) {
                    fullNameAsString = fullNameAsString + tokenizedSentence[namePiece] + " ";
                }
//                System.out.println(fullNameAsString);

                if (!names.contains(fullNameAsString)) {
                    names.add(fullNameAsString);
                }
            }
        }
        nameFinder.clearAdaptiveData();
        return names.size();
    }

    public int countNamedEntities(ArrayList<String[]> tokenizedSentences) {

        int counter = 0;

        for (String[] tokenizedSentence : tokenizedSentences) {
            Span nameSpans[] = nameFinder.find(tokenizedSentence);
            counter+=nameSpans.length;
            // do something with the names
        }
        nameFinder.clearAdaptiveData();
        return counter;
    }
}
