package parser.sampling;

import java.util.*;

import parser.*;
import utils.Utils;

public class RandomWalkSampler {
	
	Random r;
	
	public RandomWalkSampler(int seed) {
		r = new Random(seed);
	}
	
	public RandomWalkSampler() {
		r = new Random();
	}
	
	public RandomWalkSampler(Random r) {
		this.r = r;
	}
	
	
    public DependencyInstance randomWalkSampling(DependencyInstance inst,
    		LocalFeatureData lfd, boolean addLoss) {
    	int len = inst.length;
    	
		DependencyInstance predInst = new DependencyInstance(inst);
		predInst.heads = new int[len];
		predInst.depids = new int[len];
    	
    	boolean[] inTree = new boolean[len];
    	inTree[0] = true;
    	for (int i = 0; i < len; i++) {
    		predInst.heads[i] = -1;
    	}
    	
    	for (int i = 1; i < len; i++) {
    		int curr = i;
    		while (!inTree[curr]) {
    			// sample new head
    			ArrayList<Double> score = new ArrayList<Double>();
    			ArrayList<Integer> depList = new ArrayList<Integer>();

    			for (int candH = 0; candH < len; candH++) {
    				if (lfd.isPruned(candH, curr))
    					continue;
    				depList.add(candH);
    				double s = lfd.getArcScore(candH, curr);
    				if (addLoss && inst.depids[i] != candH) {
    					// cost augmented
    					s += 1.0;
    				}
    				score.add(s);
    			}

    			int sample = samplePoint(score, r);
    			predInst.depids[curr] = depList.get(sample);
    			curr = predInst.depids[curr];
    			
    			if (predInst.depids[curr] != -1 && !inTree[curr]) {
    				cycleErase(predInst.depids, curr);
    			}
    		}
    		curr = i;
    		while (!inTree[curr]) {
    			inTree[curr] = true;
    			curr = predInst.depids[curr]; 
    		}
    	}
    	
    	return predInst;
    }
    
    private void cycleErase(int[] dep, int i) {
    	while (dep[i] != -1) {
    		int next = dep[i];
    		dep[i] = -1;
    		i = next;
    	}
    }
    
    private int samplePoint(ArrayList<Double> score, Random r) {
    	double sumScore = -1e30;
    	for (int i = 0; i < score.size(); i++) {
    		sumScore = Utils.logSumExp(sumScore, score.get(i));
    	}
    	double[] prob = new double[score.size()];
    	prob[0] = Math.exp(score.get(0) - sumScore);
    	for (int i = 1; i < score.size(); i++)
    		prob[i] = prob[i - 1] + Math.exp(score.get(i) - sumScore);
    	if (Math.abs(prob[prob.length - 1] - 1.0) > 1e-4) {
    		System.out.println("sample point bug");
    		System.exit(0);
    	}
    	double p = r.nextDouble();
    	int ret = 0;
    	for (; ret < prob.length; ret++) {
    		if (p < prob[ret])
    			break;
    	}
    	return ret;
    }
}

@@ -39,17 +39,47 @@
 util.*;%0A
+import gnu.trove.list.array.*;
 %0A
-
 import p
@@ -884,32 +884,47 @@
 new head%0A    %09%09%09
+%0A    %09%09%09TDouble
 ArrayList%3CDouble
@@ -912,32 +912,24 @@
 bleArrayList
-%3CDouble%3E
  score = new
@@ -925,24 +925,31 @@
 score = new 
+TDouble
 ArrayList%3CDo
@@ -945,24 +945,16 @@
 rrayList
-%3CDouble%3E
 ();%0A    
@@ -956,16 +956,20 @@
 %0A    %09%09%09
+TInt
 ArrayLis
@@ -969,25 +969,16 @@
 rrayList
-%3CInteger%3E
  depList
@@ -984,16 +984,20 @@
 t = new 
+TInt
 ArrayLis
@@ -1001,17 +1001,8 @@
 List
-%3CInteger%3E
 ();%0A
@@ -1933,16 +1933,23 @@
 lePoint(
+TDouble
 ArrayLis
@@ -1953,16 +1953,8 @@
 List
-%3CDouble%3E
  sco

add random walk sampler