package edg.slicing;

import edg.config.Config;
import edg.constraint.Constraints;
import edg.constraint.EdgeConstraint;
import edg.graph.EDG;
import edg.graph.Edge;
import edg.graph.LAST;
import edg.graph.Node;

import java.util.*;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * Constrained slicing algorithm based on a tabular approach based on
 * <a href="https://doi.org/10.1145/199448.199462">Reps et al.</a>'s approach.
 * <br>
 * For each step, it stores the context (the formal-out node through which a
 * node was entered) or null if no context has been found already.
 * <br>
 * It needs no precomputation of summary edges and achieves the same precision
 * with similar time requirements as a recursive computation of summary edges.
 * <br>
 * This version is not suitable for recursive programs unless
 * {@link Config#MAX_STACK_SIZE} is set to a low value (&lt;10).
 * @see ConstrainedSubsumedTabularAlgorithm A more efficient implementation.
 */
public class ConstrainedTabularAlgorithm implements SlicingAlgorithm {

    /** Summary edges that have already been found, will be immediately applied
     *  when any of this map's keys are reached in the algorithm. */
    protected final Map<State, Set<State>> summaryEdges = new HashMap<>();
    protected final EDG edg;

    /** Works that haven't been processed yet, always a subset of {@link #pathEdge}. */
    protected PriorityQueue<Work> workList;
    /** Works that have been reached throughout the traversal. */
    protected Set<Work> pathEdge;

    public ConstrainedTabularAlgorithm(EDG edg) {
        this.edg = edg;
    }

    @Override
    public Set<Node> slice(Node node) {
        workList = new PriorityQueue<>();
        pathEdge = new HashSet<>();
        propagate(new Work(new State(node, new Constraints())));
        workTheList();
        return pathEdge.stream().map(w -> w.current.node).collect(Collectors.toSet());
    }

    /**
     * If you run this algorithm multiple times, some summaries will be precomputed.
     * You can create new objects of this class or run this method.
     */
    public void clearState() {
        summaryEdges.clear();
    }

    /**
     * Add a new Work to be processed, only if it hasn't been visited before.
     */
    protected void propagate(Work work) {
        if (pathEdge.add(work))
            workList.add(work);
    }

    protected void workTheList() {
        while (!workList.isEmpty()) {
            Work w = workList.remove();
            if (isEnter(w.current.node)) {
                if (w.context == null) {
                    traverseEdges(w, e -> true);
                } else {
                    traverseEdges(w, this::edgeIsIntraprocedural);
                }
            } else if (isFormalIn(w.current.node)) {
                if (w.context == null) {
                    traverseEdges(w, e -> true);
                } else {
                    traverseEdges(w, this::edgeIsIntraprocedural);
                    createSummaryEdges(w.current, w.context);
                }
            } else if (isActualOut(w.current.node)) {
                traverseEdges(w, this::edgeIsIntraprocedural);
                traverseSummaries(w);
                contextSwitch(w.current);
            } else {
                traverseEdges(w, this::edgeIsIntraprocedural);
            }
        }
    }

    /**
     * Follow edges with the simple slicing algorithm rules.
     * Excludes control-flow, non-traversable structural edges, and a customizable filter.
     * @param work   The current node and its context.
     * @param filter A filter for edges, only those that return true will be traversed.
     */
    protected void traverseEdges(Work work, Predicate<Edge> filter) {
        // Processing NodeWork
        Node node = work.current.node;
        Edge.Type leType = work.lastEdgeType;

        for (Edge edge : edg.getEdges(node, LAST.Direction.Backwards)) {
            // Edge-skipping rules
            if (edge.isControlFlowEdge() // Control Flow is not deleted from the graph, but is unnecessary
                    || edge.getType() == Edge.Type.Summary // Always ignore in-graph summary edges (those are built with grammars and are inaccurate)
                    || !filter.test(edge) // Apply the caller-specified filter
                    || !edge.isTraversable() // Skip non-traversable edges (structural)
                    || (node.getType() == Node.Type.Generator && leType != Edge.Type.Control && edge.getType() == Edge.Type.Value)
                    || (leType == Edge.Type.Structural && edge.getType() != Edge.Type.Structural)) // After a Structural edge, only traverse structural edges
                continue;
            // Processing EdgeWork
            Node nextNode = edg.getEdgeSource(edge);
            EdgeConstraint edgeCons = edge.getConstraint();
            EdgeConstraint topConstraint = work.current.stack.isEdgeConstraintsEmpty() ? null : work.current.stack.peekEdgeConstraint();
            List<Constraints> resolvedList = edgeCons.resolve(Phase.Tabular, edg, edge, (Constraints) work.current.stack.clone(), topConstraint, 0);

            for (Constraints resolved : resolvedList)
                if (edge.getType() == Edge.Type.Structural || edge.getType() == Edge.Type.Control)
                    propagate(new Work(work, new State(nextNode, resolved), edge.getType()));
                else // no need to store the type, type only matters for structural and control.
                    propagate(new Work(work, new State(nextNode, resolved)));
        }
    }

    /** Filter for intra-procedural edges */
    protected boolean edgeIsIntraprocedural(Edge e) {
        return e.getType() != Edge.Type.Input
                && e.getType() != Edge.Type.Output
                && e.getType() != Edge.Type.Call;
    }

    /**
     * Given a formal-in, formal-out pair, creates the summary edges for every corresponding
     * actual-in, actual-out pair (must be in the same call). If a newly edge is added,
     * all actual-out nodes that have been reached immediately traverse the new summary edge.
     */
    protected void createSummaryEdges(State formalIn, State formalOut) {
        List<Node> actualOutNodes = edg.getNodes(formalOut.node, LAST.Direction.Forwards, Edge.Type.Output);
        for (Node aInNode : edg.getNodes(formalIn.node, LAST.Direction.Backwards, Edge.Type.Input)) {
            Node call = edg.getParent(aInNode).getType() == Node.Type.Arguments ?
                    edg.getAncestor(aInNode, Node.Type.Call) : edg.getNodeFromRes(aInNode);
            int toBeRemoved = -1;
            for (int i = 0; i < actualOutNodes.size(); i++) {
                Node aOutNode = actualOutNodes.get(i);
                if (call.equals(edg.getNodeFromRes(aOutNode))) {
                    toBeRemoved = i;
                    State actualIn = new State(aInNode, formalIn.stack);
                    State actualOut = new State(aOutNode, formalOut.stack);
                    if (!summaryEdges.containsKey(actualOut) || !summaryEdges.get(actualOut).contains(actualIn)) {
                        if (!summaryEdges.containsKey(actualOut))
                            summaryEdges.put(actualOut, new HashSet<>());
                        summaryEdges.get(actualOut).add(actualIn);
                        Set<Work> reached = new HashSet<>();
                        for (Work w : pathEdge)
                            if (actualOut.equals(w.current))
                                reached.add(w);
                        for (Work w : reached)
                            propagate(new Work(w, actualIn));
                    }
                }
            }
            if (toBeRemoved == -1)
                throw new IllegalStateException("Could not find actual-out node, formal-out " + formalOut.node.getId() + ", actual-in " + aInNode.getId());
            actualOutNodes.remove(toBeRemoved);
        }
    }

    protected void traverseSummaries(Work work) {
        if (summaryEdges.containsKey(work.current))
            for (State actualIn : summaryEdges.get(work.current))
                propagate(new Work(work, actualIn));
    }

    /**
     * Traverses an inter-procedural output edge and switches the context to the formal-out
     * node that constitutes the edge's source. The reached node is its own context.
     * @param actualOut The edge's target.
     */
    protected void contextSwitch(State actualOut) {
        for (Node formalOutNode : edg.getNodes(actualOut.node, LAST.Direction.Backwards, Edge.Type.Output))
            propagate(new Work(new State(formalOutNode, actualOut.stack), new State(formalOutNode, actualOut.stack)));
    }

    /** Checks whether a given node acts as the Enter node for a function. */
    protected boolean isEnter(Node node) {
        return node.getType() == Node.Type.Clause;
    }

    /** Checks whether a given node is a formal-in node in a function. */
    protected boolean isFormalIn(Node node) {
        return edg.getAncestor(node, Node.Type.Parameters) != null;
    }

    /** Checks whether a given node is an actual-out node in a call. */
    protected boolean isActualOut(Node node) {
        return !edg.getEdges(node, LAST.Direction.Backwards, Edge.Type.Output).isEmpty();
    }

    /**
     * A Work in tabular slicing is a pair of states: the context and the current state.
     * The context represents the location through which a function was entered (can
     * be null if no function has been entered through a formal-out node). The current
     * state represents the node and stack that has been reached at this point in the slice.
     * <br>
     * Given a collection of works, the set of current nodes represents the slice.
     */
    protected record Work(State context, State current, Edge.Type lastEdgeType) implements Comparable<Work> {
        public Work(State current) {
            this((State) null, current, null);
        }

        public Work(State context, State current) {
            this(context, current, null);
        }

        public Work(Work work, State current) {
            this(work.context, current, null);
        }

        public Work(Work work, State current, Edge.Type lastEdgeType) {
            this(work.context, current, lastEdgeType);
        }

        /**
         * Checks whether this object is equal to another one, except for the current stack,
         * for which we check whether this current stack is subsumed by the other's.
         */
        public boolean isSubsumedBy(Work work) {
            return Objects.equals(context, work.context) && current.node.equals(work.current.node)
                    && work.current.stack.isSuffix(current.stack);
        }

        /**
         * Compares Work objects based on the size of their current stacks.
         */
        @Override
        public int compareTo(Work o) {
            // This conditional is necessary due to subclass using null values
            // and storing Work as key in a map.
            if (current.stack == null || o.current.stack == null)
                return 0;
            return current.stack.sizeEdgeConstraints() - o.current.stack.sizeEdgeConstraints();
        }
    }

    protected record State(Node node, Constraints stack) {}
}