/*
 *  Copyright (C) 2004-2019 Edward F. Valeev
 *
 *  This file is part of Libint.
 *
 *  Libint is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  Libint is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with Libint.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

#ifndef _libint2_src_bin_libint_dg_h_
#define _libint2_src_bin_libint_dg_h_

#include <iostream>
#include <string>
#include <list>
#include <map>
#include <vector>
#include <deque>
#include <algorithm>
#include <stdexcept>
#include <cassert>

#include <global_macros.h>
#include <exception.h>
#include <smart_ptr.h>
#include <key.h>
#include <dgvertex.h>

namespace libint2 {

//  class DGVertex;
  class DGArc;
  template <class T> class DGArcRel;
  template <class T> class AlgebraicOperator;
  class Strategy;
  class Tactic;
  class CodeContext;
  class MemoryManager;
  class ImplicitDimensions;
  class CodeSymbols;
  class GraphRegistry;
  class InternalGraphRegistry;

  /** DirectedGraph is an implementation of a directed graph
      composed of vertices represented by DGVertex objects. Most important operations
      will assume that this is a DAG, i.e. there are no directed cycles.

      \note The objects are allocated on free store and the graph is implemented as
      an object of type 'vertices'.
   */

  class DirectedGraph : public EnableSafePtrFromThis<DirectedGraph> {
  public:
    typedef DGVertex vertex;
    typedef DGArc arc;
    typedef SafePtr<DGVertex> ver_ptr;
    typedef SafePtr<DGArc> arc_ptr;
    typedef DGVertexKey key_type;
    typedef std::multimap<key_type,ver_ptr> VPtrAssociativeContainer;
    typedef std::list<ver_ptr> VPtrSequenceContainer;

    typedef VPtrSequenceContainer targets;
#if USE_ASSOCCONTAINER_BASED_DIRECTEDGRAPH
    typedef VPtrAssociativeContainer vertices;
#else
    typedef VPtrSequenceContainer vertices;
#endif
    typedef targets::iterator target_iter;
    typedef targets::const_iterator target_citer;
    typedef vertices::iterator ver_iter;
    typedef vertices::const_iterator ver_citer;
    //not possible: typedef vertex::Address address;
    typedef int address;
    //not possible: typedef vertex::Size size;
    typedef unsigned int size;
    typedef std::vector<address> addresses;

  private:
    /// converts what is stored in the container to a smart ptr to the vertex
    static inline const ver_ptr& vertex_ptr(const VPtrAssociativeContainer::value_type& v) {
      return v.second;
    }
    static inline const ver_ptr& vertex_ptr(const VPtrSequenceContainer::value_type& v) {
      return v;
    }
    /// converts what is stored in the container to a smart ptr to the vertex
    static inline ver_ptr& vertex_ptr(VPtrAssociativeContainer::value_type& v) {
      return v.second;
    }
    static inline ver_ptr& vertex_ptr(VPtrSequenceContainer::value_type& v) {
      return v;
    }

  public:

    /** Creates an empty DAG. Actual initialization of the graph
        must be done using append_target */
    DirectedGraph();
    ~DirectedGraph();

    /// Returns the number of vertices
    unsigned int num_vertices() const { return stack_.size(); }
#if 0
    /// Returns all vertices
    const vertices& all_vertices() const { return stack_; }
    /// Returns all targets
    const targets& all_targets() const { return targets_; }
#endif
    /// Find vertex v or it's equivalent on the graph. Return null pointer if not found.
    /// Computational cost for a graph based on a nonassociative container may be high
    const SafePtr<DGVertex>& find(const SafePtr<DGVertex>& v) const { return vertex_is_on(v); }

    /** appends v to the graph. If v's copy is already on the graph, return the pointer
	to the copy. Else return pointer to *v.
     */
    SafePtr<DGVertex> append_vertex(const SafePtr<DGVertex>& v);

    /** non-template append_target appends the vertex to the graph as a target
     */
    void append_target(const SafePtr<DGVertex>&);

    /** append_target appends I to the graph as a target vertex and applies
        RR to it. append_target can be called multiple times on the same
        graph if more than one target vertex is needed.

        I must derive from DGVertex. RR must derive from RecurrenceRelation.
        RR has a constructor which takes const I& as the only argument.
        RR must have a public member const I* child(unsigned int) .

        NOTE TO SELF : need to implement these restrictions using
        standard Bjarne Stroustrup's approach.

     */
    template <class I, class RR> void append_target(const SafePtr<I>& target) {
      target->make_a_target();
      recurse<I,RR>(target);
    }

    /** apply_to_all applies RR to all vertices already on the graph.

        RR must derive from RecurrenceRelation. RR must define TargetType
        as a typedef.
        RR must have a public member const DGVertex* child(unsigned int) .

        NOTE TO SELF : need to implement these restrictions using
        standard Bjarne Stroustrup's approach.

     */
    template <class RR> void apply_to_all() {
      typedef typename RR::TargetType TT;
      typedef vertices::const_iterator citer;
      typedef vertices::iterator iter;
      for(iter v=stack_.begin(); v!=stack_.end(); ++v) {
        ver_ptr& vptr = vertex_ptr(*v);
        if ((vptr)->num_exit_arcs() != 0)
          continue;
        SafePtr<TT> tptr = dynamic_pointer_cast<TT,DGVertex>(v);
        if (tptr == 0)
          continue;

        SafePtr<RR> rr0(new RR(tptr));
        const int num_children = rr0->num_children();

        for(int c=0; c<num_children; c++) {

          SafePtr<DGVertex> child = rr0->child(c);
          SafePtr<DGArc> arc(new DGArcRel<RR>(tptr,child,rr0));
          tptr->add_exit_arc(arc);

          recurse<RR>(child);

        }
      }
    }

    /** after all append_target's have been called, apply(strategy,tactic)
      constructs a graph. strategy specifies how to apply recurrence relations.
      The goal of strategies is to connect the target vertices to simpler, precomputable vertices.
      There usually are many ways to reduce a vertex.
      Tactic specifies which of these possibilities to choose.
     */
    void apply(const SafePtr<Strategy>& strategy,
        const SafePtr<Tactic>& tactic);

    typedef void (DGVertex::* DGVertexMethodPtr)();
    /** apply_at<method>(vertex) calls method() on vertex and all of its descendants
     */
    template <DGVertexMethodPtr method>
    void apply_at(const SafePtr<DGVertex>& vertex) const {
      ((vertex.get())->*method)();
      typedef DGVertex::ArcSetType::const_iterator aciter;
      const aciter abegin = vertex->first_exit_arc();
      const aciter aend = vertex->plast_exit_arc();
      for(aciter a=abegin; a!=aend; ++a)
        apply_at<method>((*a)->dest());
    }

    /** calls Method(v) for each v, iterating in forward direction */
    template <class Method>
    void foreach(Method& m) {
      typedef vertices::const_iterator citer;
      typedef vertices::iterator iter;
      for(iter v=stack_.begin(); v!=stack_.end(); ++v) {
        ver_ptr& vptr = vertex_ptr(*v);
        m(vptr);
      }
    }

    /** calls Method(v) for each v, iterating in forward direction */
    template <class Method>
    void foreach(Method& m) const  {
      typedef vertices::const_iterator citer;
      typedef vertices::iterator iter;
      for(citer v=stack_.begin(); v!=stack_.end(); ++v) {
        const ver_ptr& vptr = vertex_ptr(*v);
        m(vptr);
      }
    }
    /** calls Method(v) for each v, iterating in reverse direction */
    template <class Method>
    void rforeach(Method& m) {
      typedef vertices::const_reverse_iterator criter;
      typedef vertices::reverse_iterator riter;
      for(riter v=stack_.rbegin(); v!=stack_.rend(); ++v) {
        ver_ptr& vptr = vertex_ptr(*v);
        m(vptr);
      }
    }

    /** calls Method(v) for each v, iterating in reverse direction */
    template <class Method>
    void rforeach(Method& m) const {
      typedef vertices::const_reverse_iterator criter;
      typedef vertices::reverse_iterator riter;
      for(criter v=stack_.rbegin(); v!=stack_.rend(); ++v) {
        const ver_ptr& vptr = vertex_ptr(*v);
        m(vptr);
      }
    }

    /** after Strategy has been applied, simple recurrence relations need to be
        optimized away. optimize_rr_out() will replace all simple recurrence relations
        with code representing them.
     */
    void optimize_rr_out(const SafePtr<CodeContext>& context);

    /** after all apply's have been called, traverse()
        construct a heuristic order of traversal for the graph.
     */
    void traverse();

    /** update func_names_
     */
    void update_func_names();

    /// Prints out call sequence
    void debug_print_traversal(std::ostream& os) const;

    /**
    Prints out the graph in format understood by program "dot"
    of package "graphviz". If symbols is true then label vertices
    using their symbols rather than (descriptive) labels.
     */
    void print_to_dot(bool symbols, std::ostream& os = std::cout) const;

    /**
       Generates code for the current computation using context.
       dims specifies the implicit dimensions,
       args specifies the code symbols for the arguments to the function,
       label specifies the tag for the computation,
       decl specifies the stream to receive declaration code,
       code specifies the stream to receive the definition code
     */
    void generate_code(const SafePtr<CodeContext>& context, const SafePtr<MemoryManager>& memman,
        const SafePtr<ImplicitDimensions>& dims, const SafePtr<CodeSymbols>& args,
        const std::string& label,
        std::ostream& decl, std::ostream& code);

    /** Resets the graph and all vertices. The stack of unresolved recurrence
        relations is preserved.
     */
    void reset();

    /** num_children_on(rr) returns the number of children of rr which
        are already on this graph.
     */
    template <class RR>
    unsigned int
    num_children_on(const SafePtr<RR>& rr) const {
      unsigned int nchildren = rr->num_children();
      unsigned int nchildren_on_stack = 0;
      for(unsigned int c=0; c<nchildren; c++) {
        if (!vertex_is_on(rr->rr_child(c)))
          continue;
        else
          nchildren_on_stack++;
      }

      return nchildren_on_stack;
    }

    /// Returns the registry
    SafePtr<GraphRegistry>& registry() { return registry_; }
    const SafePtr<GraphRegistry>& registry() const { return registry_; }

    /// return the graph label
    const std::string& label() const { return label_; }
    /// sets label to \c new_label
    void set_label(const std::string& new_label) { label_ = new_label; }

    /// return true if there are vertices with 0 children but not pre-computed
    bool missing_prerequisites() const;

  private:

    /// contains vertices
    vertices stack_;
    /// refers to targets, cannot be an associative container -- order of iteration over targets is important
    targets targets_;
    /// addresses of blocks which accumulate targets
    addresses target_accums_;

    // graph label, used for annotating internal work, e.g. graphviz plots
    std::string label_;

    typedef std::map<std::string,bool> FuncNameContainer;
    /** Maintains the list of names of functions calls to which have been generated so far.
        It is used to generate include statements.
     */
    FuncNameContainer func_names_;

#if !USE_ASSOCCONTAINER_BASED_DIRECTEDGRAPH
    static const unsigned int default_size_ = 100;
#endif

    // maintains data about the graph which does not belong IN the graph
    SafePtr<GraphRegistry> registry_;
    // maintains private data about the graph which does not belong IN the graph
    SafePtr<InternalGraphRegistry> iregistry_;

    /// Access the internal registry
    SafePtr<InternalGraphRegistry>& iregistry() { return iregistry_; }
    const SafePtr<InternalGraphRegistry>& iregistry() const { return iregistry_; }

    /** adds a vertex to the graph. If the vertex already found on the graph
        then the vertex is not added and the function returns false */
    SafePtr<DGVertex> add_vertex(const SafePtr<DGVertex>&);
    /** same as add_vertex(), only assumes that there's no equivalent vertex on the graph (see vertex_is_on) */
    void add_new_vertex(const SafePtr<DGVertex>&);
    /// returns true if vertex if already on graph
    const SafePtr<DGVertex>& vertex_is_on(const SafePtr<DGVertex>& vertex) const;
    /// removes vertex from the graph. may throw CannotPerformOperation
    void del_vertex(vertices::iterator&);
    /** This function is used to implement (recursive) append_target().
        vertex is appended to the graph and then RR is applied to is.
     */
    template <class I, class RR> void recurse(const SafePtr<I>& vertex)  {
      SafePtr<DGVertex> dgvertex = add_vertex(vertex);
      if (dgvertex != vertex)
        return;

      SafePtr<RR> rr0(new RR(vertex));
      const int num_children = rr0->num_children();

      for(int c=0; c<num_children; c++) {

        SafePtr<DGVertex> child = rr0->child(c);
        SafePtr<DGArc> arc(new DGArcRel<RR>(vertex,child,rr0));
        vertex->add_exit_arc(arc);

        SafePtr<I> child_cast = dynamic_pointer_cast<I,DGVertex>(child);
        if (child_cast == 0)
          throw std::runtime_error("DirectedGraph::recurse(const SafePtr<I>& vertex) -- dynamic cast failed, most probably this is a logic error!");
        recurse<I,RR>(child_cast);

      }
    }

    /** This function is used to implement (recursive) apply_to_all().
        RR is applied to vertex and all its children.
     */
    template <class RR> void recurse(const SafePtr<DGVertex>& vertex)  {
      SafePtr<DGVertex> dgvertex = add_vertex(vertex);
      if (dgvertex != vertex)
        return;

      typedef typename RR::TargetType TT;
      SafePtr<TT> tptr = dynamic_pointer_cast<TT,DGVertex>(vertex);
      if (tptr == 0)
        return;

      SafePtr<RR> rr0(new RR(tptr));
      const int num_children = rr0->num_children();

      for(int c=0; c<num_children; c++) {

        SafePtr<DGVertex> child = rr0->child(c);
        SafePtr<DGArc> arc(new DGArcRel<RR>(vertex,child,rr0));
        vertex->add_exit_arc(arc);

        recurse<RR>(child);
      }
    }

    /** This function is used to implement (recursive) apply().
      strategy and tactic are applied to vertex and all its children.
     */
    void apply_to(const SafePtr<DGVertex>& vertex,
        const SafePtr<Strategy>& strategy,
        const SafePtr<Tactic>& tactic);
    /// This function insert expr of type AlgebraicOperator<DGVertex> into the graph
    SafePtr<DGVertex> insert_expr_at(const SafePtr<DGVertex>& where, const SafePtr< AlgebraicOperator<DGVertex> >& expr);
    /// This function replaces RecurrenceRelations with concrete arithemtical expressions
    void replace_rr_with_expr();
    /// This function gets rid of trivial math such as multiplication/division by 1.0, etc.
    void remove_trivial_arithmetics();
    /** This function gets rid of nodes which are connected
    to their equivalents (such as (ss|ss) shell quartet can only be connected to (ss|ss) integral)
     */
    void handle_trivial_nodes(const SafePtr<CodeContext>& context);
    /// This functions removes vertices not connected to other vertices
    void remove_disconnected_vertices();
    /** Finds (binary) subtrees. The subtrees correspond to a single-line code (no intermediates
        are used in other expressions)
     */
    void find_subtrees();
    /** Finds (binary) subtrees starting (recursively) at v.
     */
    void find_subtrees_from(const SafePtr<DGVertex>& v);
    /** If v1 and v2 are connected by DGArcDirect and all entry arcs to v1 are of the DGArcDirect type as well,
        this function will reattach all arcs extering v1 to v2 and remove v1 from the graph altogether.
	May throw CannotPerformOperation.
     */
    bool remove_vertex_at(const SafePtr<DGVertex>& v1, const SafePtr<DGVertex>& v2);

    // Which vertex is the first to compute
    SafePtr<DGVertex> first_to_compute_;
    // prepare_to_traverse must be called before actual traversal
    void prepare_to_traverse();
    // traverse_from(arc) build recurively the traversal order
    void traverse_from(const SafePtr<DGArc>&);
    // schedule_computation(vertex) puts vertex first in the computation order
    void schedule_computation(const SafePtr<DGVertex>&);

    // Compute addresses on stack assuming that quantities larger than min_size_to_alloc to be allocated on stack
    void allocate_mem(const SafePtr<MemoryManager>& memman,
        const SafePtr<ImplicitDimensions>& dims,
        unsigned int min_size_to_alloc = 1);
    // Assign symbols to the vertices
    void assign_symbols(const SafePtr<CodeContext>& context, const SafePtr<ImplicitDimensions>& dims);
    // If v is an AlgebraicOperator, assign (recursively) symbol to the operator. All other must have been already assigned
    void assign_oper_symbol(const SafePtr<CodeContext>& context, SafePtr<DGVertex>& v);
    // Print the code using symbols generated with assign_symbols()
    void print_def(const SafePtr<CodeContext>& context, std::ostream& os,
        const SafePtr<ImplicitDimensions>& dims,
        const SafePtr<CodeSymbols>& args);

    /** Returns true if cannot enclose the code in a vector loop
        Possible reason: the traversal path contains a RecurrenceRelation that generates a function call
        (most do except IntegralSet_to_Integrals; \sa RecurrentRelation::is_simple() )
        */
    bool cannot_enclose_in_outer_vloop() const;

  };

  //
  // Nonmember utilities
  //

  /// converts what is stored in the container to a smart ptr to the vertex
  inline const DirectedGraph::ver_ptr& vertex_ptr(const DirectedGraph::VPtrAssociativeContainer::value_type& v) {
    return v.second;
  }
  inline const DirectedGraph::ver_ptr& vertex_ptr(const DirectedGraph::VPtrSequenceContainer::value_type& v) {
    return v;
  }
  /// converts what is stored in the container to a smart ptr to the vertex
  inline DirectedGraph::ver_ptr& vertex_ptr(DirectedGraph::VPtrAssociativeContainer::value_type& v) {
    return v.second;
  }
  inline DirectedGraph::ver_ptr& vertex_ptr(DirectedGraph::VPtrSequenceContainer::value_type& v) {
    return v;
  }

#if USE_ASSOCCONTAINER_BASED_DIRECTEDGRAPH
  inline DirectedGraph::key_type key(const DGVertex& v);
#endif

  //
  // Nonmember predicates
  //

  /// return true if there are non-unrolled targets
  bool nonunrolled_targets(const DirectedGraph::targets& targets);

  /// extracts external symbols and RRs from the graph
  void extract_symbols(const SafePtr<DirectedGraph>& dg);

  // use these functors with DirectedGraph::foreach
  struct PrerequisitesExtractor {
    std::deque< SafePtr<DGVertex> > vertices;
    void operator()(const SafePtr<DGVertex>& v);
  };
  struct VertexPrinter {
    VertexPrinter(std::ostream& ostr) : os(ostr) {}
    std::ostream& os;
    void operator()(const SafePtr<DGVertex>& v);
  };

};


#endif
