package org.ucombinator.jaam.visualizer.layout;

import java.util.*;
import java.util.function.Function;

import javafx.scene.paint.Color;
import org.ucombinator.jaam.visualizer.graph.GraphUtils;
import org.ucombinator.jaam.visualizer.graph.Graph;
import org.ucombinator.jaam.visualizer.state.*;
import org.ucombinator.jaam.visualizer.taint.*;

public class LayerFactory
{
    // TODO: Template these functions instead of copying them.
    public static StateRootVertex getLayeredLoopGraph(Graph<StateVertex, StateEdge> graph) {
        return LayerFactory.getStronglyConnectedComponentsLoopGraph(graph);
    }

    private static StateRootVertex getStronglyConnectedComponentsLoopGraph(Graph<StateVertex, StateEdge> graph)
    {
        List<List<Integer>> sccs = GraphUtils.StronglyConnectedComponents(graph);
        HashMap<Integer, Integer> vertexToComponentIndex = getVertexToComponentMap(sccs);
        System.out.println("Strongly connected components in loop graph: " + sccs.size());

        StateRootVertex graphRoot = new StateRootVertex();
        graphRoot.setInnerGraph(graph);

        return (StateRootVertex) GraphUtils.compressGraph(graphRoot,
                v -> Integer.toString(vertexToComponentIndex.get(v.getId())),
                new Function<List<StateVertex>, StateVertex>() {
                    @Override
                    public StateVertex apply(List<StateVertex> stateVertices) {
                        StateSccVertex sccVertex = new StateSccVertex("SCC");
                        stateVertices.forEach(v -> {
                            sccVertex.getInnerGraph().addVertex(v);
                            v.setOuterGraph(sccVertex.getInnerGraph());
                        });
                        return sccVertex;
                    }
                },
                StateEdge::new);
    }

    public static TaintRootVertex getLayeredTaintGraph(Graph<TaintVertex, TaintEdge> graph) {
        System.out.println("JUAN: Found " + graph.getVertices().size() + " vertices");
        return getSuperTaintGraph(graph);
    }

    private static TaintRootVertex getStronglyConnectedComponentsTaintGraph(Graph<TaintVertex, TaintEdge> graph)
    {
        List<List<Integer>> sccs = GraphUtils.StronglyConnectedComponents(graph);
        HashMap<Integer, Integer> vertexToComponentIndex = getVertexToComponentMap(sccs);
        System.out.println("Strongly connected components in taint graph: " + sccs.size());

        TaintRootVertex graphRoot = new TaintRootVertex();
        graphRoot.setInnerGraph(graph);

        return (TaintRootVertex) GraphUtils.compressGraph(graphRoot,
                v -> Integer.toString(vertexToComponentIndex.get(v.getId())),
                new Function<List<TaintVertex>, TaintVertex>() {
                    @Override
                    public TaintVertex apply(List<TaintVertex> stateVertices) {
                        TaintSccVertex sccVertex = new TaintSccVertex("SCC");
                        stateVertices.forEach(v -> {
                            sccVertex.getInnerGraph().addVertex(v);
                            v.setOuterGraph(sccVertex.getInnerGraph());
                        });
                        return sccVertex;
                    }
                },
                TaintEdge::new);
    }

    private static HashMap<Integer, Integer> getVertexToComponentMap(List<List<Integer>> components) {
        HashMap<Integer, Integer> vertexToComponentIndex = new HashMap<>();
        for (int componentIndex = 0; componentIndex < components.size(); componentIndex++) {
            for (int vertexId : components.get(componentIndex)) {
                vertexToComponentIndex.put(vertexId, componentIndex);
            }
        }
        return vertexToComponentIndex;
    }

    private static TaintRootVertex getSuperTaintGraph(Graph<TaintVertex, TaintEdge> graph) {

        System.out.println("JUAN: Start there are " + calcSize(graph));
        TaintRootVertex classGraphRoot = getClassGroupingGraph(graph);

        System.out.println("JUAN After " + calcSize(classGraphRoot.getInnerGraph()));

        classGraphRoot.getInnerGraph().getVertices().forEach(classVertex -> {
            if (classVertex.getInnerGraph().getVertices().size() > 1) {
                TaintRootVertex methodRootVertex = getMethodGroupingGraph(classVertex.getInnerGraph());
                Graph<TaintVertex, TaintEdge> groupedMethodGraph = methodRootVertex.getInnerGraph();
                classVertex.setInnerGraph(groupedMethodGraph);
            }
        });
        return classGraphRoot;
    }

