

package kottmann.opennlp.addons.mallet;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Map;

import opennlp.tools.ml.AbstractEventTrainer;
import opennlp.tools.ml.model.DataIndexer;
import opennlp.tools.ml.model.MaxentModel;
import cc.mallet.classify.Classifier;
import cc.mallet.classify.MaxEntTrainer;
import cc.mallet.types.Alphabet;
import cc.mallet.types.FeatureVector;
import cc.mallet.types.Instance;
import cc.mallet.types.InstanceList;
import cc.mallet.types.LabelAlphabet;

public class MaxentTrainer extends AbstractEventTrainer {

  public MaxentTrainer(Map<String, String> trainParams,
      Map<String, String> reportMap) {
    super(trainParams, reportMap);
  }

  @Override
  public boolean isSortAndMerge() {
    return true;
  }

  @Override
  public MaxentModel doTrain(DataIndexer indexer) throws IOException {

    int numFeatures = indexer.getPredLabels().length;

    Alphabet dataAlphabet = new Alphabet(numFeatures);
    LabelAlphabet targetAlphabet = new LabelAlphabet();

    Collection<Instance> instances = new ArrayList<>();

    String predLabels[] = indexer.getPredLabels();
    
    int outcomes[] = indexer.getOutcomeList();
    for (int contextIndex = 0; contextIndex < indexer.getContexts().length; contextIndex++) {

      int malletFeatures[] = new int[indexer.getContexts()[contextIndex].length];
      double weights[] = new double[indexer.getContexts()[contextIndex].length];

      for (int featureIndex = 0; featureIndex < malletFeatures.length; featureIndex++) {
        malletFeatures[featureIndex] = dataAlphabet.lookupIndex(
            predLabels[indexer.getContexts()[contextIndex][featureIndex]], true);
        
        weights[featureIndex] = indexer.getNumTimesEventsSeen()[contextIndex];
      }

      FeatureVector fv = new FeatureVector(dataAlphabet, malletFeatures,
          weights);
      Instance inst = new Instance(fv, targetAlphabet.lookupLabel(
          indexer.getOutcomeLabels()[outcomes[contextIndex]], true), "name",
          "data-indexer");
      instances.add(inst);
    }

    InstanceList trainingData = new InstanceList(dataAlphabet, targetAlphabet);
    Instance inst = instances.iterator().next();

    Alphabet.alphabetsMatch(trainingData, inst);
    trainingData.addAll(instances);

    MaxEntTrainer trainer = new MaxEntTrainer();
    
    Classifier classifier = trainer.train(trainingData);

    return new ClassifierModel(classifier);
  }
}
