/* Copyright (C) 2013-2017 TU Dortmund
 * This file is part of AutomataLib, http://www.automatalib.net/.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package net.automatalib.util.graphs;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.Queue;
import java.util.function.Predicate;

import com.google.common.collect.AbstractIterator;
import net.automatalib.commons.util.mappings.MutableMapping;
import net.automatalib.graphs.IndefiniteGraph;

final class FindShortestPathsIterator<N, E> extends AbstractIterator<Path<N, E>> {

    private final Queue<N> bfsQueue = new ArrayDeque<>();
    private final IndefiniteGraph<N, E> graph;
    private final MutableMapping<N, Pred<N, E>> preds;
    private final Predicate<? super N> targetPred;
    private final int limit;
    private int count;

    FindShortestPathsIterator(IndefiniteGraph<N, E> graph,
                                     Collection<? extends N> start,
                                     int limit,
                                     Predicate<? super N> targetPred) {
        this.graph = graph;
        this.preds = graph.createStaticNodeMapping();
        this.limit = limit;
        this.targetPred = Objects.requireNonNull(targetPred);

        for (N startNode : start) {
            preds.put(startNode, new Pred<>(null, null));
            bfsQueue.add(startNode);
        }
    }

    @Override
    protected Path<N, E> computeNext() {
        while (!bfsQueue.isEmpty()) {
            if (count++ == limit) {
                return endOfData();
            }
            N curr = bfsQueue.poll();
            if (targetPred.test(curr)) {
                return makePath(curr);
            }

            for (E edge : graph.getOutgoingEdges(curr)) {
                N tgt = graph.getTarget(edge);
                Pred<N, E> pred = preds.get(tgt);
                if (pred == null) {
                    preds.put(tgt, new Pred<>(curr, edge));
                    bfsQueue.add(tgt);
                }
            }
        }

        return endOfData();
    }

    private Path<N, E> makePath(N target) {
        List<E> edges = new ArrayList<>();

        N currNode = target;
        Pred<N, E> pred = preds.get(currNode);

        while (pred != null && pred.edge != null) {
            edges.add(pred.edge);

            currNode = pred.node;
            pred = preds.get(currNode);
        }

        Collections.reverse(edges);

        return new Path<>(graph, currNode, edges);
    }

    private static final class Pred<N, E> {

        public final N node;
        public final E edge;

        Pred(N node, E edge) {
            this.node = node;
            this.edge = edge;
        }
    }

}

@@ -841,34 +841,8 @@
 st;%0A
-import java.util.Objects;%0A
 impo
@@ -899,16 +899,61 @@
 icate;%0A%0A
+import com.google.common.base.Preconditions;%0A
 import c
@@ -1227,29 +1227,8 @@
 ueue
- = new ArrayDeque%3C%3E()
 ;%0A  
@@ -1410,31 +1410,8 @@
 mit;
-%0A    private int count;
 %0A%0A  
@@ -1497,23 +1497,16 @@
         
-       
 Collecti
@@ -1558,23 +1558,16 @@
         
-       
 int limi
@@ -1595,31 +1595,24 @@
             
-  
-     
 Predicate%3C? 
@@ -1630,24 +1630,182 @@
 rgetPred) %7B%0A
+        Preconditions.checkArgument(limit %3E= 0, %22Limit must be non-negative%22);%0A        Preconditions.checkNotNull(targetPred, %22Predicate must be non-null%22);%0A%0A
         this
@@ -1932,41 +1932,74 @@
 d = 
-Objects.requireNonNull(targetPred
+targetPred;%0A%0A        this.bfsQueue = new ArrayDeque%3C%3E(start.size()
 );%0A%0A
@@ -2088,16 +2088,19 @@
 ll, null
+, 0
 ));%0A    
@@ -2246,94 +2246,8 @@
 ) %7B%0A
-            if (count++ == limit) %7B%0A                return endOfData();%0A            %7D%0A
     
@@ -2367,32 +2367,93 @@
 %0A            %7D%0A%0A
+            final int currentDepth = preds.get(curr).depth;%0A%0A
             for 
@@ -2560,33 +2560,39 @@
      Pred%3CN, E%3E 
-p
+targetP
 red = preds.get(
@@ -2617,17 +2617,23 @@
     if (
-p
+targetP
 red == n
@@ -2631,24 +2631,48 @@
 Pred == null
+ && currentDepth %3C limit
 ) %7B%0A        
@@ -2719,16 +2719,34 @@
 rr, edge
+, currentDepth + 1
 ));%0A    
@@ -2906,52 +2906,8 @@
 ) %7B%0A
-        List%3CE%3E edges = new ArrayList%3C%3E();%0A%0A
     
@@ -2975,24 +2975,78 @@
 currNode);%0A%0A
+        List%3CE%3E edges = new ArrayList%3C%3E(pred.depth);%0A%0A
         whil
@@ -3402,16 +3402,48 @@
  E edge;
+%0A        public final int depth;
 %0A%0A      
@@ -3463,16 +3463,27 @@
 , E edge
+, int depth
 ) %7B%0A    
@@ -3534,24 +3534,56 @@
 dge = edge;%0A
+            this.depth = depth;%0A
         %7D%0A  

fix semantics of ShortestPathsIterator