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

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

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

/* 
 * Implementation of the Louvain algorithm as briefly described by Mark Newman in
 * pages 511-512 of "Networks Second Edition". Uses similar methods to the leading
 * eigenvector algorithm except for the modularity calculation, it had to be changed
 * since Louvain does not split into smaller groups so basing modularity on the
 * modularity matrix would make it take O(m^2) at minimum, where m = edges. 
 * 
 * The new equation for calculating changes in modularity can be found at 
 * https://en.wikipedia.org/wiki/Louvain_method
 * 
 * The bread and butter of the algorithm is the mergeGroups() method. At the start, 
 * each node is place into its own group, then the mergeGroup() method agglomerates
 * individual nodes into groups. Then it keeps repeating the merges until no
 * additional groups could be formed. 
 * 
 * A HashMap groupMap is used to keep track of which group has which nodes in it in
 * order to massively speed up iterations. Overall, the running time is much better
 * than the leading eigenvector algorithm and the results are extremely similar if
 * not identical in some cases. For "COAD.txt", our largest available network,
 * was fully processed in 0.587 seconds when I tested it.
 * 
 * Terminology:
 * modularity: A score which quantifies the quality of connections in a split network, high modularity
 * 		means the quality of the found communities is high, low modularity means the quality is low,
 * 		negative modularity scores mean that the proposed split is worse than leaving the network as is
 * 		so the split should not be performed
 * 
 * modularity delta: The change in modularity. In this algorithm we don't really care about overall
 * 		modularity, so the initial modularity is never calculated. Since the initial step of the algorithm
 * 		places all nodes in their own community, the modularity is guaranteed to be negative if the
 * 		network has even a single connection. As such, we only care about changes in modularity (delta).
 * 
 * @author Samuel Perez
 */

public final class LouvainAlgorithm {
	private Graph graph;
	private int nodeCount;
	private int[] groupVector;
	private double[][] adjacencyMatrix;
	private HashMap<Integer, List<Node>> communities = new HashMap<>();
	private HashMap<Integer, ArrayList<Integer>> groupMap = new HashMap<>();
	
	public LouvainAlgorithm(Graph g) {
		graph = g;
		nodeCount = g.getNodeCount();
		adjacencyMatrix = getAdjacencyMatrix();
		
		// initialize the group map, which has group label as keys, and the list
		// of node integer identifiers that belong to that group as a value,
		// also a group vector purely for iteration
		groupVector = new int[graph.getNodeCount()];
		
		for(int i = 0; i < graph.getNodeCount(); i++) {
			ArrayList<Integer> newList = new ArrayList<>();
			newList.add(i);
			groupMap.put(i + 1, newList);
			groupVector[i] = i + 1;
		}
		
		// begin algorithm iterations, end them when there are no group merges which result
		// in a higher modularity
		while(true) {
			// mergeGroups performs a merge iteration and returns true if an improvement was found, else false
			if(!mergeGroups()) {
				break;
			}
		}
		
		// normalizes the group vector so that communities are labeled from 1 to C where C
		// is the total number of communities found
		normalizeGroupVector();
		
		// build the list of nodes from the final group vector
		for(int i = 0; i < this.groupVector.length; i++) {
			if(!this.communities.containsKey(this.groupVector[i])) {
				ArrayList<Node> newList = new ArrayList<>();
				newList.add(this.graph.getNodeList().get(i));
				this.communities.put(this.groupVector[i], newList);
			}
			else {
				this.communities.get(this.groupVector[i]).add(this.graph.getNodeList().get(i));
			}
		}
	}
	
