package isampling;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Random;
import java.util.Vector;

public class ImportanceSampler {

	private Vector<String> labels;
	private Vector<Integer> duplicates;
	private Vector<Pair<Double, Double>> minimalSamplingBounds;
	private Vector<MinimalSolution> orderedSols;
	private double samplingLen = -1;
	private HashSet<Integer> allNodes;
	private Vector<Integer> entitySampler;
	private int entitiesNum = -1;

	private String inputFile;
	private String seedFile;
	private String outFile;
	private int iterations;
	private int maxCoverLen;

	public ImportanceSampler(String inputF, String seedF, String outF, int itNum) {
		inputFile = inputF;
		seedFile = seedF;
		outFile = outF;
		iterations = itNum;

		this.maxCoverLen = 0;
		labels = new Vector<String>();
		duplicates = new Vector<Integer>();
		// holds the bounds for the initial minimal choice
		minimalSamplingBounds = new Vector<Pair<Double, Double>>();
		orderedSols = new Vector<MinimalSolution>();
		samplingLen = -1;
		allNodes = new HashSet<Integer>();
		entitySampler = new Vector<Integer>();
		entitiesNum = 0;
	}

	public static void main(String args[]) {

		// ImportanceSampler IS = new ImportanceSampler("example1", "seed1",
		// "out1", 100000, 2, 5);

		ImportanceSampler IS = new ImportanceSampler(args[0], args[1], args[2],
				Integer.parseInt(args[3]));

		IS.run();

	}

	/**
	 * Runs the Sampler
	 * 
	 * @return the counts for the give entities
	 */
	public float[] run() {

		// load entities to be evaluated
		loadEntities();

		// compute counts
		float[] counts = getCounts();

		// write counts
		writeCounts(counts);

		return counts;
	}

	/**
	 * Writes the counts to the specified output file (global outFile parameter)
	 * 
	 * @param counts
	 */
	private void writeCounts(float counts[]) {

		FileWriter fw = null;

		try {
			fw = new FileWriter(outFile);
			int i = 0;
			for (float c : counts) {
				fw.write(labels.get(i) + "\t" + c + "\n");
				++i;
			}
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(1);
		} finally {
			try {
				fw.close();
			} catch (Exception ex) {
				ex.printStackTrace();
				System.exit(1);
			}
		}

	}

	/**
	 * Retrieves the kmin of an Sk
	 * 
	 * Looking for the first minimal solution that is a subset of Sk. This means
	 * that the sol
	 * 
	 * @param Sk
	 *            : Sk
	 * 
	 * @param k
	 *            : This is the index of Hk. We already know that the Hk that we
	 *            flipped to get Sk is a subset of Sk, so there is now point in
	 *            looking beyond this.
	 * 
	 * @return
	 */
	private boolean canonicalCheck(MinimalSolution Sk, int k) {
		// check minimal sols in order
		for (int i = 0; i <= k; ++i) {

			// get the indices of the 1-bits in the current solution
			HashSet<Integer> myNodes = this.orderedSols.get(i).nodes;

			if (Sk.nodes.containsAll(myNodes)) {

				if (i == k)
					return true;
				else
					return false;

			}
		}
		// System.err.println("NO SUBSETS OF Sk FOUND. THIS SHOULD NEVER HAPPEN");
		System.exit(1);
		return false;
	}

	/**
	 * 
	 */
	private float[] getCounts() {

		Random R = new Random();

		this.readMinimals();

		samplingLen = this.getSamplingBounds();

		float counts[] = new float[entitiesNum];

		int total = 0;
		while (total < iterations) {

			// if (total % 100 == 0) {
			// System.out.println(total);
			// }
			++total;

			// ---------------------------------------------------------
			// get the respective minimal solution
			MinimalSolution sol = sampleMinimal(R);

			int extensions = Math.min(sol.missing.size(), maxCoverLen
					- sol.nodes.size());

			HashMap<Integer, Integer> extendedNodeSet = new HashMap<Integer, Integer>();
			for (int N : sol.nodes) {
				extendedNodeSet.put(N, 1);
			}
			extendMinimal(extensions, R, extendedNodeSet);
			// ---------------------------------------------------------

			if (duplicationCheck(sol, extendedNodeSet, R)
					&& this.canonicalCheck(sol, sol.order)) {
				for (int N : sol.nodes) {
					counts[N] += extendedNodeSet.get(N) * 1f
							/ (duplicates.get(N));
				}
			}

		}

		return counts;
	}

	/**
	 * Samples a minimal cover based on the precomputed bounds
	 * 
	 * @param R
	 * @return a copy of the sampled cover
	 */
	private MinimalSolution sampleMinimal(Random R) {
		double coin = R.nextDouble() * samplingLen;

		int pos = -1;
		if (samplingLen == 0) {
			pos = R.nextInt(orderedSols.size());
		}

		else {

			// for each bound-pair
			for (int j = 0; j < minimalSamplingBounds.size(); ++j) {
				// if the coin falls within the bounds
				if (coin >= minimalSamplingBounds.get(j).key
						&& coin < minimalSamplingBounds.get(j).val) {
					pos = j;
					break;
				}
			}
		}

		return new MinimalSolution(orderedSols.get(pos));

	}

