package svm.instances.dependency;

import static svm.instances.dependency.Edmonds.edmonds;

import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import svm.SVMStruct;
import svm.SVMStructInstance;
import base.SparseVector;

public class DependencyParserInstance implements SVMStructInstance<TokenSequence, DependencyTree>
{
	private static final int BUFFERSIZE = 1000;
	public ConcurrentHashMap<TokenSequence, SparseVector[][]> precomputed = new ConcurrentHashMap<>();

	@Override
	public SparseVector phi(TokenSequence x, DependencyTree y)
	{
		SparseVector phi = new SparseVector();
		boolean known = false;

		known = precomputed.containsKey(x);

		if (known)
		{
			for (int dependent = 0; dependent < y.heads.length; dependent++)
			{
				phi.add(precomputed.get(x)[y.heads[dependent]][dependent + 1]);
			}
		}
		else
		{// Features for every arc
			SparseVector[][] precompute = new SparseVector[x.lemmas.length][x.lemmas.length];
			for (int head = 0; head < x.lemmas.length; head++)
			{
				for (int dependent = 0; dependent < x.lemmas.length; dependent++)
				{
					precompute[head][dependent] = new SparseVector();
					addEdge(x, head, dependent, precompute[head][dependent]);
				}
			}
			for (int dependent = 0; dependent < y.heads.length; dependent++)
			{
				phi.add(precompute[y.heads[dependent]][dependent + 1]);
			}
			if (precomputed.size() < BUFFERSIZE)
			{
				precomputed.put(x, precompute);
				// System.out.println(precomputed.size() + " "+ Runtime.getRuntime().freeMemory()/1024/1024);
			}

		}
		return phi;
	}

	private static String concat(String... a)
	{
		StringBuilder b = new StringBuilder();
		for (String x : a)
		{
			b.append(x);
			b.append("::");
		}
		return b.toString();
	}

	/**
	 * @param x
	 * @param head
	 * @param dependent
	 * @param phi
	 */
	private void addEdge(TokenSequence x, int head, int dependent, SparseVector phi)
	{
		String attachementDirection = head < dependent ? "RIGHT" : "LEFT";
		int dist = Math.abs(head - dependent);
		String distBool = "0";
		if (dist > 10)
		{
			distBool = "10";
		}
		else if (dist > 5)
		{
			distBool = "5";
		}
		else
		{
			distBool = Integer.toString(dist - 1);
		}

		for (String suffix : new String[] {"","::DIR-" + attachementDirection + "-" + distBool })
		{

			// Edge Features
			phi.plusOne(concat("THD", x.tokens[head], x.tokens[dependent], suffix));
			phi.plusOne(concat("LHD", x.lemmas[head], x.lemmas[dependent], suffix));
			phi.plusOne(concat("PHD", x.posTag[head], x.posTag[dependent], suffix));
			phi.plusOne(concat("P2HD", x.posTag2[head], x.posTag2[dependent], suffix));
			phi.plusOne(concat("TPHD", x.tokens[head], x.tokens[dependent], x.posTag[head], x.posTag[dependent], suffix));
			phi.plusOne(concat("TP2HD", x.tokens[head], x.tokens[dependent], x.posTag2[head], x.posTag2[dependent], suffix));
			phi.plusOne(concat("LPHD", x.lemmas[head], x.lemmas[dependent], x.posTag[head], x.posTag[dependent], suffix));
			phi.plusOne(concat("LP2HD", x.lemmas[head], x.lemmas[dependent], x.posTag2[head], x.posTag2[dependent], suffix));

//			// Head Features
			phi.plusOne(concat("TH", x.tokens[head], suffix));
			phi.plusOne(concat("LH", x.lemmas[head], suffix));
			phi.plusOne(concat("PH", x.posTag[head], suffix));
			phi.plusOne(concat("P2H", x.posTag2[head], suffix));
//
			phi.plusOne(concat("TPH", x.tokens[head], x.posTag[head], suffix));
			phi.plusOne(concat("TP2H", x.tokens[head], x.posTag2[head], suffix));
//
//			// Dependent Features
			phi.plusOne(concat("TD", x.tokens[dependent], suffix));
			phi.plusOne(concat("LD", x.lemmas[dependent], suffix));
			phi.plusOne(concat("PD", x.posTag[dependent], suffix));
			phi.plusOne(concat("P2D", x.posTag2[dependent], suffix));

			phi.plusOne(concat("TPD", x.tokens[dependent], x.posTag[dependent], suffix));
			phi.plusOne(concat("TP2D", x.tokens[dependent], x.posTag2[dependent], suffix));

			// in between features
			for (int i = Math.min(head, dependent) + 1; i < Math.max(head, dependent); i++)
			{
				phi.plusOne(concat("THBD", x.posTag[head], x.posTag[i], x.posTag[dependent], suffix));
				phi.plusOne(concat("T2HBD", x.posTag2[head], x.posTag2[i], x.posTag2[dependent], suffix));
			}

			// surrounding words features
			
			// LHDR
			phi.plusOne(concat("TN1", 
					(head > 0 ? x.posTag[head - 1] : "<null>"), 
					x.posTag[head], 
					x.posTag[dependent],
					(dependent + 1 < x.posTag.length ? x.posTag[dependent + 1] : "<null>"), 
					suffix));
			phi.plusOne(concat("T2N1",
					(head > 0 ? x.posTag2[head - 1] : "<null>"), 
					x.posTag2[head], 
					x.posTag2[dependent], 
					(dependent + 1 < x.posTag2.length ? x.posTag2[dependent + 1] : "<null>")
					+ suffix));
			// HRLD
			phi.plusOne(concat("TN2", 
					x.posTag[head],
					(head + 1 < x.posTag.length ? x.posTag[head + 1] : "<null>"),
					(dependent > 0 ? x.posTag[dependent - 1] : "<null>"),
					x.posTag[dependent],
					suffix));
			phi.plusOne(concat("T2N2", 
					x.posTag2[head],
					(head + 1 < x.posTag2.length ? x.posTag2[head + 1] : "<null>"),
					(dependent > 0 ? x.posTag2[dependent - 1] : "<null>"),
					x.posTag2[dependent] ,
					suffix));
			// HRDR
			phi.plusOne(concat("TN3", 
					x.posTag[head],
					(head + 1 < x.posTag.length ? x.posTag[head + 1] : "<null>"),
					x.posTag[dependent],
					(dependent + 1 < x.posTag.length ? x.posTag[dependent + 1] : "<null>"),
					suffix));
			phi.plusOne(concat("T2N3", 
					x.posTag2[head],
					(head + 1 < x.posTag2.length ? x.posTag2[head + 1] : "<null>"),
					x.posTag2[dependent],
					(dependent + 1 < x.posTag2.length ? x.posTag2[dependent + 1] : "<null>"),
					suffix));
			// LHLD
			phi.plusOne(concat("TN4", 
					(head > 0 ? x.posTag[head - 1] : "<null>"),
					x.posTag[head],
					(dependent > 0 ? x.posTag[dependent - 1] : "<null>"),
					x.posTag[dependent],
					suffix));
			phi.plusOne(concat("T2N4", 
					(head > 0 ? x.posTag2[head - 1] : "<null>"),
					x.posTag2[head],
					(dependent > 0 ? x.posTag2[dependent - 1] : "<null>"),
					x.posTag2[dependent],
					suffix));
		}
	}