    private static int calcSize(Graph<TaintVertex, TaintEdge> graph) {
        int total = 0;

        for(TaintVertex v : graph.getVertices())
        {
            if (v.getInnerGraph().getVertices().size() > 1) total += calcSize(v.getInnerGraph());
            else total++;
        }

        return total;
    }

    public static long longHash(String string) {
        long h = 98764321261L;
        int l = string.length();
        char[] chars = string.toCharArray();

        for (int i = 0; i < l; i++) {
            h = 31*h + chars[i];
        }
        return h;
    }

    private static TaintRootVertex getClassGroupingGraph(Graph<TaintVertex, TaintEdge> graph)
    {
        TaintRootVertex graphRoot = new TaintRootVertex();
        graphRoot.setInnerGraph(graph);

        return (TaintRootVertex) GraphUtils.compressGraph(graphRoot,
                v -> {
                    String className = v.getClassName();

                    if (className == null) { className = v.toString(); }

                    System.out.println("JUAN: Called hash function on class" + className );

                    return longHash((className));
                },
                new Function<List<TaintVertex>, TaintVertex>() {
                    @Override
                    public TaintVertex apply(List<TaintVertex> stateVertices) {

                        String className = stateVertices.stream().findFirst().get().getClassName();

                        assert  className != null;

                        TaintSccVertex classVertex = new TaintSccVertex(className, LayoutAlgorithm.LAYOUT_ALGORITHM.DFS);
                        classVertex.setColor(Color.ORANGE);
                        stateVertices.forEach(v -> {
                            classVertex.getInnerGraph().addVertex(v);
                            v.setOuterGraph(classVertex.getInnerGraph());
                        });

                        System.out.println("Creating class vertex " + className);


                        return classVertex;
                    }
                },
                TaintEdge::new);
    }

    private static TaintRootVertex getMethodGroupingGraph(Graph<TaintVertex, TaintEdge> graph)
    {
        TaintRootVertex graphRoot = new TaintRootVertex();
        graphRoot.setInnerGraph(graph);

        return (TaintRootVertex) GraphUtils.compressGraph(graphRoot,
                v -> {
                    String methodName = v.getMethodName();
                    if (methodName == null) {
                        methodName = v.toString() + v.getLabel(); // Should be unique
                    }
                    return longHash(methodName);
                },
                new Function<List<TaintVertex>, TaintVertex>() {
                    @Override
                    public TaintVertex apply(List<TaintVertex> stateVertices) {

                        String methodName = stateVertices.stream().findFirst().get().getMethodName();
                        assert methodName != null;

                        TaintSccVertex methodVertex = new TaintSccVertex(methodName, LayoutAlgorithm.LAYOUT_ALGORITHM.DFS);
                        methodVertex.setColor(Color.HOTPINK);
                        stateVertices.forEach(v -> {
                            methodVertex.getInnerGraph().addVertex(v);
                            v.setOuterGraph(methodVertex.getInnerGraph());
                        });
                        return methodVertex;
                    }
                },
                TaintEdge::new);
    }

}

@@ -1937,110 +1937,45 @@
     
-System.out.println(%22JUAN: Found %22 + graph.getVertices().size() + %22 vertices%22);%0A        return getSuper
+return getStronglyConnectedComponents
 Tain
@@ -3744,13 +3744,22 @@
  get
-Super
+ClassClustered
 Tain
@@ -3817,227 +3817,68 @@
     
-System.out.println(%22JUAN: Start there are %22 + calcSize(graph));%0A        TaintRootVertex classGraphRoot = getClassGroupingGraph(graph);%0A%0A        System.out.println(%22JUAN After %22 + calcSize(classGraphRoot.getInnerGraph())
+TaintRootVertex classGraphRoot = getClassGroupingGraph(graph
 );%0A%0A

Set SCC Taint graph as default