package svm.instances.dependency.kernel;

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

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Hashtable;
import java.util.Map;
import java.util.Map.Entry;

import svm.AtomicConstraint;
import svm.instances.dependency.DependencyTree;
import svm.instances.dependency.TokenSequence;
import base.AtomicKernel;
import base.Kernel;
import base.SparseVector;

//AtomicKernel
public class DependencyTreeKernel implements AtomicKernel<TokenSequence, DependencyTree>, Serializable
{
	/**
	 * 
	 */
	private static final long serialVersionUID = 1213;
	public static HashMap<String,String> word2vec;
	
	static{
		word2vec = new HashMap<String,String>();
		BufferedReader r;
		try
		{
			r = new BufferedReader(new FileReader("data/wordclusters.txt"));
			for(String l = r.readLine();l!=null;l = r.readLine())
			{
				String[] a = l.split(" ");
				word2vec.put(a[0], a[1]);
			}
		}
		catch (IOException e)
		{
			e.printStackTrace();
		}
		
	}
	
	public String hashFunction(int head, int dependent, TokenSequence x)
	{
		return attachDir(head,dependent,twoPosTags2(head,dependent,x));
	}
	
	public String hashFunction2(int head, int dependent, TokenSequence x)
	{
		return twoTokens(head, dependent, x);
	}
	