	@Override
	public DependencyTree predict(TokenSequence x, SVMStruct svm)
	{
		boolean known = false;
		known = precomputed.containsKey(x);

		double[][] e = new double[x.lemmas.length][x.lemmas.length];
		for (int head = 0; head < x.lemmas.length; head++)
		{
			for (int dependent = 0; dependent < x.lemmas.length; dependent++)
			{
				SparseVector phiij;
				if (known)
					phiij = precomputed.get(x)[head][dependent];
				else
				{
					phiij = new SparseVector();
					addEdge(x, head, dependent, phiij);
				}

				e[head][dependent] = phiij.multiply(svm.w, svm.wMap);
			}
		}
		DependencyTree yPred = edmonds(x, e);
		return yPred;
	}

	@Override
	public DependencyTree marginRescaling(TokenSequence x, DependencyTree y, SVMStruct svm)
	{
		boolean known = false;

		known = precomputed.containsKey(x);
		// TODO: Retrieve Edges
		double[][] e = new double[x.lemmas.length][x.lemmas.length];
		for (int head = 0; head < x.lemmas.length; head++)
		{
			for (int dependent = 0; dependent < x.lemmas.length; dependent++)
			{
				SparseVector phiij;
				if (known)
					phiij = precomputed.get(x)[head][dependent];
				else
				{
					phiij = new SparseVector();
					addEdge(x, head, dependent, phiij);
				}

				// addEdge(x, head, dependent, phiij);
				e[head][dependent] = phiij.multiply(svm.w, svm.wMap);
				if (dependent == 0 || head != y.heads[dependent - 1])
					e[head][dependent] += 1.0 / y.heads.length;
			}
		}
		DependencyTree yPred = edmonds(x, e);
		return yPred;
	}

	

	@Override
	public DependencyTree slackRescaling(TokenSequence x, DependencyTree y, SVMStruct svm)
	{
		// TODO Auto-generated method stub
		throw new RuntimeException("not yet implemented");
	}

	@Override
	public double loss(DependencyTree y, DependencyTree predY)
	{
		assert predY != null;

		if (y.heads.length != predY.heads.length)
		{
			System.out.println("dafuq");
			return 1;
		}
		int errors = 0;
		for (int i = 0; i < y.heads.length; i++)
			if (y.heads[i] != predY.heads[i])
				errors++;
		return 1.0 * errors / y.heads.length;
	}

	@Override
	public Set<String> features(TokenSequence x, DependencyTree y)
	{
		// TODO Auto-generated method stub
		SparseVector s = new SparseVector();
		for (int dependent = 0; dependent < y.heads.length; dependent++)
		{
			addEdge(x, y.heads[dependent], dependent + 1, s);
		}
		return s.keySet();
	}
}