	// performs a single iteration of the algorithm, which consists of merging groups into
	// larger communities if possible
	private boolean mergeGroups() {
		boolean improved = false;
		
		// keeps track of which groups have had nodes moved into them, since if a group had nodes
		// moved into it this iteration, it should not have any nodes moved out
		ArrayList<Integer> grew = new ArrayList<>();
		ArrayList<Integer> checked = new ArrayList<>();
		
		// iterate through each group; using group vector instead of group map to avoid any hassle
		// with the group map since you can't edit the group map while iterating through it;
		// to avoid extra computations, the checked array list is used to keep track of which
		// groups have been processed so they aren't processed again
		for(int i = 0; i < groupVector.length; i++) {
			// if this group is the result of a merge during this iteration, or
			// this group was already checked this iteration, don't bother checking
			if(grew.contains(groupVector[i]) || checked.contains(groupVector[i])) {
				continue;
			}
			else {
				checked.add(groupVector[i]);
			}
			
			// save the original group so it can be restored after each swap to try a new group
			GroupData originalData = new GroupData(groupVector[i], groupMap.get(groupVector[i]));
			
			// initial value doesn't matter
			double maxDelta = 0;
			
			// value holding the best group to change the group nodes to, if it stays at -1
			// that means no good merge was found so the group vector should not be updated
			int maxGroup = -1;
			
			for(int targetGroupID : groupMap.keySet()) {
				// 1. If there are no neighbors in common skip this group; existNeighbors
				//    also returns false if the groups are the same
				if(originalData.groupLabel == targetGroupID || !existsNeighbor(originalData.groupLabel, targetGroupID)) {
					continue;
				}
				
				// 2. Calculate the modularity delta
				double currDelta = modularityDelta(originalData.groupLabel, targetGroupID);
				
				// 3. If delta is better than current max, update current max as well as group for which the max was found
				if(currDelta > maxDelta) {
					maxDelta = currDelta;
					maxGroup = targetGroupID;
				}
			}
			
			// 4. After all iterations are done, if an improvement was found, commit the changes
			if(maxGroup != -1) {
				// change the group vector and group map
				for(int nodeID : originalData.memberIndices) {
					this.groupVector[nodeID] = maxGroup;
					this.groupMap.get(maxGroup).add(nodeID);
				}

				groupMap.remove(originalData.groupLabel);

				// update which cluster grew so we don't try to destroy it by merging it into another group
				// during this iteration
				grew.add(maxGroup);
				
				// at least 1 improvement was found, so we can run another iteration
				improved = true;
			}
		}
		
		return improved;
	}
	
	// returns true if the two groups identified by their labels have at least 1 neighbor in common
	// running time complexity is O(g1n + g2), where g1n is the number of neighbors of all nodes in
	// group 1 and g2 is the number of nodes in group 2.
	private boolean existsNeighbor(int group1, int group2) {
		ArrayList<Integer> group1Indices = groupMap.get(group1);
		ArrayList<Integer> group2Indices = groupMap.get(group2);
		
		HashMap<Node, Boolean> group1Neighbors = new HashMap<>();
		
		for(int index : group1Indices) {
			List<Node> neighbors = graph.getNodeList().get(index).getNeighbors();
			
			for(Node neighbor : neighbors) {
				if(!group1Neighbors.containsKey(neighbor)) {
					group1Neighbors.put(neighbor, true);
				}
			}
		}
		
		for(int index : group2Indices) {
			if(group1Neighbors.containsKey(graph.getNodeList().get(index))) {
				return true;
			}
		}
		
		return false;
	}
	
	// simple algorithm to generate an adjacency matrix from the graph instance variable
	private double[][] getAdjacencyMatrix() {
		double[][] adjMatrix = new double[nodeCount][nodeCount];
		int sourceIntegerIdentifier = 0;
		int targetIntegerIdentifier = 0;
		double edgeWeight = 0;
		
		for(Edge edge: graph.getEdgeList()) {
			sourceIntegerIdentifier = edge.getSource().getIntegerIdentifier();
			targetIntegerIdentifier = edge.getTarget().getIntegerIdentifier();
			edgeWeight = edge.getData();
			
			adjMatrix[sourceIntegerIdentifier][targetIntegerIdentifier] = edgeWeight;
			adjMatrix[targetIntegerIdentifier][sourceIntegerIdentifier] = edgeWeight;
		}
		
		return adjMatrix;
	}
	
