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

import java.util.ArrayList;
import java.util.Arrays;
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 leading eigenvector algorithm as described by Mark Newman. 
 * Works by splitting a network using the signs of the leading eigenvector of the
 * network's modularity matrix.
 * 
 * Terminology:
 * modularity matrix: A matrix generated from the adjacency matrix of a network which
 * 		shows how "strong" a connection is between two nodes. If modMatrix[i][j] is large,
 * 		that means nodes i and j belong together, if it's small (can be negative) that means
 * 		nodes i and j probably don't belong together.
 * 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
 * leading eigenvector: The eigenvector corresponding to the most positive eigenvalue of a matrix. This
 * 		algorithm calculates the leading eigenvector of the modularity matrices at each step and then
 * 		splits groups according to the signs of the elements in the leading eigenvector. For example,
 * 		if leading eigenvector = [0.2, -1.11, 5.7, 6.8, -3.22] that means nodes 0, 2, and 3 belong in
 * 		one group and nodes 1 and 4 belong to another group.
 * 
 * @author Samuel Perez
 */

public final class LeadingEigenvectorAlgorithm {
	private Graph graph;
	private int[] groupVector;
	private int nodeCount;
	private int groupLabel = 1;
	private double[][] modularityMatrix;
	private HashMap<Integer, List<Node>> communities = new HashMap<>();
	
