using System.Collections.Generic;

namespace Microsoft.Dafny {

  public class Graph<Node>
  {
    enum VisitedStatus { Unvisited, OnStack, Visited }
    class Vertex {
      public readonly Node N;
      public readonly List<Vertex!>! Successors = new List<Vertex!>();
      public Vertex SccRepresentative;  // null if not computed
      public List<Vertex!> SccMembers;  // non-null only for the representative of the SCC
      public int SccId;  // valid only for SCC representatives; indicates position of this representative vertex in the graph's topological sort
      // the following field is used during the computation of SCCs and of reachability
      public VisitedStatus Visited;
      // the following fields are used during the computation of SCCs:
      public int DfNumber;
      public int LowLink;
      // the following field is used during a Reaches computation
      public int Gen;  // generation <= Gen means this vertex has been visited in the current generation
      
      public Vertex(Node n) {
        N = n;
      }
      public void AddSuccessor(Vertex! v) {
        Successors.Add(v);        
      }
    }
    Dictionary<Node, Vertex!>! vertices = new Dictionary<Node, Vertex!>();
    bool sccComputed = false;
    List<Vertex!> topologicallySortedRepresentatives;  // computed by the SCC computation
    invariant sccComputed ==> topologicallySortedRepresentatives != null;
    public int SccCount {
      get {
        ComputeSCCs();
        assert topologicallySortedRepresentatives != null;  // follows from postcondition of ComputeSCCs and the object invariant
        return topologicallySortedRepresentatives.Count;
      }
    }
    int generation = 0;
    
    public Graph()
    {
    }
    
    /// <summary>
    /// Idempotently adds a vertex 'n' to the graph.
    /// </summary>
    public void AddVertex(Node n) {
      GetVertex(n);
    }

    /// <summary>
    /// Idempotently adds a vertex 'n' to the graph and then returns the Vertex for it.
    /// </summary>    
    Vertex! GetVertex(Node n) {
      Vertex v = FindVertex(n);
      if (v == null) {
        v = new Vertex(n);
        vertices.Add(n, v);
        if (sccComputed) {
          assert topologicallySortedRepresentatives != null;  // follows from object invariant
          v.SccRepresentative = v;
          v.SccMembers = new List<Vertex!>();
          v.SccMembers.Add(v);
          v.SccId = topologicallySortedRepresentatives.Count;
          topologicallySortedRepresentatives.Add(v);
        }
      }
      return v;
    }
    
    /// <summary>
    /// Returns the vertex for 'n' if 'n' is in the graph.  Otherwise, returns null.
    /// </summary>
    Vertex FindVertex(Node n) {
      Vertex v;
      if (vertices.TryGetValue(n, out v)) {
        assert v != null;  // follows from postcondition of TryGetValue (since 'vertices' maps to the type Vertex!)
        return v;
      } else {
        return null;
      }
    }

    /// <summary>
    /// Idempotently adds verices 'from' and 'to' the graph, and then
    /// adds an edge from 'from' to 'to'.
    /// </summary>
    public void AddEdge(Node from, Node to) {
      Vertex v0 = GetVertex(from);
      Vertex v1 = GetVertex(to);
      v0.AddSuccessor(v1);
      sccComputed = false;  // the addition of an edge may invalidate any previous computation of the graph's SCCs
    }
    
    /// <summary>
    /// Idempotently adds 'n' as a vertex and then returns a Node that is the representative element of the
    /// strongly connected component containing 'n'.
    /// </summary>
    public Node GetSCCRepresentative(Node n) {
      return GetSCCRepr(n).N;
    }
    
    /// <summary>
    /// Idempotently adds 'n' as a vertex.  Then, returns the number of SCCs before the SCC of 'n' in the
    /// topologically sorting of SCCs.
    /// </summary>
    public int GetSCCRepresentativeId(Node n) {
      return GetSCCRepr(n).SccId;
    }
    
    Vertex! GetSCCRepr(Node n) {
      Vertex v = GetVertex(n);
      ComputeSCCs();
      assert v.SccRepresentative != null;  // follows from what ComputeSCCs does
      return v.SccRepresentative;
    }
    