	// implements the modularity delta equation found at https://en.wikipedia.org/wiki/Louvain_method
	// group1 is the group being merged, group2 is the group being merged into. Running time is
	// O(max(group1Size, group2Size)^2), which in the grand scheme of things is insignificant
	// due to the fact that the algorithm starts with groups of size 1 and groups never really
	// get too large, so for large sparse networks this might as well be constant.
	private double modularityDelta(int group1, int group2) {
		// calculate the sum of the weight of all edges in the target group
		double S_in = 0;
		ArrayList<Integer> group2Indices = groupMap.get(group2);
		
		for(int i = 0; i < group2Indices.size(); i++) {
			for(int j = 0; j < group2Indices.size(); j++) {
				S_in += adjacencyMatrix[group2Indices.get(i)][group2Indices.get(j)];
			}
		}
		
		// divide by 2 since the adjacency matrix is symmetric so edge weights are added twice
		S_in = S_in / 2;
		
		// calculate the sum of the weight of all vertices in the target group
		double S_tot = 0;
		
		for(int i = 0; i < group2Indices.size(); i++) {
			S_tot += graph.getNodeList().get(group2Indices.get(i)).getWeight();
		}
		
		// calculate the sum of the weight of all nodes in the source group
		double C_tot = 0;
		ArrayList<Integer> group1Indices = groupMap.get(group1);
		
		for(int i = 0; i < group1Indices.size(); i++) {
			C_tot += graph.getNodeList().get(group1Indices.get(i)).getWeight();
		}
		
		// calculate the weight of the links between source nodes and target nodes
		double K_in = 0;
		
		for(int i = 0; i < group1Indices.size(); i++) {
			for(int j = 0; j < group2Indices.size(); j++) {
				K_in += adjacencyMatrix[group1Indices.get(i)][group2Indices.get(j)];
			}
		}
		
		// calculate each component of the modularity delta function separately;
		// five is missing due to the fact that in this method it'll always be 
		// zero so it's not calculated, but some derivations found online include it
		double one = (S_in + (2 * K_in)) / (2 * graph.getGraphWeight());
		double two = Math.pow((S_tot + C_tot) / (2 * graph.getGraphWeight()), 2);
		double three = (S_in / (2 * graph.getGraphWeight()));
		double four = Math.pow((S_tot / (2 * graph.getGraphWeight())), 2);
		double six =  Math.pow(C_tot / (2 * graph.getGraphWeight()), 2);
		
		// put the components together according to the equation
		double result = (one - two) - ((three - four) - six);
		
		return result;
	}
	
	// normalize the group vector so that if there are n communities, the values in the group vector are limited to
	// numbers between 1 and n
	private void normalizeGroupVector() {
		// list of unique values in the group vector
		ArrayList<Integer> values = new ArrayList<>();
		
		for(int i = 0; i < this.groupVector.length; i++) {
			if(!values.contains(this.groupVector[i])) {
				values.add(this.groupVector[i]);
			}
		}
		
		// map such that map[n] = n + 1, so essentially an array of size n + 1 that starts at 1 and ends at n + 1
		int[] map = new int[values.size()];
		
		for(int i = 0; i < map.length; i++) {
			map[i] = i + 1;
		}
		
		// indexes group vector elements in the generated integer map to convert them to the appropriate value
		// for example, if 5 is the first unique value found in the group vector, all 5s in the group vector will be converted to 1s
		// and if 11 is the third unique value found, then all 11s will be converted to 3s
		for(int i = 0; i < this.groupVector.length; i++) {
			this.groupVector[i] = map[values.indexOf(groupVector[i])];
		}
	}
	
	public HashMap<Integer, List<Node>> getCommunities(){
		return communities;
	}
	
	private class GroupData{
		int groupLabel;
		ArrayList<Integer> memberIndices;
		
		private GroupData(int groupLabel, ArrayList<Integer> memberIndices){
			this.groupLabel = groupLabel;
			this.memberIndices = memberIndices;
		}
	}
}
