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

import edu.claflin.finder.algo.clustering.struct.louvain_struct.LouvainAlgorithm;
import edu.claflin.finder.io.graph.SimpleGraphIO;
import edu.claflin.finder.logic.Edge;
import edu.claflin.finder.logic.Graph;
import edu.claflin.finder.logic.Node;

import java.io.File;
import java.util.*;

public class RelativeUndirectedBitLength {

    private Graph g;
    List<Edge> edges;
    List<Node> nodes;
    Map<Node, Integer> nodeToModule;
    private double totalWeight;
    private List<Double> relativeWeights;
    private Map<Integer, Double> moduleWeights;
    private Map<Integer, Double> moduleExitWeights;
    private double totalExitWeight;

    public static RelativeUndirectedBitLength fromNodeToModule(Graph graph, Map<Node, Integer> nodeToModule) {
        return new RelativeUndirectedBitLength(graph, nodeToModule);
    }

/*    public static RelativeUndirectedBitLength fromModules(Graph graph, List<Set<Node>> sourceModules) {
        Map<Integer, List<Node>> modules = new HashMap<>();
        for (int i = 0; i < sourceModules.size(); i++) {
            Set<Node> module = sourceModules.get(i);
            if (module.size() > 0) {
                List<Node> nodes = new ArrayList<>(module.size());
                nodes.addAll(module);
                modules.put(i, nodes);
            }
        }
        return new RelativeUndirectedBitLength(graph, modules);
    }*/

    public RelativeUndirectedBitLength(Graph graph, Map<Node, Integer> nodeToModule) {
        this.g = graph;
        this.nodeToModule = nodeToModule;
        this.edges = g.getEdgeList();
        this.nodes = g.getNodeList();
        calculateTotalWeight();
        calculateRelativeWeights();
    }

    public void setNodeToModule(Map<Node, Integer> nodeToModule) {
        this.nodeToModule = nodeToModule;
    }

    public double calculateNewModule(Map<Node, Integer> ntm) {
        setNodeToModule(ntm);
        calculateModuleWeights();
        calculateModuleExitWeights();
        return getBitLength();
    }

    public void calculateRelativeWeights(){
        relativeWeights= new ArrayList<>(nodes.size());
        for (Node n : nodes) {
            relativeWeights.add(n.getWeight()/ totalWeight);
        }
    }
    public void calculateTotalWeight() {
        totalWeight = 0;
        for (Node n : nodes) {
            totalWeight += n.getWeight();
        }
    }

    public void calculateModuleWeights(){
        moduleWeights = new HashMap<>();
        for (Map.Entry<Node, Integer> entry : nodeToModule.entrySet()) {
            int moduleId = entry.getValue();
            Node n = entry.getKey();

            double currentModuleWeight = moduleWeights.getOrDefault(moduleId, 0.0);
            moduleWeights.put(moduleId, currentModuleWeight + relativeWeights.get(n.getIntegerIdentifier()));
        }
    }

    public double getBitLength() {
        return totalExitWeight * log2(totalExitWeight)
                - 2 * moduleExitWeights.values().stream().reduce(0.0, (total, moduleExitWeight) -> total + moduleExitWeight * log2(moduleExitWeight))
                - relativeWeights.stream().reduce(0.0, (total, nodeWeight) -> total + nodeWeight * log2(nodeWeight))
                + sum(moduleWeights, moduleExitWeights);
    }

    public void calculateModuleExitWeights() {
        moduleExitWeights = new HashMap<>();
        totalExitWeight = 0;
        for (Edge e : edges) {
            int sourceModule = nodeToModule.get(e.getSource());
            int targetModule = nodeToModule.get(e.getTarget());
            if (sourceModule != targetModule) {
                double edgeWeight = e.getData();
                double relativeWeight = edgeWeight / totalWeight;
                double sourceWeight = moduleExitWeights.getOrDefault(sourceModule, 0.0) + relativeWeight;
                double targetWeight = moduleExitWeights.getOrDefault(targetModule, 0.0) + relativeWeight;

                moduleExitWeights.put(sourceModule, sourceWeight);
                moduleExitWeights.put(targetModule, targetWeight);

                totalExitWeight += 2 * relativeWeight;
            }
        }
    }

    public void calculateTotalExitWeight() {
        totalExitWeight = 0;
        for (Map.Entry<Integer, Double> entry : moduleExitWeights.entrySet()) {
            double weight = entry.getValue();
            totalExitWeight += weight;
        }
    }

    private static Map<Node, Integer> generateNodeToModule(Map<Integer, List<Node>> modules) {
        Map<Node, Integer> nodeToModule = new HashMap<>();
        for (Map.Entry<Integer, List<Node>> entry : modules.entrySet()) {
            int module = entry.getKey();
            List<Node> moduleNodes = entry.getValue();
            for (Node n : moduleNodes) {
                nodeToModule.put(n, module);
            }
        }
        return nodeToModule;
    }

    private static double log2(double x) {
        return Math.log(x) / Math.log(2);
    }

    private static double sum(Map<Integer, Double> moduleWeights, Map<Integer, Double> moduleExitWeights) {
        double total = 0;
        for (int moduleId : moduleWeights.keySet()) {
            double weight = moduleWeights.getOrDefault(moduleId, 0.0);
            double exitWeight = moduleExitWeights.getOrDefault(moduleId, 0.0);
            total += (exitWeight + weight) * log2(exitWeight + weight);
        }
        return total;
    }

    public static void main(String[] args) {
        /*
        String fname = "two_triangles.txt"; // two_triangles.txt nine_triangles.txt
        SimpleGraphIO sgio = new SimpleGraphIO();

        Graph sourceGraph = sgio.parseGraph(new File("./Test Graphs/" + fname), true);
        assert sourceGraph != null;

        Graph g;
        Map<Integer, List<Node>> modules;
        RelativeUndirectedBitLength bitLengthCalculator;

        // BAD
        g = sourceGraph.copy();
        modules = Map.of(
                1, List.of(g.getNode("5"), g.getNode("6")),
                2, List.of(g.getNode("4"), g.getNode("1")),
                3, List.of(g.getNode("2"), g.getNode("3"))
        );
        bitLengthCalculator = new RelativeUndirectedBitLength(g, modules);
        System.out.println("BAD: " + bitLengthCalculator.calculateAvgBitLength());

        //OK
        g = sourceGraph.copy();
        modules = Map.of(
                1, List.of(g.getNode("5")),
                2, List.of(g.getNode("4"), g.getNode("6")),
                3, List.of(g.getNode("1"), g.getNode("2"), g.getNode("3"))
        );
        bitLengthCalculator = new RelativeUndirectedBitLength(g, modules);
        System.out.println("OK: " + bitLengthCalculator.calculateAvgBitLength());


        //GOOD
        g = sourceGraph.copy();
        LouvainAlgorithm louvain = new LouvainAlgorithm(g);
        modules = louvain.getCommunities();
        bitLengthCalculator = new RelativeUndirectedBitLength(g, modules);
        System.out.println("GOOD: " + bitLengthCalculator.calculateAvgBitLength());
        */
    }
}
