import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.StringReader;
import java.util.Collections;

import opennlp.tools.namefind.NameFinderME;
import opennlp.tools.namefind.NameSample;
import opennlp.tools.namefind.NameSampleDataStream;
import opennlp.tools.namefind.TokenNameFinderModel;
import opennlp.tools.util.ObjectStream;
import opennlp.tools.util.PlainTextByLineStream;
import opennlp.tools.util.featuregen.AggregatedFeatureGenerator;
import opennlp.tools.util.featuregen.PreviousMapFeatureGenerator;
import opennlp.tools.util.featuregen.TokenClassFeatureGenerator;
import opennlp.tools.util.featuregen.TokenFeatureGenerator;
import opennlp.tools.util.featuregen.WindowFeatureGenerator;

/**
 * 中文命名實體識別模型訓練元件
 * 
 * @author ddlovehy
 *
 */
public class NamedEntityMultiFindTrainer {

	// 預設引數
	private int iterations = 80;
	private int cutoff = 5;
	private String langCode = "general";
	private String type = "default";

	// 待設定的引數
	private String nameWordsPath; // 命名實體詞庫路徑
	private String dataPath; // 訓練集已分詞語料路徑
	private String modelPath; // 模型儲存路徑

	public NamedEntityMultiFindTrainer() {
		super();
		// TODO Auto-generated constructor stub
	}

	public NamedEntityMultiFindTrainer(String nameWordsPath, String dataPath,
			String modelPath) {
		super();
		this.nameWordsPath = nameWordsPath;
		this.dataPath = dataPath;
		this.modelPath = modelPath;
	}

	public NamedEntityMultiFindTrainer(int iterations, int cutoff,
			String langCode, String type, String nameWordsPath,
			String dataPath, String modelPath) {
		super();
		this.iterations = iterations;
		this.cutoff = cutoff;
		this.langCode = langCode;
		this.type = type;
		this.nameWordsPath = nameWordsPath;
		this.dataPath = dataPath;
		this.modelPath = modelPath;
	}

	/**
	 * 生成定製特徵
	 * 
	 * @return
	 */
	public AggregatedFeatureGenerator prodFeatureGenerators() {
		AggregatedFeatureGenerator featureGenerators = new AggregatedFeatureGenerator(
				new WindowFeatureGenerator(new TokenFeatureGenerator(), 2, 2),
				new WindowFeatureGenerator(new TokenClassFeatureGenerator(), 2,
						2), new PreviousMapFeatureGenerator());

		return featureGenerators;
	}

	/**
	 * 將模型寫入磁碟
	 * 
	 * @param model
	 * @throws Exception
	 */
	public void writeModelIntoDisk(TokenNameFinderModel model) throws Exception {
		File outModelFile = new File(this.getModelPath());
		FileOutputStream outModelStream = new FileOutputStream(outModelFile);
		model.serialize(outModelStream);
	}

	/**
	 * 讀出標註的訓練語料
	 * 
	 * @return
	 * @throws Exception
	 */
	public String getTrainCorpusDataStr() throws Exception {

		// TODO 考慮入持久化判斷直接載入標註資料的情況 以及增量式訓練

		String trainDataStr = null;
		trainDataStr = NameEntityTextFactory.prodNameFindTrainText(
				this.getNameWordsPath(), this.getDataPath(), null);

		return trainDataStr;
	}

	/**
	 * 訓練模型
	 * 
	 * @param trainDataStr
	 *            已標註的訓練資料整體字串
	 * @return
	 * @throws Exception
	 */
	public TokenNameFinderModel trainNameEntitySamples(String trainDataStr)
			throws Exception {
		ObjectStream<NameSample> nameEntitySample = new NameSampleDataStream(
				new PlainTextByLineStream(new StringReader(trainDataStr)));
		
		System.out.println("**************************************");
		System.out.println(trainDataStr);

		TokenNameFinderModel nameFinderModel = NameFinderME.train(
				this.getLangCode(), this.getType(), nameEntitySample,
				this.prodFeatureGenerators(),
				Collections.<String, Object> emptyMap(), this.getIterations(),
				this.getCutoff());

		return nameFinderModel;
	}

	/**
	 * 訓練元件總呼叫方法
	 * 
	 * @return
	 */
	public boolean execNameFindTrainer() {

		try {
			String trainDataStr = this.getTrainCorpusDataStr();
			TokenNameFinderModel nameFinderModel = this
					.trainNameEntitySamples(trainDataStr);
			// System.out.println(nameFinderModel);
			this.writeModelIntoDisk(nameFinderModel);

			return true;
		} catch (Exception e) {
			// TODO Auto-generated catch block
			e.printStackTrace();

			return false;
		}
	}
｝