package se.lth.cs.nlp.opinions;

import java.util.*;

import se.lth.cs.nlp.nlputils.ml_long.*;

class Viterbi {

	static Labeling search1(Sentence x, Labeling y,
			boolean maxLoss, boolean sqrt, FeatureTemplateSet fe, int[] range, 
			Model model, double[] scoresOut) {

		DPItem[] items = new DPItem[1];
		DPItem dummy = new DPItem();
		dummy.totalScore = 0;
		dummy.prev = dummy;
		dummy.label = FeatureTemplate.OUTSIDE_LABEL;
		items[0] = dummy;
		int[] ofs = new int[2];
		
		//double normalizer = 1.0 / Math.sqrt(x.features.length);
		double normalizer = sqrt? 1.0 / Math.sqrt(x.features.length): 1.0;
		
		SparseVector sv = new SparseVector();
		for(int i = 0; i < x.features.length; i++) {
			DPItem[] prevItems = items;
			items = new DPItem[range.length];			
			for(int j = 0; j < range.length; j++) {
				int label = range[j];				
				DPItem item = makeDPItem(x, y,
						fe, model, sv, 
						i, label, prevItems, ofs);

				// TODO detta ser konstigt ut!
				//if(maxLoss)
				//	item.totalScore += 1.0 / Math.sqrt(x.features.length);

				// Bug-fixed version.
				if(maxLoss && label != y.labels[i])
					if(sqrt)
						item.totalScore += 1.0 * normalizer;
					else
						item.totalScore += 1.0;
				
				items[j] = item;
			}			
		}
		
		DPItem finalItem = makeDPItem(x, y,
				fe, model, sv, 
				x.features.length, FeatureTemplate.OUTSIDE_LABEL, items, ofs);
		
		int[] labels = new int[x.features.length];
		DPItem item = finalItem;
		for(int i = x.features.length - 1; i >= 0; i--) {
			item = item.prev;
			labels[i] = item.label;
		}
		scoresOut[0] = finalItem.totalScore;
		return new Labeling(labels);
	}
	
	private static DPItem makeDPItem(Sentence x, Labeling y,
			FeatureTemplateSet fe, Model model, SparseVector sv, 
			int position, int label, DPItem[] prevItems,
			int[] ofs) {

		sv.clear();

		//fe.extractStaticFeatures(x, position, sv, label);
		ofs[0] = label;
		fe.extractFeatures(sv, x.features, ofs, position, 0); // TODO lift this out
		fe.extractFeatures(sv, x.features, ofs, position, 1);

		double staticScore = model.score(sv);
		
		double maxDynScore = Double.NEGATIVE_INFINITY;
		DPItem bestPrev = null;

		for(int k = 0; k < prevItems.length; k++) {
			sv.clear();
			DPItem prev = prevItems[k];
			ofs[1] = prev.label;
			//fe.extractDynamicFeatures(x, position, sv, label, prev.label);
			fe.extractFeatures(sv, x.features, ofs, position, 2);
			double dynScore = prev.totalScore + model.score(sv);
			if(dynScore > maxDynScore) {
				maxDynScore = dynScore;
				bestPrev = prev;
			}
		}

		if(bestPrev == null) {
			System.out.println("x = " + x);
			System.out.println("position = " + position);
			System.out.println("prevItems.length = " + prevItems.length);
			
			throw new RuntimeException("bestPrev = null");
		}
		
		DPItem item = new DPItem();
		item.prev = bestPrev;
		item.totalScore = staticScore + maxDynScore;
		item.label = label;
		
		return item;
	}
	
