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

import java.util.*;

import edu.claflin.finder.algo.clustering.struct.louvain_struct.LouvainAlgorithm;
import edu.claflin.finder.algo.huffman.HuffmanNode;
import edu.claflin.finder.logic.comp.HuffComparator;
import edu.claflin.finder.algo.huffman.HuffmanCode;
import edu.claflin.finder.algo.clustering.struct.infomap_struct.MapEquation;
import edu.claflin.finder.struct.PrioritySet;
import edu.claflin.finder.logic.Edge;
import edu.claflin.finder.logic.Graph;
import edu.claflin.finder.logic.Node;


public class InfomapAlgorithm {

    private Graph graph;
    private int nodeCount;
    private final int traverseGraph;
    private Node node = null;
    List<Node> neighbors;
    List<Node> nodeList;
    List<List<Node>> communities = new ArrayList<>();
    PrioritySet<HuffmanNode> ps = new PrioritySet<>(new HuffComparator(),true);
    HuffmanCode tree = new HuffmanCode();

    public InfomapAlgorithm(Graph g){
        graph = g;
        Random rand = new Random();
        int currentNode = 0;
        nodeCount = g.getNodeCount();
        nodeList = g.getNodeList();
        traverseGraph = 5 * nodeCount;
        ArrayList<String> nodePathString = new ArrayList(traverseGraph);
        ArrayList<Integer> nodePathInt = new ArrayList(traverseGraph);
        ArrayList<Integer> nodePathFreq = new ArrayList(traverseGraph);
        ArrayList<Node> nodePath = new ArrayList(traverseGraph);

        /*
        This is a random walker that traverses the graph 5 * the amount of nodes in a given graph (could be subject to change)
        The idea is that the random walker will begin at a certain node in a graph, then move randomly across the graph for a given amount of steps/walks
        We use this so that we can begin checking for communities that will later be calculated by the map equation, as well as incrementing a given nodes frequency to be used for a huffman node
        For more information on huffman nodes, please look at the huffman classes created
         */
        for(int i = 0; i < traverseGraph; i++){
            neighbors = graph.getNodeList().get(currentNode).getNeighbors();
            node = neighbors.get(rand.nextInt(neighbors.size()));
            node.incrementFreq();
            currentNode = node.getIntegerIdentifier();
            nodePath.add(node);
            nodePathString.add(node.getIdentifier());
            nodePathInt.add(currentNode);
            nodePathFreq.add(node.getFreq());
        }

        //HuffmanEncoder(nodeList, ps);

        //System.out.println(nodePathFreq);
        //System.out.println(nodePathInt);
        //System.out.println(nodePathString);
        System.out.println(nodePath);

        AvgBitLength(nodePath);
    }

    /*
    This method should theoretically serve to calculate the map equation (whether that be by using other classes or methods),
    and returning the number of optimal partitions a graph should have in order to minimize the cost function
    In other words, using the map equation, we should get an average bit length that is the shortest bit length when partitioning the
    graph into a certain amount of modules (i.e. 4 modules), and returning the number of modules that were created that
    resulted in the shortest bit length
     */

