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())); // TODO: random distribution based on edge weight
            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); // TODO: why use Louvian?
        // TODO: enter and exit codes are not used
        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);

            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();
                    }
                }
                //relative weight is the total weight of the links connected to the node divided by twice the total weight of all links in the network,
                wa = node.getWeight() / (2 * graph.getGraphWeight());
                w.add(wa);
            }

            wiy = wiy/(2 * graph.getGraphWeight());
            wy.add(wiy);
            wi.add(w);
            System.out.println("wiy: " + wiy);
            System.out.println("w: " + wy);
            System.out.println("wy: " + wi);
        }

        /*
        current community is used to detect which community the random walker is currently in
        this value should be used in the for loop below so that when you traverse the nodePathInt array, and compare it
        to which community it is in, you can change the value accordingly
        */
        System.out.println("current community is: " + currentCommunity);

        for(int i = 0; i < traverseGraph; i++){

            /*
            use the nodePathString arraylist to access the path the random walker took
            we start off at our first node, then we check afterwards to see if the next node it entered is in the same community
            if it is, nothing happens, but if the node it entered is in another community, increase entrance frequency of the
            community it entered, and exit frequency of the community it exited
             */
            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();
                    }
                }
            }
        }

        /*
        TODO: change to map equation implementation from paper
        We would like to work on finding a way to get the huffman codes for the entrance and exit probabilities as seen
        in the map equation paper by Martin Rosvall, however that will be something to accomplish at a later time

        HuffmanEncoder(exitFreq, exit);
        HuffmanEncoder(enterFreq, enter);
        */

        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;
    }
}