	static Labeling search2(Sentence x, Labeling y,
			boolean maxLoss, boolean sqrt, FeatureTemplateSet fe, int[] range, 
			Model model, double[] scoresOut) {

		final int senLength = x.features.length;
		DPItem[] items = new DPItem[1];
		DPItem dummy = new DPItem();
		dummy.totalScore = 0;
		dummy.prev = dummy;
		dummy.label = FeatureTemplate.OUTSIDE_LABEL;
		items[0] = dummy;
		int[] ofs = new int[3];
		
		//double normalizer = 1.0 / x.features.length; // why differs?
		double normalizer = sqrt? 1.0 / Math.sqrt(x.features.length): 1.0;
		
		int nprev1 = 1, nprev2 = 1;
		
		SparseVector sv = new SparseVector();
		for(int i = 0; i < senLength; i++) {

			DPItem[] prevItems = items;
			items = new DPItem[nprev1 * range.length];
			
			for(int j = 0; j < range.length; j++) {

				int label = range[j];

				sv.clear();
				
				//fe.extractStaticFeatures(x, i, sv, label);
				ofs[0] = label;
				fe.extractFeatures(sv, x.features, ofs, i, 0); // TODO lift this out
				fe.extractFeatures(sv, x.features, ofs, i, 1);

				double staticScore = model.score(sv);

				// Bug-fixed version.
				if(maxLoss && label != y.labels[i])
					if(sqrt)
						staticScore += 1.0 * normalizer;
					else
						staticScore += 1.0;
					
				for(int k = 0; k < nprev1; k++) {

					sv.clear();
					//fe.extractDynamicFeatures(x, i, sv, label, prevItems[k*nprev2].label);
					ofs[1] = prevItems[k*nprev2].label;
					fe.extractFeatures(sv, x.features, ofs, i, 2);
					double dynScore1 = model.score(sv);
					
					double maxDynScore2 = Double.NEGATIVE_INFINITY;
					DPItem maxItem = null;
					
					for(int l = 0; l < nprev2; l++) {
						DPItem item = prevItems[k*nprev2 + l];
						if(item.label != ofs[1])
							throw new RuntimeException("!!!");
						sv.clear();
						//fe.extractDynamicFeatures2(x, i, sv, label, item.label, item.prev.label);
						ofs[2] = item.prev.label;
						fe.extractFeatures(sv, x.features, ofs, i, 3);
						double dynScore2 = model.score(sv) + item.totalScore;
						if(dynScore2 > maxDynScore2) {
							maxDynScore2 = dynScore2;
							maxItem = item;
						}
					}
					
					DPItem newItem = new DPItem();
					newItem.label = label;
					newItem.prev = maxItem;
					newItem.totalScore = staticScore + dynScore1 + maxDynScore2;

					//if(maxLoss) // TODO bug
					//	newItem.totalScore += 1.0; // / Math.sqrt(x.ranges.length);
					
					items[j*nprev1 + k] = newItem;
				}
				
			}			
			nprev2 = nprev1;
			nprev1 = range.length;
		}
		
		DPItem finalItem;
		{
			int i = senLength;
			DPItem[] prevItems = items;
			
			int label = FeatureTemplate.OUTSIDE_LABEL;

			sv.clear();
			//fe.extractStaticFeatures(x, i, sv, label);
			ofs[0] = label;
			fe.extractFeatures(sv, x.features, ofs, i, 0); // TODO lift this out
			fe.extractFeatures(sv, x.features, ofs, i, 1);
			double staticScore = model.score(sv);

			double maxDynScore = Double.NEGATIVE_INFINITY;
			DPItem maxItem = null;

			for(int k = 0; k < nprev1; k++) {

				sv.clear();
				//fe.extractDynamicFeatures(x, i, sv, label, prevItems[k*nprev2].label);
				ofs[1] = prevItems[k*nprev2].label;
				fe.extractFeatures(sv, x.features, ofs, i, 2);
				double dynScore1 = model.score(sv);

				for(int l = 0; l < nprev2; l++) {
					DPItem item = prevItems[k*nprev2 + l];
					if(item.label != ofs[1])
						throw new RuntimeException("!!!");
					sv.clear();
					//fe.extractDynamicFeatures2(x, i, sv, label, item.label, item.prev.label);
					ofs[2] = item.prev.label;
					fe.extractFeatures(sv, x.features, ofs, i, 3);
					
					double dynScore = model.score(sv) + item.totalScore + dynScore1;
					if(dynScore > maxDynScore) {
						maxDynScore = dynScore;
						maxItem = item;
					}
				}					
			}
			
			finalItem = new DPItem();
			finalItem.label = label;
			finalItem.prev = maxItem;
			finalItem.totalScore = staticScore + maxDynScore;
		}
		
		int[] labels = new int[senLength];
		DPItem item = finalItem;
		for(int i = senLength - 1; i >= 0; i--) {
			item = item.prev;
			labels[i] = item.label;
		}
		scoresOut[0] = finalItem.totalScore;
		return new Labeling(labels);
	}
	