    /*
    quick note: when talking about per step, it is talking about how many steps the random walker took (traverseGraph = 5 * node count = 410)
    knowing this, we can divide the amount of times the index code book was used by the amount number of steps the random walker took
     */
    public double AvgBitLength(ArrayList<Node> nodePath){
        double avgBitLength = 0;
        int currentCommunity = 0;
        nodeCount = graph.getNodeCount();
        nodeList = graph.getNodeList();

        LouvainAlgorithm louvain = new LouvainAlgorithm(graph);
        PrioritySet<HuffmanNode> enter = new PrioritySet<>(new HuffComparator(),true);
        PrioritySet<HuffmanNode> exit = new PrioritySet<>(new HuffComparator(),true);
        List<PrioritySet<HuffmanNode>> huff = new ArrayList<>();
        List<Node> enterFreq = new ArrayList<>(communities.size());
        List<Node> exitFreq = new ArrayList<>(communities.size());

        ArrayList<ArrayList<Double>> wi = new ArrayList<>();
        ArrayList<Double> w = new ArrayList<>(); //arraylist holding all wa (relative weight of node at a)
        ArrayList<Double> wy = new ArrayList<>(); //arraylist holding all wiy (relative weight of all links exiting module i)
        double wiy = 0; //sum up all edge weights exiting module i, then store into wy arraylist
        double wa = 0;

        //System.out.println(louvain.getCommunities());

        System.out.println("=======================================");
        for(Integer key : louvain.getCommunities().keySet()) {
            List<Node> community = new ArrayList<>();
            ArrayList<Double> p = new ArrayList<>();
            PrioritySet<HuffmanNode> huffmanNodes = new PrioritySet<>(new HuffComparator(),true);
            huff.add(huffmanNodes);
            community.addAll(louvain.getCommunities().get(key));
            communities.add(community);
            exitFreq.add(new Node(""));
            enterFreq.add(new Node(""));
            HuffmanEncoder(community, huffmanNodes);
            System.out.println("=======================================");

            for(int i = 0; i < community.size(); i++) {
                if (louvain.getCommunities().get(key).get(i).getIntegerIdentifier() == nodePath.get(0).getIntegerIdentifier()) {
                    currentCommunity = key;
                }
            }

            for (Node node : community){
                for (int i = 0; i < node.getEdges().size(); i++){
                    wa += node.getEdges().get(i).getData();
                    if(!community.contains(node.getEdges().get(i).getTarget())){
                        wiy += node.getEdges().get(i).getData();
                    }
                }
                wa = node.getWeight() / (2 * graph.getGraphWeight());
                w.add(wa);
            }

            wiy = wiy/(2 * graph.getGraphWeight());
            wy.add(wiy);
            wi.add(w);
        }

        for(int i = 0; i < traverseGraph; i++){
            if(!communities.get(currentCommunity-1).contains(nodePath.get(i))){
                for(int j = 0; j < communities.size(); j++){
                    if(communities.get(j).contains(nodePath.get(i))){
                        exitFreq.get(currentCommunity-1).incrementFreq();
                        currentCommunity = j+1;
                        enterFreq.get(currentCommunity-1).incrementFreq();
                    }
                }
            }
        }

        avgBitLength = MapEquationUndirected(wy, w, wi, nodeCount, communities.size());
        System.out.println("The average bit length of this partitioning is: " + avgBitLength);

        return avgBitLength;
    }

    /*
    Responsible for encoding all the nodes in the graph into unique huffman codes (look at huffman classes for more info)
    Takes in a parameter of type List<Node> in order to better modify how many iterations the huffman encoder should run,
    and takes in a parameter of type PrioritySet<HuffmanNode> in order to generate the huffman nodes for the values in a
    specific priority set (In this case, we need a few for each different module, exit probability and entrance probability
     */
    public void HuffmanEncoder(List<Node> nodeList, PrioritySet<HuffmanNode> ps){
        for(Node n : nodeList) {
            HuffmanNode hn = new HuffmanNode();
            hn.setName(n.getIdentifier());
            hn.setFreq(n.getFreq());
            hn.setLeft(null);
            hn.setRight(null);
            ps.add(hn);
        }
        System.out.println("Size of priority set: " + ps.size());
        HuffmanNode root = tree.generateTree(ps);
        tree.showTree(root, "");
    }

    /*
     * Method that will use the map equation helper class to calculate the map equation.
     * @param q an array list of the exit probabilities.
     * @param p an array list of the probabilities of a node to be visited within a module.
     * @param nodeNum number of nodes.
     * @param modNum number of modules.
     */
    public double mapEquation(ArrayList<Double> q, ArrayList<Double> p, ArrayList<ArrayList<Double>> pProb, int nodeNum, int modNum) {
        MapEquation map = new MapEquation();
        double part1 = map.sigmaNotationQi(q, modNum);
        double part2 = map.sigmaNotationQiLogQi(q, modNum);
        double part3 = map.sigmaNotationPaLogPa(p, nodeNum);
        double part4 = map.sigmaNotationPa(p, nodeNum);
        double part5 = map.longSigmaNotation(q, modNum, communities, pProb);
        double result = part1 * Math.log(part1) - (2 * part2) - part3 + part5;
        return result;
    }

    /*
     * Method that will use the map equation helper class to calculate the map equation.
     * @param wy an array list containing each wiy (relative weight of links exiting module i)
     * @param w an array list of wa
     * @param wi an array list that holds another array list containing each wa (relative weight of node a) within a given module
     * @param nodeNum number of nodes.
     * @param modNum number of modules.
     */
    public double MapEquationUndirected(ArrayList<Double> wy, ArrayList<Double> w,
                                        ArrayList<ArrayList<Double>> wi, int nodeNum,
                                        int modNum){
        MapEquation map = new MapEquation();
        double part1 = map.SigmaWyLogSigmaWy(modNum, wy);
        double part2 = map.SigmaNotationWiyLogWiy(modNum, wy);
        double part3 = map.SigmaWaLogWa(nodeNum, w);
        double part4 = map.LongSigmaWiyAndWi(modNum, wy, communities, wi);
        return part1 - (2 * part2) - part3 + part4;
    }
}