	/**
	 * An extra check used to guarantee that the correct counts are obtained in
	 * the presence of multiplicities.
	 * 
	 * @param mSol
	 * @param extNodeSet
	 * @param R
	 * @return
	 */
	private boolean duplicationCheck(MinimalSolution mSol,
			HashMap<Integer, Integer> extNodeSet, Random R) {
		double prod = 1;

		for (int N : extNodeSet.keySet()) {
			if (mSol.nodes.contains(N)) {
				prod *= extNodeSet.get(N);
			} else {
				mSol.nodes.add(N);
			}
		}

		prod = 1 / prod;

		if (R.nextDouble() <= prod) {
			return true;
		}
		return false;
	}

	/**
	 * 
	 * Read the minimal covers. One cover per line. Space separated lines. Each
	 * entry represents the order of the respective entity in the inputFile
	 * (starting from 0)
	 * 
	 */
	private void readMinimals() {

		BufferedReader br = null;

		try {
			br = new BufferedReader(new FileReader(seedFile));

			for (int i = 0; i < entitiesNum; ++i) {
				allNodes.add(i);
			}

			int lineCnt = 0;
			// for each line (solution)
			while (br.ready()) {

				String line = br.readLine().trim();
				if (line.length() > 0) {

					// holds the indices of the 1-bits for each solution
					HashSet<Integer> myAces = new HashSet<Integer>();

					// go over the nodes in the solution and set the
					// corresponding bits to 1. Remember the 1-indexes of each
					// sol.
					double product = 1;
					String[] toks = line.split(" ");
					this.maxCoverLen = Math.max(this.maxCoverLen, toks.length);
					for (String tok : toks) {
						int id = Integer.parseInt(tok);
						myAces.add(id);
						product *= duplicates.get(id);
					}

					MinimalSolution solution = new MinimalSolution(myAces,
							allNodes, lineCnt);

					solution.product = product;

					orderedSols.add(solution);

				}

				lineCnt++;
			}
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(1);

		} finally {
			try {
				br.close();
			} catch (Exception ex) {
				ex.printStackTrace();
				System.err.println("Could not close file");
				System.exit(1);
			}
		}

	}

	/**
	 * Extends the set of nodes in a minimal cover
	 * 
	 * @param extensionsNum
	 * @param R
	 * @param entitiesInCover
	 */
	private void extendMinimal(int extensionsNum, Random R,
			HashMap<Integer, Integer> entitiesInCover) {

		Collections.shuffle(entitySampler);

		// for each node
		for (int x = 0; x < extensionsNum; ++x) {

			int ent = entitySampler.get(R.nextInt(entitySampler.size()));

			Integer freq = entitiesInCover.get(ent);

			// all copies of this entity already in there
			if (freq != null && freq >= duplicates.get(ent)) {
				--x;
				continue;
			}

			if (R.nextFloat() <= 0.5) {
				if (freq == null) {
					entitiesInCover.put(ent, 1);
				} else {
					entitiesInCover.put(ent, freq + 1);
				}
			}
		}
	}

	/**
	 * Prepare the sampling bounds for the minimal covers.
	 */
	private double getSamplingBounds() {

		int maxBase = Integer.MIN_VALUE;
		int minExpon = Integer.MAX_VALUE;
		int size, base, expon;
		double start = 0;

		// go over the minimal solution and compute the minimum exponent
		for (int j = 0; j < orderedSols.size(); ++j) {

			size = orderedSols.get(j).nodes.size();

			base = size;// (V - size);//*orderedSols.get(j).product ;
			expon = (entitiesNum - size);
			maxBase = Math.max(base, maxBase);
			minExpon = Math.min(expon, minExpon);
		}

		// for each minimal solution
		for (int j = 0; j < orderedSols.size(); ++j) {

			int x = orderedSols.get(j).nodes.size();

			double d0 = ((entitiesNum - x) * (maxCoverLen - maxBase));// /
			double d1 = ((entitiesNum - maxBase) * (maxCoverLen - x));

			double d2 = (entitiesNum - x) * Math.E / (maxCoverLen - x);

			double len = Math.pow(d0 / d1, maxCoverLen - maxBase)
					* Math.pow(d2, maxBase - x) * orderedSols.get(j).product;

			Pair<Double, Double> newPair = new Pair<Double, Double>(start,
					start + len);

			minimalSamplingBounds.add(newPair);
			start = start + len;

		}
		return start;

	}

	/**
	 * Reads the entities to be evaluated. One line per entity. Tab-separated:
	 */
	private void loadEntities() {

		BufferedReader br = null;
		String line, toks[];
		int cnt = 0;
		try {

			br = new BufferedReader(new FileReader(inputFile));
			while (br.ready()) {
				line = br.readLine();
				toks = line.split("\t");
				labels.add(toks[0]);
				duplicates.add(Integer.parseInt(toks[1]));
				for (int i = 0; i < Integer.parseInt(toks[1]); ++i) {
					entitySampler.add(cnt);
				}
				++cnt;
				entitiesNum++;
			}

		} catch (Exception e) {
			e.printStackTrace();
			System.exit(1);
		} finally {
			try {
				br.close();
			} catch (Exception ex) {
				ex.printStackTrace();
				System.err.println("Could not close file");
				System.exit(1);
			}
		}

	}

}
