package de.tu.darmstadt.lt.ner.writer;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;

import org.apache.uima.analysis_engine.AnalysisEngineProcessException;
import org.apache.uima.jcas.JCas;
import org.apache.uima.util.Level;
import org.uimafit.component.JCasConsumer_ImplBase;
import org.uimafit.descriptor.ConfigurationParameter;
import org.uimafit.util.JCasUtil;

import de.tu.darmstadt.lt.ner.types.GoldNamedEntity;
import de.tudarmstadt.ukp.dkpro.core.api.ner.type.NamedEntity;
import de.tudarmstadt.ukp.dkpro.core.api.segmentation.type.Sentence;

public class EvaluatedNERWriter
    extends JCasConsumer_ImplBase
{
    public static final String OUTPUT_FILE = "OutputFile";
    @ConfigurationParameter(name = OUTPUT_FILE, mandatory = true)
    private File OutputFile = null;

    public static final String NOD_OUTPUT_FILE = "nodOutputFile";
    @ConfigurationParameter(name = NOD_OUTPUT_FILE, mandatory = false)
    private File nodOutputFile = null;

    public static final String IS_GOLD = "isGold";
    @ConfigurationParameter(name = IS_GOLD, mandatory = false)
    private boolean isGold = false;

    public static final String SENTENCES_ID = "sentencesId";
    @ConfigurationParameter(name = SENTENCES_ID, mandatory = false)
    private List<String> sentencesId = null;

    public static final String LF = System.getProperty("line.separator");
    public static final String TAB = "\t";
    private static final String ORG = "ORG";
    private static final String PER = "PER";
    private static final String B_ORG = "B-ORG";
    private static final String B_PER = "B-PER";
    private static final String I_ORG = "I-ORG";
    private static final String I_PER = "I-PER";
    private static final String TYPE_SEP = "$";
    private static final String ENT_SEP = ",";

    @Override
    public void process(JCas jCas)
        throws AnalysisEngineProcessException
    {
        try {
            FileWriter outputWriter = new FileWriter(OutputFile);
            Map<Sentence, Collection<NamedEntity>> sentencesNER = JCasUtil.indexCovered(jCas,
                    Sentence.class, NamedEntity.class);
            
            FileWriter nodOutputWriter = null;
            if (nodOutputFile != null) {
                nodOutputWriter = new FileWriter(nodOutputFile);
            }
            int sentenceIndex = 0;
            
            List<Sentence> sentences = new ArrayList<Sentence>(sentencesNER.keySet());
            // sort sentences by sentence 
            Collections.sort(sentences, new Comparator<Sentence>() {                                
                public int compare(Sentence arg0, Sentence arg1) {
                    return arg0.getBegin() - arg1.getBegin();
                }
            });
            for (Sentence sentence : sentences) {

                List<String> personSb = new ArrayList<String>();
                List<String> orgSb = new ArrayList<String>();
                String prevNeType = "O";
                String namedEntity = "";

                for (NamedEntity neAnnotation : sentencesNER.get(sentence)) {

                    String text = neAnnotation.getCoveredText().replace(" ", "");
                    String neType = neAnnotation.getValue();

                    StringBuilder sb = new StringBuilder();
                    sb.append(text);
                    sb.append(" ");
                    if (isGold) {
                        sb.append(JCasUtil.selectCovered(jCas, GoldNamedEntity.class, neAnnotation)
                                .get(0).getNamedEntityType());
                    }
                    sb.append(" ");
                    sb.append(neAnnotation.getValue());
                    sb.append(LF);
                    outputWriter.write(sb.toString());

                    // save the sentence-ne file for NoD
                    if (nodOutputFile != null) {
                        if (neType.equals("O")) {
                            if (prevNeType == "O") {
                                continue;
                            }
                            else if (prevNeType.equals(PER)) {
                                personSb.add(namedEntity);
                                prevNeType = "O";
                                namedEntity = "";
                            }
                            else {
                                orgSb.add(namedEntity);
                                prevNeType = "O";
                                namedEntity = "";
                            }

                        }
                        // if ORG is followed by PER
                        else if (prevNeType.equals(ORG)
                                && (neType.equals(B_PER) || neType.equals(I_PER))) {
                            orgSb.add(namedEntity);
                            prevNeType = PER;
                            namedEntity = text;
                        }
                        // if PER is followed by ORG
                        else if (prevNeType.equals(PER)
                                && (neType.equals(B_ORG) || neType.equals(I_ORG))) {
                            personSb.add(namedEntity);
                            prevNeType = ORG;
                            namedEntity = text;
                        }
                        else if (neType.equals(B_PER) || neType.equals(I_PER)) {
                            prevNeType = PER;
                            namedEntity = namedEntity.trim() + " " + text;
                        }
                        else if (neType.equals(B_ORG) || neType.equals(I_ORG)) {
                            prevNeType = ORG;
                            namedEntity = namedEntity.trim() + " " + text;
                        }
                        else {
                            prevNeType = "O";
                        }

                    }
                }
                outputWriter.write(LF);
                if (nodOutputFile != null && sentencesId != null) {
                    nodOutputWriter.write(sentencesId.get(sentenceIndex) + TAB
                            + (personSb.size() == 0 ? "O" : listNames(personSb)) + TAB + TYPE_SEP + TAB
                            + (orgSb.size() == 0 ? "O" : listNames(orgSb))+LF);
                }
                sentenceIndex++;
            }
            outputWriter.close();
            nodOutputWriter.close();

            getContext().getLogger().log(Level.INFO,
                    "Output written to: " + OutputFile.getAbsolutePath());
        }
        catch (IOException e) {
            e.printStackTrace();
        }
    }

    private String listNames(List<String> names)
    {
        if (names.size() == 1) {
            return names.get(0);
        }
        String name = "";
        for (String nameInList : names) {
            if (name.equals("")) {
                name = nameInList +ENT_SEP;
            }
            else {
                name = name + nameInList;
            }
        }
        return name;
    }
}
@@ -4666,32 +4666,39 @@
 .add(namedEntity
+.trim()
 );%0D%0A            
@@ -5085,32 +5085,39 @@
 .add(namedEntity
+.trim()
 );%0D%0A            

remove trailing white space