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

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

import java.util.*;

public class InfomapImplementation {
    private final static int MAX_ITERATIONS = 10;
    private final Graph originalGraph;
    private Map<Node, Integer> nodeToModule;
    private Map<Node, Set<Node>> moduleMembers;
    private RelativeUndirectedBitLength bitLengthCalculator;
    private double bestBitLength = Double.MAX_VALUE;

    public InfomapImplementation(Graph graph) {
        this.originalGraph = graph;
        nodeToModule = new HashMap<>();
        moduleMembers = new HashMap<>();
        bitLengthCalculator = new RelativeUndirectedBitLength(graph, nodeToModule);
    }

    public Map<Integer, List<Node>> cluster() {
        Graph graph = originalGraph;
        oneModulePerNode(graph);

        int iteration = 0;
        while (iteration < MAX_ITERATIONS) {
            iteration++;
            int numNodes = graph.getNodeCount();
            optimize(graph, nodeToModule);
            int numModules = (int) nodeToModule.values().stream().distinct().count();
            if (numModules == 1 || numModules == numNodes)
                break;

            // aggregate modules to next level graph
            Map<Integer, Node> nextLevelNodes = new HashMap<>(numModules);
            Map<Node, Set<Node>> newModuleMembers = new HashMap<>();
            int moduleId = 0;
            for (Map.Entry<Node, Integer> entry : nodeToModule.entrySet()) {
                Node n = entry.getKey();
                int module = entry.getValue();
                Node nextLevelNode = nextLevelNodes.get(module);
                if (nextLevelNode == null) {
                    nextLevelNode = new Node("inter-" + iteration + "-" + moduleId++);
                }
                Set<Node> nextLevelMembers = newModuleMembers.getOrDefault(nextLevelNode, new HashSet<>());
                if (moduleMembers.containsKey(n)) {
                    // module node
                    Set<Node> members = moduleMembers.get(n);
                    nextLevelMembers.addAll(members);

                } else {
                    // original node
                    nextLevelMembers.add(n);
                }
                newModuleMembers.put(nextLevelNode, nextLevelMembers);
                nextLevelNodes.put(module, nextLevelNode);
            }

            List<Edge> nextLevelEdges = new ArrayList<>();
            for (Edge e : graph.getEdgeList()) {
                int sourceModule = getNodeModule(e.getSource());
                int targetModule = getNodeModule(e.getTarget());
                if (sourceModule != targetModule) {
                    Node nextLevelSource = nextLevelNodes.get(sourceModule);
                    Node nextLevelTarget = nextLevelNodes.get(targetModule);
                    Edge nextLevelEdge = new Edge(nextLevelSource, nextLevelTarget, e.getData(), e.isUndirected());
                    int idx = nextLevelEdges.indexOf(nextLevelEdge);
                    if (idx > 0) {
                        Edge existingEdge = nextLevelEdges.get(idx);
                        double data = existingEdge.getData() + e.getData();
                        existingEdge.setData(data);
                    } else {
                        nextLevelEdges.add(nextLevelEdge);
                    }
                }
            }

            moduleMembers = newModuleMembers;
            graph = new Graph("iteration " + iteration, new ArrayList<>(nextLevelNodes.values()), nextLevelEdges);

            oneModulePerNode(graph);
        }
        return getCommunities();
    }

    public Map<Integer, List<Node>> getCommunities() {
        Map<Integer, List<Node>> modules = new HashMap<>();
        for (Map.Entry<Node, Integer> entry : nodeToModule.entrySet()) {
            Node n = entry.getKey();
            int module = entry.getValue();
            List<Node> moduleNodes = modules.getOrDefault(module, new ArrayList<>());
            moduleNodes.add(n);
            modules.put(module, moduleNodes);
        }
        return modules;
    }

    private void moveNode(Node node, int targetModule) {
        if (moduleMembers.containsKey(node)) {
            // module node
            for (Node m : moduleMembers.get(node)) {
                moveNode(m, targetModule);
            }
        } else {
            // original node
            nodeToModule.put(node, targetModule);
        }
    }

    /**
     * Puts each node in its own module
     */
    private void oneModulePerNode(Graph g) {
        List<Node> nodeList = g.getNodeList();
        for (int i = 0; i < nodeList.size(); i++) {
            Node node = nodeList.get(i);
            moveNode(node, i);
        }
    }

    private int getNodeModule(Node node) {
        if (moduleMembers.containsKey(node)) {
            // module node
            return getNodeModule(moduleMembers.get(node).iterator().next());
        } else {
            // original node
            return nodeToModule.get(node);
        }
    }

    private void optimize(Graph g, Map<Node, Integer> nodeToModule) {
        List<Node> nodeList = g.getNodeList();
        boolean shouldContinue = true;
        while (shouldContinue) {
            shouldContinue = false;
            for (Node node : nodeList) {
                int currentModule = getNodeModule(node);
                int bestModule = currentModule;
                double bestDelta = 0.0;
                double currentBitLength = bitLengthCalculator.calculateNewModule(nodeToModule);
                // try moving this node to every neighbor's module and identify the best option
                for (Node neighbor : node.getNeighbors()) {
                    int neighborModule = getNodeModule(neighbor);
                    if (neighborModule == currentModule) {
                        continue;
                    }

                    moveNode(node, neighborModule);
                    double newBitLength = bitLengthCalculator.calculateNewModule(nodeToModule);
                    moveNode(node, currentModule); // move node back to current module

                    if (newBitLength < bestBitLength) {
                        bestBitLength = newBitLength;
                    }
                    double delta = currentBitLength - newBitLength;

                    if (delta > bestDelta) {
                        bestDelta = delta;
                        bestModule = neighborModule;
                    }
                }
                // move node to the best module option, don't move if current module is the best
                if (bestModule != currentModule) {
                    shouldContinue = true;
                    moveNode(node, bestModule);
                }
            }
        }
    }

    private double modularityGain(Node node, int currentModule, int targetModule) {
        // Implement the calculation of modularity gain here

        double currentBitLength = bitLengthCalculator.calculateNewModule(nodeToModule);

        moveNode(node, targetModule);
        double newBitLength = bitLengthCalculator.calculateNewModule(nodeToModule);
        moveNode(node, currentModule); // move node back to current module

        if (newBitLength < bestBitLength) {
            bestBitLength = newBitLength;
        }

        return currentBitLength - newBitLength;
    }

    public double getBestBitLength() {
        return bestBitLength;
    }
}
