package gsearch;

import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;

import misc.Common;
import graph.Graph;

public class RWR {
	public static int MAX_ITER=100;
	public static double L1_THRESH = 0.001;
	public static double LAZY = 0.5;
	
	public static double[] eilb(Graph g, 
							  HashMap<Integer, Double> rest, 
							  double r, HashSet<Integer> K, HashSet<Integer> F ) {
		if (null==K) {
			K = new HashSet<Integer>();
			for (int i=0;i<g.getn();i++) K.add(i);
			F = new HashSet<Integer>();
		}
		double[] p = rwlb(g,rest,r,K,F);
		for (Integer i: F) if (g.wn[i]>0) p[i] /= g.wn[i];
		for (Integer i: K) if (g.wn[i]>0) p[i] /= g.wn[i];
		return p;
	}

	public static double[] rwlb(Graph g, 
			  HashMap<Integer, Double> rest, 
			  double r, HashSet<Integer> K, HashSet<Integer> F) {
		if (null==K) {
			K = new HashSet<Integer>();
			for (int i=0;i<g.getn();i++) K.add(i);
			F = new HashSet<Integer>();
		}
		assert (0 <= r && r <= 1.0) : "Invalid Restart Pobability";
		assert (Math.abs(Common.sum(rest.values()) - 1.0) < L1_THRESH) : "Invalid Restart Distribution";
		for (Integer n: rest.keySet()) {
			if (rest.get(n) > 0) assert (K.contains(n)): "Unknown restart node" + n;
		}
		
		double[] p = new double[g.getn()];
		double[] pn = new double[g.getn()];
		for(Integer n: rest.keySet()) {
			p[n] = rest.get(n);
		}
		
		
		boolean conv = false;
		int it = 0;
		while (!conv) {
			// p = (1-r)*A
			for (Integer i: K) {
				for (int j = g.ind[i]; j < g.ind[i+1]; j++) {
					pn[g.endv[j]] += (1.0 - r)*g.we[j]*p[i]/g.wn[i];
				}
			}
			// sinks have loops
			for (Integer n: F) {
				pn[n] += (1.0 - r)*p[n];
			}
			// Handle restart
			for (Integer n: rest.keySet()) pn[n] += r*rest.get(n);
			
			//assert (Math.abs(Common.sum(pn) - 1.0) <= L1_THRESH) : "Not normalized new vector p_new. norm=" + Common.sum(pn);
			conv = (Common.l1(pn,p) <= L1_THRESH) || (it > MAX_ITER);
			it++;
			for(int i=0; i< p.length; i++) p[i] = pn[i];
			Arrays.fill(pn, 0.0);
		}
		// Estimate the sinks as the sum of the incoming flow
		if (F.size() > 0) {
			for (Integer n: F) p[n] = 0.0;
			for (Integer i: K) {
				for (int j = g.ind[i]; j < g.ind[i+1]; j++) {
					if (F.contains(g.endv[j])) {
						p[g.endv[j]] += (1.0 - r)*g.we[j]*p[i]/g.wn[j];
					}
				}
			}
		}
		return p;
	}
	
	public static double[] eiub(Graph g, 
			  HashMap<Integer, Double> rest, 
			  double r, HashSet<Integer> K, HashSet<Integer> F ) {
		if (null==K) {
			K = new HashSet<Integer>();
			for (int i=0;i<g.getn();i++) K.add(i);
			F = new HashSet<Integer>();
		}
		double[] p = rwub(g,rest,r,K,F);
		for (Integer i: F) if (g.wn[i]>0) p[i] /= g.wn[i];
		for (Integer i: K) if (g.wn[i]>0) p[i] /= g.wn[i];
		return p;
	}
	
	public static double[] rwub(Graph g, 
			  HashMap<Integer, Double> rest, 
			  double r, HashSet<Integer> K, HashSet<Integer> F) {
		if (null==K) {
			K = new HashSet<Integer>();
			for (int i=0;i<g.getn();i++) K.add(i);
			F = new HashSet<Integer>();
		}
		assert (0 <= r && r <= 1.0) : "Invalid Restart Pobability";
		assert (Math.abs(Common.sum(rest.values()) - 1.0) < L1_THRESH) : "Invalid Restart Distribution";
		for (Integer n: rest.keySet()) {
			if (rest.get(n) > 0) assert (K.contains(n)): "Unknown restart node" + n;
		}
		
		double rl = r/(2-r);
		double[] p = new double[g.getn()];
		double[] s = new double[g.getn()];
		
		for (Integer n: rest.keySet()) s[n] = rest.get(n);
		
		double tdelta;
		tdelta = 1.0;
		while(tdelta > L1_THRESH) {
			tdelta = 0.0;
			for (Integer i:K) {
				p[i] += rl*s[i]; 
				tdelta += Math.abs(rl*s[i]);
				for (int j = g.ind[i]; j < g.ind[i+1]; j++) {
					tdelta += Math.abs(((1.0-rl)*s[i]*g.we[j])/(2.0*g.wn[i]));
					s[g.endv(j)] += ((1.0-rl)*s[i]*g.we[j])/(2.0*g.wn[i]);
					assert(F.contains(g.endv[j]) || K.contains(g.endv[j])): "Incorrect K and F sets";
				}
				tdelta += Math.abs(s[i]-(1.0-rl)*s[i]/2.0);
				s[i] = (1.0-rl)*s[i]/2.0;
			}
		}
		
		HashSet<Integer> KF = new HashSet<Integer>();
		KF.addAll(K); KF.addAll(F);
		double eps = 0.0;
		for (Integer i:KF) {
			if (s[i]/g.wn[i] >= eps) 
				eps = s[i]/g.wn[i];  
		}
		double[] res = new double[p.length];
		for (Integer i:KF) res[i] = p[i] + eps*s[i];
		return res;
	}
	
	
}