    /// <summary>
    /// Returns a list of the topologically sorted SCCs, each represented in the list by its representative node.
    /// </summary>
    public List<Node>! TopologicallySortedComponents() {
      ComputeSCCs();
      assert topologicallySortedRepresentatives != null;  // follows from object invariant
      List<Node> nn = new List<Node>();
      foreach (Vertex v in topologicallySortedRepresentatives) {
        nn.Add(v.N);
      }
      return nn;
    }

    /// <summary>
    /// Idempotently adds 'n' as a vertex and then returns the set of Node's in the strongly connected component
    /// that contains 'n'.
    /// </summary>
    public List<Node>! GetSCC(Node n) {
      Vertex v = GetVertex(n);
      ComputeSCCs();
      Vertex repr = v.SccRepresentative;
      assert repr != null && repr.SccMembers != null;  // follows from postcondition of ComputeSCCs
      List<Node> nn = new List<Node>();
      foreach (Vertex w in repr.SccMembers) {
        nn.Add(w.N);
      }
      return nn;
    }
    
    /// <summary>
    /// Idempotently adds 'n' as a vertex and then returns the size of the set of Node's in the strongly connected component
    /// that contains 'n'.
    /// </summary>
    public int GetSCCSize(Node n)
      ensures 1 <= result;
    {
      Vertex v = GetVertex(n);
      ComputeSCCs();
      Vertex repr = v.SccRepresentative;
      assert repr != null && repr.SccMembers != null;  // follows from postcondition of ComputeSCCs
      return repr.SccMembers.Count;
    }
    
    /// <summary>
    /// This method sets the SccRepresentative fields of the graph's vertices so that two
    /// vertices have the same representative iff they are in the same strongly connected
    /// component.
    /// As a side effect, this method may change the Visited, DfNumber, and LowLink fields
    /// of the vertices.
    /// </summary>
    void ComputeSCCs()
      ensures sccComputed;
    {
      if (sccComputed) { return; }  // check if already computed

      // reset all SCC information
      topologicallySortedRepresentatives = new List<Vertex!>();
      foreach (Vertex v in vertices.Values) {
        v.Visited = VisitedStatus.Unvisited;
        v.SccMembers = null;
      }
      Stack<Vertex!> stack = new Stack<Vertex!>();
      int cnt = 0;
      foreach (Vertex v in vertices.Values) {
        if (v.Visited == VisitedStatus.Unvisited) {
          SearchC(v, stack, ref cnt);
        }
      }
      assert cnt == vertices.Count;  // sanity check that everything has been visited

      sccComputed = true;
    }
    
    /// <summary>
    /// This is the 'SearchC' procedure from the Aho, Hopcroft, and Ullman book 'The Design and Analysis of Computer Algorithms'.
    /// </summary>
    void SearchC(Vertex! v, Stack<Vertex!>! stack, ref int cnt)
      requires v.Visited == VisitedStatus.Unvisited;
      requires topologicallySortedRepresentatives != null;
      ensures v.Visited != VisitedStatus.Unvisited;
    {
      v.DfNumber = cnt;
      cnt++;
      v.LowLink = v.DfNumber;
      stack.Push(v);
      v.Visited = VisitedStatus.OnStack;
      
      foreach (Vertex w in v.Successors) {
        if (w.Visited == VisitedStatus.Unvisited) {
          SearchC(w, stack, ref cnt);
          v.LowLink = min{v.LowLink, w.LowLink};
        } else if (w.Visited == VisitedStatus.OnStack) {
          assert w.DfNumber < v.DfNumber || v.LowLink <= w.DfNumber;  // the book also has the guard 'w.DfNumber < v.DfNumber', but that seems unnecessary to me, so this assert is checking my understanding
          v.LowLink = min{v.LowLink, w.DfNumber};
        }
      }
      
      if (v.LowLink == v.DfNumber) {
        // The SCC containing 'v' has now been computed.
        v.SccId = topologicallySortedRepresentatives.Count;
        topologicallySortedRepresentatives.Add(v);
        v.SccMembers = new List<Vertex!>();
        while (true) {
          Vertex x = stack.Pop();
          x.Visited = VisitedStatus.Visited;
          x.SccRepresentative = v;
          v.SccMembers.Add(x);
          if (x == v) { break; }
        }
      }
    }
    