	private String attachDir(int head, int dependent, String hash)
	{
		StringBuilder stringBuilder = new StringBuilder(hash);
		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);
		}
		stringBuilder.append("::");
		stringBuilder.append(distBool);
		stringBuilder.append(attachementDirection);
		return stringBuilder.toString();
	}
	private String twoLemmas(int head, int dependent, TokenSequence x)
	{
		// TODO Auto-generated method stub
		StringBuilder stringBuilder = new StringBuilder(2 + x.lemmas[dependent].length() + x.lemmas[head].length());
		stringBuilder.append(x.lemmas[dependent]);
		stringBuilder.append("::");
		stringBuilder.append(x.lemmas[head]);
		return stringBuilder.toString();
	}
	@SuppressWarnings("unused")
	private String twoTokens(int head, int dependent, TokenSequence x)
	{
		StringBuilder stringBuilder = new StringBuilder(2 + x.tokens[dependent].length() + x.tokens[head].length());
		stringBuilder.append(x.tokens[dependent]);
		stringBuilder.append("::");
		stringBuilder.append(x.tokens[head]);
		return stringBuilder.toString();
	}
	
	@SuppressWarnings("unused")
	private String twoPosTags(int head, int dependent, TokenSequence x)
	{
		StringBuilder stringBuilder = new StringBuilder(2 + x.posTag[dependent].length() + x.posTag[head].length());
		stringBuilder.append(x.posTag[dependent]);
		stringBuilder.append("::");
		stringBuilder.append(x.posTag[head]);
		return stringBuilder.toString();
	}
	
	private String twoPosTags2(int head, int dependent, TokenSequence x)
	{
		StringBuilder stringBuilder = new StringBuilder(2 + x.posTag[dependent].length() + x.posTag[head].length());
		stringBuilder.append(x.posTag2[dependent]);
		stringBuilder.append("::");
		stringBuilder.append(x.posTag2[head]);
		return stringBuilder.toString();
//		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);
//		}
//		StringBuilder stringBuilder = new StringBuilder(2 + x.posTag2[dependent].length() + x.posTag2[head].length());
//		stringBuilder.append(x.posTag2[dependent]);
//		stringBuilder.append("::");
//		stringBuilder.append(x.posTag2[head]);
//		stringBuilder.append("::");
//		stringBuilder.append(distBool);
//		stringBuilder.append(attachementDirection);
//		return stringBuilder.toString();
	}
	@SuppressWarnings("unused")
	private String twoClusters(int head, int dependent, TokenSequence x)
	{
		String a = word2vec.get(x.tokens[head]);
		String b = word2vec.get(x.tokens[dependent]);
		
		StringBuilder stringBuilder = new StringBuilder(12);
		if(a!=null)
			stringBuilder.append(a);
		stringBuilder.append("::");
		if(b!=null)
			stringBuilder.append(b);
		return stringBuilder.toString();
	}


	@Override
	public double k(DependencyTree a, DependencyTree b)
	{	
		Hashtable<String, Integer> mergeTable = mergeTable1(a);
		Hashtable<String, Integer> mergeTable2 = mergeTable2(a);
		int sum = matchTables(b, mergeTable, mergeTable2);
		return 1.0 * sum;// / (2 * Math.sqrt(a.heads.length * b.heads.length));
	}
	private int matchTables(DependencyTree b, Hashtable<String, Integer> mergeTable, Hashtable<String, Integer> mergeTable2)
	{
		int sum = 0;
		for (int dependentB = 0; dependentB < b.x.tokens.length - 1; dependentB++)
		{
			String hashFunction  = hashFunction( b.heads[dependentB], dependentB + 1, b.x);
			if (mergeTable.containsKey(hashFunction))
				sum+=mergeTable.get(hashFunction);
//			String hashFunction2 = hashFunction2(b.heads[dependentB], dependentB + 1, b.x);
//			if (mergeTable2.containsKey(hashFunction2))
//				sum+=mergeTable2.get(hashFunction2);
		}
		return sum;
	}
	private Hashtable<String, Integer> mergeTable2(DependencyTree a)
	{
		Hashtable<String,Integer> mergeTable2 = new Hashtable<>(a.x.tokens.length+1,1);
		for (int dependentA = 0; dependentA < a.x.tokens.length - 1; dependentA++)
		{
			String hashFunction2 = hashFunction2(a.heads[dependentA], dependentA + 1, a.x);
			Integer before = mergeTable2.put(hashFunction2,1);
			if(before!=null) mergeTable2.put(hashFunction2,1+before);
				
		}
		return mergeTable2;
	}
	private Hashtable<String, Integer> mergeTable1(DependencyTree a)
	{
		Hashtable<String,Integer> mergeTable = new Hashtable<>(a.x.tokens.length+1,1);
		for (int dependentA = 0; dependentA < a.x.tokens.length - 1; dependentA++)
		{
			String hashFunction = hashFunction(a.heads[dependentA], dependentA + 1, a.x);
			Integer before = mergeTable.put(hashFunction,1);
			if(before!=null) mergeTable.put(hashFunction,1+before);
				
		}
		return mergeTable;
	}

	@Override
	public DependencyTree preImage(Map<DependencyTree, Double> alphaY, TokenSequence x)
	{
		
		
		//Compute which hash values are possible for the input x and create a map that maps hashvalues to edges with said value
		HashMap<String, ArrayList<int[]>> reverse = new HashMap<>();
		for (int i = 0; i < x.tokens.length; i++)
		{
			for (int j = 0; j < x.tokens.length; j++)
			{
				String hash = hashFunction(i, j, x);
				if (!reverse.containsKey(hash))
					reverse.put(hash, new ArrayList<int[]>(1));
				reverse.get(hash).add(new int[] { i, j });
			}
		}
		//same for the second hash function
		HashMap<String, ArrayList<int[]>> reverse2 = new HashMap<>();
		for (int i = 0; i < x.tokens.length; i++)
		{
			for (int j = 0; j < x.tokens.length; j++)
			{
				String hash = hashFunction2(i, j, x);
				if (!reverse2.containsKey(hash))
					reverse2.put(hash, new ArrayList<int[]>(1));
				reverse2.get(hash).add(new int[] { i, j });
			}
		}
		
		double[][] e = new double[x.lemmas.length][x.lemmas.length];

		
		//Iterate over each element of the given pre-image linear combination
		for (Entry<DependencyTree, Double> entry : alphaY.entrySet())
		{
			DependencyTree y = entry.getKey();
//			double scale = 2.0*Math.sqrt(y.heads.length*(x.tokens.length-1));
			//walk over each edge of the given tree and update the graph of possible edges, as that edge will contribute to the kernel dot product
			for (int i = 0; i < y.heads.length; i++)
			{
				
				String hash = hashFunction(y.heads[i], i + 1, y.x);
				//lookup which edges in the full graph have the same hashvalue
				if (reverse.containsKey(hash))
				{
					for (int[] headDependent : reverse.get(hash))
					{
						//update the edges of the graph.
						e[headDependent[0]][headDependent[1]] += entry.getValue();///scale;
					}
				}
				String hash2 = hashFunction2(y.heads[i], i + 1, y.x);
				if (reverse2.containsKey(hash2))
				{
					for (int[] headDependent : reverse2.get(hash2))
					{
						e[headDependent[0]][headDependent[1]] += entry.getValue();///scale;
					}
				}
			}
		}

		DependencyTree yPred = edmonds(x, e);
		return yPred;
	}

	@Override
	public DependencyTree preImagePenelized(Map<DependencyTree, Double> alphaY, DependencyTree yy, TokenSequence x)
	{
		HashMap<String, ArrayList<int[]>> reverse = new HashMap<>();
		for (int i = 0; i < x.tokens.length; i++)
		{
			for (int j = 0; j < x.tokens.length; j++)
			{
				String hash = hashFunction(i, j, x);
				if (!reverse.containsKey(hash))
					reverse.put(hash, new ArrayList<int[]>(1));
				reverse.get(hash).add(new int[] { i, j });
			}
		}

		HashMap<String, ArrayList<int[]>> reverse2 = new HashMap<>();
		for (int i = 0; i < x.tokens.length; i++)
		{
			for (int j = 0; j < x.tokens.length; j++)
			{
				String hash = hashFunction2(i, j, x);
				if (!reverse2.containsKey(hash))
					reverse2.put(hash, new ArrayList<int[]>(1));
				reverse2.get(hash).add(new int[] { i, j });
			}
		}
		
		double[][] e = new double[x.lemmas.length][x.lemmas.length];

		for (Entry<DependencyTree, Double> entry : alphaY.entrySet())
		{
			DependencyTree y = entry.getKey();
//			double scale = 2.0*Math.sqrt(y.heads.length*(x.tokens.length-1));
			for (int i = 0; i < y.heads.length; i++)
			{
				String hash = hashFunction(y.heads[i], i + 1, y.x);
				if (reverse.containsKey(hash))
				{
					for (int[] headDependent : reverse.get(hash))
					{
						e[headDependent[0]][headDependent[1]] += entry.getValue();///scale;
					}
				}
				String hash2 = hashFunction2(y.heads[i], i + 1, y.x);
				if (reverse2.containsKey(hash2))
				{
					for (int[] headDependent : reverse2.get(hash2))
					{
						e[headDependent[0]][headDependent[1]] += entry.getValue();///scale;
					}
				}
			}
		}
		for(int head = 0;head<x.tokens.length;head++)
			for(int dependent=0;dependent<x.tokens.length;dependent++)
				if (dependent == 0 || head != yy.heads[dependent - 1])
					e[head][dependent] += 1.0 / yy.heads.length;
		DependencyTree yPred = edmonds(x, e);
		return yPred;
	}
	@Override
	public SparseVector featureMap(TokenSequence x, DependencyTree y)
	{
		Hashtable<String, Integer> mergeTable = mergeTable1(y);
//		Hashtable<String, Integer> mergeTable2 = mergeTable2(y);
		SparseVector ret = new SparseVector();
		for(Entry<String,Integer> e: mergeTable.entrySet())
			ret.put("1::"+e.getKey(),e.getValue()*1d);///(Math.sqrt(2*y.heads.length)));
//		for(Entry<String,Integer> e: mergeTable2.entrySet())
//			ret.put("2::"+e.getKey(),e.getValue()*1d);//(Math.sqrt(2*y.heads.length)));
		return ret;
	}
	@Override
	public DependencyTree preImage(AtomicConstraint<TokenSequence,DependencyTree> w, TokenSequence x,Kernel<TokenSequence> inputKernel)
	{
		double[][] e = new double[x.lemmas.length][x.lemmas.length];	
		
		for(int head = 0;head<x.tokens.length;head++)
			for(int dependent=0;dependent<x.tokens.length;dependent++)
			{
				String hash1 = "1::"+hashFunction(head, dependent, x);
				String hash2 = "2::"+hashFunction2(head, dependent, x);
				if(w!=null)//for(int i=0;i<constraints.size();i++)
				{
					//AtomicConstraint<TokenSequence,DependencyTree> c = (AtomicConstraint<TokenSequence,DependencyTree>) constraints.get(i);
					//double alphai = alphas.get(i);
					//if(alphai<0.000001)
					//	continue;
					HashMap<TokenSequence, Double> bla = w.phi.get(hash1);
					if(bla!=null)
						for(Entry<TokenSequence,Double> entry:bla.entrySet())
						{
							e[head][dependent]+=inputKernel.k(entry.getKey(),x) * entry.getValue();// * alphai;//  / Math.sqrt(2*(x.tokens.length-1));
						}
//					bla = w.phi.get(hash2);
//					if(bla!=null)
//						for(Entry<TokenSequence,Double> entry:bla.entrySet())
//						{
//							e[head][dependent]+=inputKernel.k(entry.getKey(),x) * entry.getValue();// * alphai;//  / Math.sqrt(2*(x.tokens.length-1));
//						}
				}
			}
//		DependencyTree yPred = new DependencyTree();
//		yPred.x = x;
//		yPred.heads = new int[x.tokens.length-1];
//		for(int i=1;i<x.tokens.length;i++)
//		{
//			int max = 0;
//			for(int head=0;head<x.tokens.length;head++)
//			{
//				if(e[head][i]>e[max][i])
//					max = head;
//				yPred.heads[i-1]=max;
//			}
//		}
		DependencyTree yPred = edmonds(x, e);
		return yPred;
	}
	@Override
	public DependencyTree preImageL2Penelized(AtomicConstraint<TokenSequence,DependencyTree> w, DependencyTree yy, TokenSequence x,Kernel<TokenSequence> inputKernel)
	{	
		double[][] e = new double[x.lemmas.length][x.lemmas.length];	
		
		for(int head = 0;head<x.tokens.length;head++)
			for(int dependent=0;dependent<x.tokens.length;dependent++)
			{
				String hash1 = "1::"+hashFunction(head, dependent, x);
				String hash2 = "2::"+hashFunction2(head, dependent, x);
				if(w!=null)//for(int i=0;i<constraints.size();i++)
				{
				//	AtomicConstraint<TokenSequence,DependencyTree> c = (AtomicConstraint<TokenSequence,DependencyTree>) constraints.get(i);
					//double alphai = alphas.get(i);
					//if(alphai<0.000001)
					//	continue;
					HashMap<TokenSequence, Double> bla = w.phi.get(hash1);
					if(bla!=null)
						for(Entry<TokenSequence,Double> entry:bla.entrySet())
						{
							e[head][dependent]+=inputKernel.k(entry.getKey(),x) * entry.getValue();// * alphai;// / Math.sqrt(2*(x.tokens.length-1));
						}
//					bla = w.phi.get(hash2);
//					if(bla!=null)
//						for(Entry<TokenSequence,Double> entry:bla.entrySet())
//						{
//							e[head][dependent]+=inputKernel.k(entry.getKey(),x) * entry.getValue();// * alphai;// / Math.sqrt(2*(x.tokens.length-1));
//						}
				}
				if (dependent == 0 || head != yy.heads[dependent - 1])
					e[head][dependent] += 1.0 / yy.heads.length;
			}
//		DependencyTree yPred = new DependencyTree();
//		yPred.x = x;
//		yPred.heads = new int[x.tokens.length-1];
//		for(int i=1;i<x.tokens.length;i++)
//		{
//			int max = 0;
//			for(int head=0;head<x.tokens.length;head++)
//			{
//				if(e[head][i]>e[max][i])
//					max = head;
//				yPred.heads[i-1]=max;
//			}
//		}
		DependencyTree yPred = edmonds(x, e);
		return yPred;
	}

}
