package lkformat;

import java.io.*;
import java.util.*;
import java.util.regex.*;
import se.lth.cs.nlp.depsrl.format.*;
import se.lth.cs.nlp.nlputils.depgraph.*;
import se.lth.cs.nlp.nlputils.core.*;
import gnu.trove.*;

public class CoNLL2008ToLK2 {
    
    private static final String ENCODING = "UTF-8";
    
    public static void main(String[] argv) {
        String lkTextFile = argv[0];
        String lkTokenFile = argv[1];
        String conll2008File = argv[2];

        String pbDir = argv[3];
        String nbDir = argv[4];
        
        String basename = lkTokenFile.replaceAll("\\.tokens\\.xml", "");
        String depOutFile = basename + ".depsyntax.xml";
        String paOutFile = basename + ".predargs.xml";

        try {
            BufferedReader tInput = new BufferedReader(new FileReader(lkTokenFile));
            Scanner conllInput = new Scanner(new FileReader(conll2008File));
            
            PrintWriter depOut = new PrintWriter(new OutputStreamWriter(new FileOutputStream(depOutFile), ENCODING));
            PrintWriter paOut = new PrintWriter(new OutputStreamWriter(new FileOutputStream(paOutFile), ENCODING));

            LexicalDB propBank = new LexicalDB(pbDir);
			LexicalDB nomBank = new LexicalDB(nbDir);
            
            TIntArrayList ids = new TIntArrayList();

            Pattern pattern = Pattern.compile("id=\"(.*?)\"");
            
            String line = tInput.readLine();
            boolean started = false;
            while(line != null) {
                line = line.trim();
                if(line.equals("")) {
                    // skip
                } if(!started && line.startsWith("<annotation") && line.contains("TOKENS"))
                    started = true;
                else if(started && line.startsWith("</annotation"))
                    break;
                else if(started) {
                    if(!line.startsWith("<e "))
                        throw new RuntimeException("unexpected: "  + line);
                    Matcher m = pattern.matcher(line);
                    if(!m.find())
                        throw new RuntimeException("unexpected: " + line);
                    int id = Integer.parseInt(m.group(1));
                    ids.add(id);
                }
                line = tInput.readLine();
            }
            
            depOut.println("<?xml version=\"1.0\" encoding=\"" + ENCODING + "\"?>");
            depOut.println("<lk-annotation>");
            depOut.println("<meta-info>");
            depOut.println("  <tag name=\"base\">" + lkTextFile + "</tag>");
            depOut.println("  <tag name=\"annotator\">LTH-DEP-SRL</tag>");
            depOut.println("</meta-info>");
            depOut.println("<annotation scope=\"" + lkTokenFile + "\" provides=\"DEP-SYNTAX\">");

            int depIndex = 0;

            ArrayList<String[]> predOut = new ArrayList();
            ArrayList<String[]> argOut = new ArrayList();
            ArrayList<String[]> linkOut = new ArrayList();

            Triple<DepGraph, DepGraph, List<PAStructure>> t = CoNLL2008Format.readNextGraph(conllInput);
            while(t != null) {
                DepGraph dg = t.first;
                printDepGraph(depOut, dg, depIndex, ids);

                SRLPostProcess.processPAs(predOut, argOut, linkOut, t.third, dg, propBank, nomBank);

                depIndex += dg.nodes.length - 1;
                
                depOut.println();
                t = CoNLL2008Format.readNextGraph(conllInput);
            }

            depOut.println("</annotation>");
            depOut.println("</lk-annotation>");
            depOut.close();

            paOut.println("<?xml version=\"1.0\" encoding=\"" + ENCODING + "\"?>");
            paOut.println("<lk-annotation>");
            paOut.println("<meta-info>");
            paOut.println("  <tag name=\"base\">" + lkTextFile + "</tag>");
            paOut.println("  <tag name=\"annotator\">LTH-DEP-SRL</tag>");
            paOut.println("</meta-info>");
            paOut.println("<annotation scope=\"" + lkTokenFile + "\" provides=\"SRL-preds\">");            
            for(String[] ps: predOut) {
            	paOut.println(" <e id=\"" + ps[0] + "\" on=\"#" 
                        + ps[1] + "\">" + ps[2] + "</e>");
            }            
            paOut.println("</annotation>");
            
            /*
            paOut.println("<annotation provides=\"SRL-arg-heads\">");            
            for(String[] ps: argOut) {
            	paOut.println(" <e id=\"" + ps[0] + "\" on=\"#" 
                        + ps[1] + "\"/>");
            }            
            paOut.println("</annotation>");
            */
            paOut.println("<annotation provides=\"SRL-arg-links\">");
            for(String[] ps: linkOut) {
            	String linkId = ps[0];
            	String predId = ps[1];
            	
            	String argEvId = ps[2];
            	String argTokenId = ps[3];
            	
            	String pbLabel = ps[4];
            	String vnLabel = ps[5];
            	String pbDescr = ps[6];
            	String connective = ps[7];
            	
            	String argId;
            	if(argEvId != null)
            		argId = "#" + argEvId;
            	else
            		argId = lkTokenFile + "#" + argTokenId;
            	
            	paOut.print(" <e id=\"" + linkId + "\" from=\"#" + predId
            			+ "\" to=\"" + argId + "\">");
            	
            	paOut.print("<a rid=\"" + pbLabel + "\"");
            	if(pbDescr != null && !pbDescr.isEmpty())
                	paOut.print(" rd=\"" + pbDescr + "\"");
            	if(vnLabel != null && !vnLabel.isEmpty())
                	paOut.print(" vn=\"" + vnLabel + "\"");
            	if(connective != null && !connective.isEmpty())
                	paOut.print(" con=\"" + connective + "\"");            	
            	paOut.println("/></e>");
            }            
            paOut.println("</annotation>");
            
            paOut.println("</lk-annotation>");
            paOut.close();
            
        } catch(Exception e) {
            e.printStackTrace();
            System.exit(1);
        }
        
    }

    private static void printDepGraph(PrintWriter out, DepGraph dg, 
            int depIndexStart, TIntArrayList ids) {

        for(int i = 1; i < dg.nodes.length; i++) {
            DepNode n = dg.nodes[i];
            if(n.parents.length != 1)
                throw new IllegalArgumentException("Only single-head dependency trees allowed yet");
            DepNode p = dg.nodes[i].parents[0];
            int childIndex = depIndexStart + n.position - 1;
            int childTokenId = ids.get(childIndex);
            
            if(p.position == 0) {
                out.println(" <e id=\"" + childTokenId + "\" on=\"#" 
                        + childTokenId + "\">" + n.relations[0] + "</e>");
            } else {
                int parentIndex = depIndexStart + p.position - 1;
                int parentTokenId = ids.get(parentIndex);
                out.println(" <e id=\"" + childTokenId + "\" from=\"#" 
                        + parentTokenId + "\" to=\"#" + childTokenId
                        + "\">" + n.relations[0] + "</e>");
            }
        }        
    }
    

    
}