package edu.claflin.finder.algo.clustering.struct.infomap_struct;

import edu.claflin.finder.logic.Graph;
import edu.claflin.finder.logic.Node;

import java.util.*;

public class AvgBitLength {

	private static final int MULTIPLIER = 10;
	private Graph graph;
	private ArrayList<Node> randomPath;
	private Map<Integer, List<Node>> modules; // i.e. communities
	private Map<Node, Integer> nodeToModule;

	public AvgBitLength(Graph graph, Map<Integer, List<Node>> modules) {
		this.graph = graph;
		this.randomPath = randomWalker(graph);
//		this.randomPath = testRandomWalker(graph);
		this.modules = modules;
		this.nodeToModule = new HashMap<>();
		for (Map.Entry<Integer, List<Node>> entry : modules.entrySet()) {
			int module = entry.getKey();
			List<Node> nodes = entry.getValue();
			for (Node n : nodes) {
				nodeToModule.put(n, module);
			}
		}
	}

	public void setGraph(Graph graph) {
		this.graph = graph;
	}

	public void setModules(Map<Integer, List<Node>> modules) {
		this.modules = modules;
	}

	/**
	 * Generates a path taken by random walker.
	 * Will be used to calculate average bit length
	 */
	private static ArrayList<Node> randomWalker(Graph g) {

		Random rand = new Random();
		int currentNode = rand.nextInt(g.getNodeCount());
		int nodeCount = g.getNodeCount();
		int traverseGraph = MULTIPLIER * nodeCount;
		ArrayList<Node> randomPath = new ArrayList(traverseGraph);
		randomPath.add(g.getNodeList().get(currentNode));

		for (int i = 0; i < traverseGraph; i++) {
			List<Node> neighbors = g.getNodeList().get(currentNode).getNeighbors();
			Node node = neighbors.get(rand.nextInt(neighbors.size())); // TODO: random distribution based on edge weight
			node.incrementFreq();
			currentNode = node.getIntegerIdentifier();
			randomPath.add(node);
		}
//		System.out.println();
//		for(int i = 0; i < randomPath.size(); i++) {
//			System.out.print("[" + randomPath.get(i).getIdentifier() + "] ");
//		}
//		System.out.println();

		return randomPath;
	}

	private static ArrayList<Node> testRandomWalker(Graph g) {
		int[] path = new int[]{6, 4, 5, 6, 4, 5, 6, 4, 6, 4, 6, 4, 6, 5, 6, 5, 4, 6, 5, 4, 1, 4, 6, 4, 1, 4, 5, 4, 6, 4, 6};
		ArrayList<Node> randomPath = new ArrayList(path.length);
		for (int n : path) {
			randomPath.add(g.getNode(String.valueOf(n)));
		}

//		System.out.println();
//		for(int i = 0; i < randomPath.size(); i++) {
//			System.out.print("[" + randomPath.get(i).getIdentifier() + "] ");
//		}
//		System.out.println();

		return randomPath;
	}

	/**
	 * Calculates average bit length for random path of this graph
	 *
	 * @return
	 */
	public double avgBitLength() {

		Map<Integer, Integer> moduleFrequency = new HashMap<>();
		int currentModule = -1;
		for (Node n : randomPath) {
			// increment node frequency
			n.incrementFreq();

			// increment module frequency
			int module = nodeToModule.get(n);
			if (module != currentModule) {
				int freq = moduleFrequency.getOrDefault(module, 0);
				moduleFrequency.put(module, freq + 1);
				currentModule = module;
			}
		}

		double totalModuleFreq = totalModuleFreq(moduleFrequency);
		double totalExitCodes = totalExitCodes(moduleFrequency);
		double totalNodeFreq = totalNodeFreq(graph.getNodeList());
		double totalModuleCode = totalModuleCode(modules, moduleFrequency);

		double totalBitLength = totalModuleFreq * Math.log(totalModuleFreq)
				- 2 * totalExitCodes - totalNodeFreq
				+ totalModuleCode;

		return totalBitLength / randomPath.size();
	}

	private double totalModuleCode(Map<Integer, List<Node>> modules, Map<Integer, Integer> moduleFrequency) {
		double result = 0;
		for (Map.Entry<Integer, List<Node>> entry : modules.entrySet()) {
			int module = entry.getKey();
			List<Node> moduleNodes = entry.getValue();

			int moduleFreq = moduleFrequency.getOrDefault(module, 0);
			double totalModuleNodeFreq = totalModuleNodeFreq(moduleNodes);
			double moduleCode = moduleFreq + totalModuleNodeFreq;
			result += moduleCode * Math.log(moduleCode);
		}
		return result;
	}

	private double totalModuleFreq(Map<Integer, Integer> moduleFrequency) {
		double result = 0;
		for (int freq : moduleFrequency.values()) {
			result += freq;
		}
		return result;
	}

	private double totalExitCodes(Map<Integer, Integer> moduleFrequency) {
		double result = 0;
		for (int freq : moduleFrequency.values()) {
			if (freq != 0)
				result += freq * Math.log(freq);
		}
		return result;
	}

	private double totalNodeFreq(List<Node> nodes) {
		double result = 0;
		for (Node n : nodes) {
			if (n.getFreq() != 0)
				result += n.getFreq() * Math.log(n.getFreq());
		}
		return result;
	}

	private double totalModuleNodeFreq(List<Node> moduleNodes) {
		double result = 0;
		for (Node n : moduleNodes) {
			result += n.getFreq();
		}
		return result;
	}

}