	public LeadingEigenvectorAlgorithm(Graph g) {
		// Initialize all instance variables
		this.graph = g.uniqueCopy();
		this.nodeCount = this.graph.getNodeCount();
		this.groupVector = new int[this.nodeCount];
		this.modularityMatrix = getModularityMatrix();
		
		this.groupVector = getGroupVector(this.modularityMatrix);
		stepOptimization();

		// updating the instance group vector with appropriate labels; this just converts -1 to 1 and 1 to 2
		// i.e. if groupVector[15] == 1, 1 gets converted to 2, which means node 15 is in group 2
		for(int i = 0; i < this.groupVector.length; i++) {
			if(this.groupVector[i] == -1) {
				this.groupVector[i] = this.groupLabel;
			}
			else {
				this.groupVector[i] = this.groupLabel + 1;
			}
		}
		
		// updates the group label tracker so that the next splits have a different label
		this.groupLabel += 2;
		
		// checks if the modularity score is positive; if it's negative do nothing
		if(getModularityScore(this.modularityMatrix, this.groupVector) > 0) {
			split(1);
			split(2);
			
			// normalizes the group vector so that if there are n communities, the community labels go from 1 to n
			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));
				}
			}
		}
		else {
			this.communities.put(1, this.graph.getNodeList());
		}
	}
	
	// Recursive function which does the bulk of the work for the algorithm.
	//
	// @param	groupLabel	the label for the group you want to split
	private void split(int groupLabel) {
		// gets the modularity matrix for the group you want to split
		MappedModularityMatrix groupModMatrix = getGroupModMatrix(groupLabel);
		
		if(groupModMatrix.matrix.length != 0) {
			// using the modularity matrix calculated above, calculate the localGroupVector for that group,
			// which tells you how to split the group
			int[] localGroupVector = getGroupVector(groupModMatrix.matrix);
			
			// calculate modularity for the division
			double unscaledModularityDelta = 0;
			
			// this is a O(N) way of checking if the group vector only has 1 unique value
			int groupVectorSum = 0;
			
			for(int i = 0; i < localGroupVector.length; i++) {
				groupVectorSum += localGroupVector[i];
			}
			
			// only happens when all 1s or all -1s, meaning the best split is no split at all
			if(Math.abs(groupVectorSum) == localGroupVector.length) {
				unscaledModularityDelta = -1d;
			}
			// groupVector has both 1s and -1s, so calculate the change in modularity 
			else {
				unscaledModularityDelta = getModularityScore(groupModMatrix.matrix, localGroupVector);
			}
			
			// if the generated split results in a positive change to modularity, perform the split by updating the instance
			// group vector and recursively split the 2 newly generated splits. The new generated splits have this.groupLabel - 2
			// and this.groupLabel - 1 as labels because updateGroupVector, after assigning labels, adds 2 to the groupLabel
			// instance variable
			if(unscaledModularityDelta > 0) {
				updateGroupVector(groupModMatrix.map, localGroupVector);
			
				// perform the node moving step optimization for the group vector after updating it
				stepOptimization();
				
				int group1Size = 0;
				int group2Size = 0;
				
				for(int i = 0; i < this.groupVector.length; i++) {
					if(this.groupVector[i] == this.groupLabel - 2) {
						group1Size += 1;
					}
					if(this.groupVector[i] == this.groupLabel - 1) {
						group2Size += 1;
					}
				}
				
				// if either group has size 0 that means the other group has all the elements of the original group
				// so there's no point in processing it again since it will yield the same results
				if(group1Size != 0 && group2Size != 0) {
					split(this.groupLabel - 2);
					split(this.groupLabel - 1);
				}
			}
		}
	}
	
	// simple algorithm to generate an adjacency matrix from the graph instance variable
	// currently edgeCount serves no purpose; I just added it in case I wanted to add the
	// option to ignore edge weights.
	private double[][] getAdjacencyMatrix() {
		double[][] adjMatrix = new double[this.nodeCount][this.nodeCount];
		int sourceIntegerIdentifier = 0;
		int targetIntegerIdentifier = 0;
		double edgeWeight = 0;
		
		for(Edge edge: this.graph.getEdgeList()) {
			sourceIntegerIdentifier = edge.getSource().getIntegerIdentifier();
			targetIntegerIdentifier = edge.getTarget().getIntegerIdentifier();
			edgeWeight = edge.getData();
			
			adjMatrix[sourceIntegerIdentifier][targetIntegerIdentifier] = edgeWeight;
			adjMatrix[targetIntegerIdentifier][sourceIntegerIdentifier] = edgeWeight;
		}
		
		return adjMatrix;
	}
	
	// generates a symmetric modularity matrix from an adjacency matrix; implements equation 27 in 
	// "Finding community structure in networks using the eigenvectors of matrices" by Mark Newman.
	private double[][] getModularityMatrix() {
		double[][] adjMatrix = getAdjacencyMatrix();
		double[][] modMatrix = new double[this.nodeCount][this.nodeCount];
		
		for(int i = 0; i < this.nodeCount; i++) {
			for(int j = 0; j < this.nodeCount; j++) {		
				modMatrix[i][j] = adjMatrix[i][j] - (double)(this.graph.getNodeList().get(i).getWeight() * this.graph.getNodeList().get(j).getWeight()) / (2 * this.graph.getGraphWeight());
			}
		}
		
		return modMatrix;
	}
	
	// Calculates the modularity score of a modularity matrix given a group vector which shows
	// how the group is being split. For example, if groupVector = [-1, -1, 1], it means the
	// first two nodes belong to the same group, and the last node belongs in its own group.
	
	// If modMatrix is this.modularityMatrix, the result is the modularity score of the network
	// However if modMatrix is the modularity matrix of a subgroup, the result is the change in modularity,
	// also known as the modularity delta. The result is unscaled because the exact modularity score is
	// irrelevant; only the sign matters, so division by graph weight is unnecessary.
	
	// @param	modMatrix	a NxN modularity matrix used to calculate the modularity score
	// @param	groupVector	an array of length N (same length as modMatrix) which shows how the group is being split
	private double getModularityScore(double[][] modMatrix, int[] groupVector) {	
		double unscaledModularityScore = 0;
		
		for(int i = 0; i < modMatrix.length; i++) {
			for(int j = 0; j < modMatrix.length; j++) {
				int multiplier = groupVector[i] == groupVector[j] ? 1 : -1;
				unscaledModularityScore += modMatrix[i][j] * multiplier;
			}
		}
		
		return unscaledModularityScore;
	}
	
	// Given a node to swap, the old group of that node, and the new group you want to swap it to,
	// it returns the change in modularity. Runs in O(N) and I derived it myself from the complete
	// modularity delta equation, since the Louvain shortcut doesn't work and Newman didn't offer
	// a better alternative for the optional node swapping step. It essentially calculates how much
	// modularity would change if the given node was removed from the old group and added to the new one.
	// It does this by substracting the contributions of the node to the old group and adding the 
	// contributions of the node to the new group. The results are multiplied by 2 since, for example,
	// if the contribution to oldGroup was 9, you need to substract 9 to make it 0, and 9 more to
	// make the new modularity contribution -9, since when groups change, modularity contributions
	// aren't zeroed out, but rather their sign is flipped.
	private double getModularityDelta(int nodeID, int oldGroup, int newGroup) {
		double delta = 0;
		
		for(int i = 0; i < this.groupVector.length; i++) {
			if(nodeID != i) {
				if(this.groupVector[i] == oldGroup) {
					delta -= this.modularityMatrix[i][nodeID];
					delta -= this.modularityMatrix[nodeID][i];
				}
				
				if(this.groupVector[i] == newGroup) {
					delta += this.modularityMatrix[i][nodeID];
					delta += this.modularityMatrix[nodeID][i];
				}
			}
		}
		
		return delta * 2;
	}
	
	// Returns a subgroup modularity matrix for the group identified by groupLabel. For example,
	// if the parameter groupLabel = 6, it returns a modularity matrix for all nodes which belong
	// to the group with label 6.
	//
	// @param	groupLabel	the label of the group for which the modularity matrix is being created
	// @return	groupModMatrix	a mapped modularity matrix, which includes the modularity matrix along
	//							with an integer map which keeps track of which column/row in the group
	//							modularity matrix is which node in the full network.
	//
	// Implements equation 51 of "Finding community structure in networks using the eigenvectors of matrices"
	// by Mark Newman.
	private MappedModularityMatrix getGroupModMatrix(int groupLabel) {
		int groupSize = 0;
		
		for(int i = 0; i < this.groupVector.length; i++) {
			if(this.groupVector[i] == groupLabel) {
				groupSize++;
			}
		}
		
		MappedModularityMatrix groupModMatrix = new MappedModularityMatrix(groupSize);
		
		int gi = 0;
		int gj = 0;
		
		for(int i = 0; i < this.groupVector.length; i++) {
			if(this.groupVector[i] != groupLabel) continue;
			
			double groupModularitySum = 0;
			
			for(int k = 0; k < groupVector.length; k++) {
				if(this.groupVector[k] == groupLabel) {
					groupModularitySum += this.modularityMatrix[i][k];
				}
			}
			
			for(int j = 0; j < this.groupVector.length; j++) {
				if(this.groupVector[j] != groupLabel) continue;
				
				groupModMatrix.matrix[gi][gj] = this.modularityMatrix[i][j] - groupModularitySum;
				gj++;
			}
			
			groupModMatrix.map[gi] = i;
			
			gj = 0;
			gi++;
		}
		
		return groupModMatrix;
	}
	
	// Finds the most positive eigenvalue (i.e. the leading eigenvalue) in a modularity matrix and
	// its corresponding eigenvector (i.e. the leading eigenvector) using JBLAS, then assigns a group 
	// to each node depending on the sign of its corresponding element in the leading eigenvector.
	// For example, if leadingEigenvector[2] = -1 and leadingEigenvector[3] = 1, that means the node
	// which corresponds to row 2/column 2 in the modularity matrix belongs to group -1, and the
	// node which corresponds to row 3/column 3 in the modularity matrix belongs to group 1.
	private int[] getGroupVector(double[][] modMatrix) {
		int[] localGroupVector = new int[modMatrix.length];
		
        Matrix jmat = new Matrix(modMatrix);
        EigenvalueDecomposition eigenDecomposition = new EigenvalueDecomposition(jmat);
        double[] eigenvalues = eigenDecomposition.getRealEigenvalues();
        
        int maxEigenvalueIndex = 0;
        
        for(int i = 0; i < eigenvalues.length; i++) {
        	if(eigenvalues[i] > eigenvalues[maxEigenvalueIndex]) {
        		maxEigenvalueIndex = i;
        	}
        }
        
        double[] leadingEigenvector = new double[modMatrix.length];
        
        double[][] decompositionMatrix = eigenDecomposition.getV().getArray();
        
        for(int i = 0; i < leadingEigenvector.length; i++) {
        	leadingEigenvector[i] = decompositionMatrix[i][maxEigenvalueIndex];
        }
        
        for(int i = 0; i < leadingEigenvector.length; i++) {
        	if(leadingEigenvector[i] <= 0) {
        		localGroupVector[i] = -1;
        	}
        	else {
        		localGroupVector[i] = 1;
        	}
        }
		
		return localGroupVector;
	}
	
	// optimizes the network group vector, called every time a change is made to the instance group vector
	
	/* Description:
	 * 1. Go through each node, swapping its group, recalculating modularity, and swapping the group back
	 * 2. After finding the swap which yielded the highest modularity, make that single swap and mark
	 * 	  the node as swapped, since swapped nodes cannot be swapped again.
	 * 3. Repeat this process until no swaps occur.
	 * 
	 * In order to drastically speed up the algorithm, only groups which contain neighbors are checked.
	 * 
	 * */
	private void stepOptimization() {
		double totalDelta = 0;
		ArrayList<VectorInfo> iterVectors = new ArrayList<>();
		iterVectors.add(new VectorInfo(0, Arrays.copyOf(this.groupVector, this.groupVector.length)));

		boolean[] swapped = new boolean[this.groupVector.length];
		
		// get a list all current groups identified by their integer label
		ArrayList<Integer> groupIDs = new ArrayList<>();
		
		for(int i = 0; i < this.groupVector.length; i++) {
			if(!groupIDs.contains(this.groupVector[i])) {
				groupIDs.add(this.groupVector[i]);
			}
		}
		
		boolean done = false;
		
		// keep iterating until all swaps have been made
		while(!done) {
			
			ArrayList<SwapInfo> iterSwaps = new ArrayList<>();
			
			// try swapping the group of every node to another group
			for(int i = 0; i < this.groupVector.length; i++) {
				// skip i if it has already been swapped this iteration
				if(swapped[i]) {
					continue;
				}
				
				int oldGroup = this.groupVector[i];
				
				double maxDelta = Integer.MIN_VALUE;
				int maxGroup = oldGroup;
				
				for(int groupID : groupIDs) {
					// check all groups except own group and only if that group has at least 1 neighbor
					if(groupID != this.groupVector[i] && hasNeighbor(i, groupID)) {
						double currDelta = getModularityDelta(i, oldGroup, groupID);
						
						if(currDelta > maxDelta) {
							maxDelta = currDelta;
							maxGroup = groupID;
						}
					}
				}
				
				// at this point we have the best possible swap for i
				if(maxDelta != Integer.MIN_VALUE) {
					iterSwaps.add(new SwapInfo(i, maxGroup, maxDelta));
				}
			}
			
			// at this point we have an arraylist containing information about the best swaps for every node
			
			// make the best possible swap
			if(iterSwaps.size() > 0) {
				SwapInfo bestSwap = iterSwaps.get(0);
				
				for(SwapInfo swap : iterSwaps) {
					if(swap.delta > bestSwap.delta) {
						bestSwap = swap;
					}
				}
				
				this.groupVector[bestSwap.index] = bestSwap.newGroup;
				swapped[bestSwap.index] = true;
				totalDelta += bestSwap.delta;
				iterVectors.add(new VectorInfo(totalDelta, Arrays.copyOf(this.groupVector, this.groupVector.length)));
			}
			else {
				done = true;
			}
		}
		
		// at this point, we have every possible group vector state for the algorithm along with its running modularity delta, choose the best
		VectorInfo bestVector = iterVectors.get(0);
		
		for(int i = 0; i < iterVectors.size(); i++) {
			if(iterVectors.get(i).delta > bestVector.delta) {
				bestVector = iterVectors.get(i);
			}
		}
		
		this.groupVector = bestVector.vector;
	}
	
	// returns true if node with ID nodeID has neighbors in the group with group label groupToCheck
	private boolean hasNeighbor(int nodeID, int groupToCheck) {
		List<Node> neighbors = this.graph.getNodeList().get(nodeID).getNeighbors();
		
		for(Node neighbor : neighbors) {
			if(this.groupVector[neighbor.getIntegerIdentifier()] == groupToCheck) {
				return true;
			}
		}
		
		return false;
	}
	
	// Updates the instance group vector given an integer map and the group vector of a subgroup.
	private void updateGroupVector(int[] map, int[] localGroupVector) {
		for(int i = 0; i < localGroupVector.length; i++) {
			if(localGroupVector[i] == -1) {
				this.groupVector[map[i]] = this.groupLabel;
			}
			else if(localGroupVector[i] == 1) {
				this.groupVector[map[i]] = this.groupLabel + 1;
			}
		}
		
		this.groupLabel += 2;
	}
	
	// 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 this.communities;
	}
	
	public Graph getGraph(){
		return this.graph;
	}
	
	// Wrapper class which stores a trimmed modularity matrix and a map which can be indexed
	// to identify which row/column in the trimmed matrix matches with which which node index identifier.
	// For example, if map[0] = 15, that means that row 0 and column 0 in the trimmed matrix match
	// the node with nodeIntegerIdentifier = 15
	private class MappedModularityMatrix{
		int[] map;
		double[][] matrix;
		
		private MappedModularityMatrix(int nodeCount) {
			this.matrix = new double[nodeCount][nodeCount];
			this.map = new int[nodeCount];
		}
	}
	
	private class SwapInfo{
		int index;
		int newGroup;
		double delta;
		
		private SwapInfo(int index, int newGroup, double delta) {
			this.index = index;
			this.newGroup = newGroup;
			this.delta = delta;
		}
	}
	
	private class VectorInfo{
		double delta;
		int[] vector;
		
		private VectorInfo(double delta, int[] vector){
			this.delta = delta;
			this.vector = vector;
		}
	}
}