	static int search1KBest(Sentence x, Labeling y, int k,
			Labeling[] results, double[] scores,
			boolean maxLoss, boolean sqrt, FeatureTemplateSet fe, int[] range, 
			Model model) {

		ArrayList<DPItem>[] items = new ArrayList[1];
		items[0] = new ArrayList();
		
		DPItem dummy = new DPItem();
		dummy.totalScore = 0;
		dummy.prev = dummy;
		dummy.label = FeatureTemplate.OUTSIDE_LABEL;
		items[0].add(dummy);
		int[] ofs = new int[2];
		
		double normalizer = sqrt? 1.0 / Math.sqrt(x.features.length): 1.0;
		
		SparseVector sv = new SparseVector();
		for(int i = 0; i < x.features.length; i++) {
			ArrayList<DPItem>[] prevItems = items;
			items = new ArrayList[range.length];
			for(int j = 0; j < range.length; j++) {				
				int label = range[j];
				items[j] = makeDPItemKBest(x, k, fe, model, sv,
						i, label, prevItems, ofs);

				if(maxLoss && label != y.labels[i])
					for(DPItem item: items[j])
						item.totalScore += normalizer;
			}
		}
		
		ArrayList<DPItem> finalItems = makeDPItemKBest(x, k, fe, model, sv, 
				x.features.length, FeatureTemplate.OUTSIDE_LABEL, items, ofs);
		
		for(int i = 0; i < finalItems.size(); i++) {
			int[] labels = new int[x.features.length];
			DPItem item = finalItems.get(i);
			scores[i] = item.totalScore;
			for(int j = x.features.length - 1; j >= 0; j--) {
				item = item.prev;
				labels[j] = item.label;
			}
			results[i] = new Labeling(labels);
		}
		return finalItems.size();
	}
	
	private static class HeapItem implements Comparable<HeapItem> {
		
		DPItem item;

		ArrayList<DPItem> prevList;
		int position;
		
		HeapItem(DPItem item, ArrayList<DPItem> prevList, int position) {
			this.item = item;
			this.prevList = prevList;
			this.position = position;
		}
		
		public int compareTo(HeapItem o) {
			return -Double.compare(item.totalScore, o.item.totalScore);
		}
		
		public boolean equiv(Object o) {
			HeapItem other = (HeapItem) o;
			return prevList == other.prevList && position == other.position
			&& item.equals(other.item);
		}
		
	}
	
	static int countPops = 0;
	static int countRedundantPops = 0;
	
	private static ArrayList<DPItem> makeDPItemKBest(Sentence x, int k,
			FeatureTemplateSet fe, Model model, SparseVector sv, 
			int position, int label, ArrayList<DPItem>[] prevItems,
			int[] ofs) {

		ArrayList<DPItem> out = new ArrayList();
		
		sv.clear();

		ofs[0] = label;
		fe.extractFeatures(sv, x.features, ofs, position, 0); // TODO lift this out
		fe.extractFeatures(sv, x.features, ofs, position, 1);

		double staticScore = model.score(sv);
		
		PriorityQueue<HeapItem> q = new PriorityQueue();

		// initialization of the queue
		for(int j2 = 0; j2 < prevItems.length; j2++) {
			sv.clear();
			DPItem prev = prevItems[j2].get(0);
			ofs[1] = prev.label;
			fe.extractFeatures(sv, x.features, ofs, position, 2);
			double dynScore = prev.totalScore + model.score(sv);			
			DPItem item = new DPItem();
			item.prev = prev;
			item.totalScore = staticScore + dynScore;
			item.label = label;			
			q.add(new HeapItem(item, prevItems[j2], 0));
		}

		HeapItem pr = null;
		while(q.size() > 0 && out.size() < k) {
			HeapItem hi = q.poll();

			countPops++;
			if(pr != null && hi.equiv(pr)) {
				countRedundantPops++;
				continue;
			}
			pr = hi;
			
			out.add(hi.item);
			
			if(out.size() < k && hi.position < hi.prevList.size() - 1) {
				DPItem prev = hi.prevList.get(hi.position + 1);
				ofs[1] = prev.label;
				sv.clear();
				fe.extractFeatures(sv, x.features, ofs, position, 2);
				double dynScore = prev.totalScore + model.score(sv);
				DPItem item = new DPItem();
				item.prev = prev;
				item.totalScore = staticScore + dynScore;
				item.label = label;				
				q.add(new HeapItem(item, hi.prevList, hi.position + 1));
			}
		}
		
		if(out.size() == 0) {
			System.out.println("x = " + x);
			System.out.println("position = " + position);
			System.out.println("prevItems.length = " + prevItems.length);			
			throw new RuntimeException("!!!");
		}
		
		return out;
	}
	
}