    /// <summary>
    /// Returns null if the graph has no cycles.  If the graph does contain some cycle, returns the list of
    /// vertices on one such cycle.
    /// </summary>
    public List<Node> TryFindCycle() {
      // reset all visited information
      foreach (Vertex v in vertices.Values) {
        v.Visited = VisitedStatus.Unvisited;
      }
      
      foreach (Vertex v in vertices.Values) {
        assert v.Visited != VisitedStatus.OnStack;
        if (v.Visited == VisitedStatus.Unvisited) {
          List<Vertex!> cycle = CycleSearch(v);
          if (cycle != null) {
            List<Node> nodes = new List<Node>();
            foreach (Vertex v in cycle) {
              nodes.Add(v.N);
            }
            return nodes;  // a cycle is found
          }
        }
      }
      return null;  // there are no cycles
    }
    
    /// <summary>
    /// A return of null means there are no cycles involving any vertex in the subtree rooted at v.
    /// A non-null return means a cycle has been found.  Then:
    /// If v.Visited == Visited, then the entire cycle is described in the returned list.
    /// If v.Visited == OnStack, then the cycle consists of the vertices strictly deeper than
    /// w on the stack followed by the vertices (in reverse order) in the returned list, where
    /// w is the first vertex in the list returned.
    /// </summary>
    List<Vertex!> CycleSearch(Vertex! v)
      requires v.Visited == VisitedStatus.Unvisited;
      ensures v.Visited != VisitedStatus.Unvisited;
      ensures result == null ==> v.Visited == VisitedStatus.Visited;
      ensures result != null ==> result.Count != 0;
    {
      v.Visited = VisitedStatus.OnStack;
      foreach (Vertex succ in v.Successors) {
        // todo:  I would use a 'switch' statement, but there seems to be a bug in the Spec# compiler's type checking.
        if (succ.Visited == VisitedStatus.Visited) {
          // there is no cycle in the subtree rooted at succ, hence this path does not give rise to any cycles
        } else if (succ.Visited == VisitedStatus.OnStack) {
          // we found a cycle!
          List<Vertex!> cycle = new List<Vertex!>();
          cycle.Add(succ);
          if (v == succ) {
            // entire cycle has been found
            v.Visited = VisitedStatus.Visited;
          }
          return cycle;
        } else {
          assert succ.Visited == VisitedStatus.Unvisited;
          List<Vertex!> cycle = CycleSearch(succ);
          if (cycle != null) {
            if (succ.Visited == VisitedStatus.Visited) {
              // the entire cycle has been collected
              v.Visited = VisitedStatus.Visited;
              return cycle;
            } else {
              cycle.Add(succ);
              if (v == cycle[0]) {
                // the entire cycle has been collected and we are the first to find out
                v.Visited = VisitedStatus.Visited;
              }
              return cycle;
            }
          }
        }
      }
      v.Visited = VisitedStatus.Visited;  // there are no cycles from here on
      return null;
    }
    
    /// <summary>
    /// Returns whether or not 'source' reaches 'sink' in the graph.
    /// 'source' and 'sink' need not be in the graph; if neither is, the return value
    /// is source==sink.
    /// </summary>
    public bool Reaches(Node source, Node sink) {
      Vertex a = FindVertex(source);
      Vertex b = FindVertex(sink);
      if (a == null || b == null) {
        return source.Equals(sink);
      }
      generation++;
      return ReachSearch(a, b);
    }
    
    bool ReachSearch(Vertex! source, Vertex! sink) {
      if (source == sink) {
        return true;
      } else if (source.Gen == generation) {
        // already visited
        return false;
      } else {
        source.Gen = generation;
        return exists{Vertex succ in source.Successors; ReachSearch(succ, sink)};
      }
    }
  }
}