#ifndef LLVM_TRANSFORMS_SCALAR_LOOPCHAIN_PRE_H
#define LLVM_TRANSFORMS_SCALAR_LOOPCHAIN_PRE_H

#include <map>
#include <string>
#include <vector>

#include "llvm/ADT/DirectedGraph.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/IR/PassManager.h"

#include "llvm/Analysis/MemorySSA.h"
#include "llvm/Analysis/AliasAnalysis.h"
#define DEBUG_TYPE "pre-loopchain"
#define PRE_LOOPCHAIN_PASS_PREFIX "loopchain_"
namespace llvm{

  class PreLCNode;
  class PreLCEdge;
  using PreLCNodeBase = DGNode<PreLCNode, PreLCEdge>;
  using PreLCEdgeBase = DGEdge<PreLCNode, PreLCEdge>;
  using PreLCBase = DirectedGraph<PreLCNode, PreLCEdge>;
}


namespace llvm {
struct PreLoopChainOptions{
  std::optional<bool> FilterJustLeapfrog;
  std::optional<unsigned> MinimalLCSize;

  PreLoopChainOptions() = default;

  //Enable or disable filtering for leapfrog routine name
  PreLoopChainOptions &setFilterJustLeapfrog(bool O){
    FilterJustLeapfrog = O;
    return *this;
  }
  // Set minimal number of loops in the loopchain
  PreLoopChainOptions &setMinimalLCSize(int O){
    MinimalLCSize = O;
    return *this;
  }
};

class PreLCNode : public PreLCNodeBase {

public:
  llvm::Loop* Loop = nullptr;
  int nestingLvl = 0;
  int num = 0;
  PreLCNode* parent = nullptr;
  PreLCEdge* edgeToParent = nullptr;

  int maxLoopChain = 0;
  bool partOfLoopChain = false;

  PreLCNode() = default;
  PreLCNode(PreLCNode &N) = delete;
  PreLCNode(llvm::Loop* L) : PreLCNodeBase(), Loop(L) {}
  PreLCNode(const bool &isroot) : PreLCNodeBase(), isRoot(isroot) {}
  //~PreLCNode();

  bool isRootNode() const { return isRoot; }
  PreLCNode*& getParent() { return parent; }
  PreLCEdge*& getEdgeToParent() { return edgeToParent; }
  int& getMaxLoopChain()  { return maxLoopChain; }
  
  //Update the maximal LoopChain size of the node
  void updateMaxLoopChainSize();
  
  //Detect whether Node is already part of LoopChain
  bool isPartOfLoopChain() const ;
  
  //Find edge to node
  bool findEdgeTo(const PreLCNode &N, PreLCEdge* &EdgeToN);

  protected:
    bool isRoot=false; 
};

class PreLCEdge : public PreLCEdgeBase {

public:
  int LCId = -1; 
  PreLCEdge(PreLCNode &N) : PreLCEdgeBase(N) {}
  //~PreLCEdge();

};

struct LCStorageNode{
  PreLCNode* node = nullptr;
  int iterator_begin = -1;
  int LCsize = -1;
  int LCId = -1;
};

template <typename NodeType> class PreLoopChainGraphInfo {
public:

  PreLoopChainGraphInfo() = delete;
  PreLoopChainGraphInfo(const PreLoopChainGraphInfo &G) = delete;
  PreLoopChainGraphInfo(const std::string name): Name(name){}
  virtual ~PreLoopChainGraphInfo() = default;

  /// Return the label that is used to name this graph.
  StringRef getName() const { return Name; }

protected:
  // Name of the graph.
  std::string Name;

};

using PreLCInfo = PreLoopChainGraphInfo<PreLCNode>;

class PreLoopChainGraph final : public PreLCBase, public PreLCInfo {

public:
  using NodeType = PreLCNode;
  using EdgeType = PreLCEdge;

  PreLoopChainGraph(const PreLoopChainGraph &) = delete;
  PreLoopChainGraph(Function &F, LoopInfo &li) : PreLoopChainGraphInfo(F.getName().str()), LI(li) {}
  ~PreLoopChainGraph(){}

  bool buildPreLoopChainGraph();
  void zeroOutMaxLCForChildren(PreLCNode &node);
  void createEdgesForNode(NodeType &N, unsigned const &nestingLevel, llvm::Loop* L);
  /// Add the given node \p N to the graph.
  bool addNode(PreLCNode &N);
  int numberOfNodes(){ return Nodes.size(); }

  //Greedy algorithm for detecting Loop Chains
  bool detectLoopChainsGreedy(const int &minLCSize = 2);

  NodeType* createNode(Loop *L){
    if (L == nullptr) {
      errs() << "Error createNode called with invalid Loop arg\n"; 
      return nullptr;
    }
    NodeType* node = new NodeType(L);
    return node;
  }

  NodeType* createRootNode(){
    NodeType *root = new NodeType(true);
    //We can a root node only if the graph is empty
    if (!Nodes.empty())
      return nullptr;

    if (!addNode(*root))
      return nullptr;
    
    return root;
  }

  EdgeType* createEdge(NodeType *N){
    if (N == nullptr) {
      errs() << "Error createEdge called with invalid Node arg\n";
      return nullptr;
    }
    EdgeType* edge = new EdgeType(*N);
    return edge;
  }
  
  //Find node with larges loop chain
  NodeType* findNodeWithLargestLC(int &size);
  
  //For a given node find a loop chain of its subloops that has at least LCSize.
  //The LC will be stored in LCStorage. The edges to the subloops will be marked
  //with the LCId to mark them as part of LC.
  bool findLargestLCForNode(PreLCNode &node, const int &LCsize, const int &LCId);

  void addLC(PreLCNode &N, const int &LCId, const int &LCSize, const int &iterator_begin){
    LCStorageNode mark;
    mark.node = &N;
    mark.LCId = LCId;
    mark.LCsize = LCSize;
    mark.iterator_begin = iterator_begin;
    LCStorage.push_back(mark);
  }

  //Updates the maxLCSize for the node and all its parents
  //We need to update the information to all childs nodes that they
  //should not be considered for other loop chains. This is done
  //by findNodeWithLargestLC routine as this is faster.
  bool updateLCInformation(NodeType* node, const int &LCId){
    NodeType* current;
    if (!node){
      errs() << "Internal Error in updateLCInformation node is nullptr\n";
      return false;
    }
    
    node->updateMaxLoopChainSize();
    current = node;

    //Update LoopChain information for all nodes on the path up to the root
    while(!current->isRootNode()){
      NodeType* parent = current->getParent();
      EdgeType* edgeToChild = nullptr;
      
      if (!parent->findEdgeTo(*current, edgeToChild))
        errs() << "Internal error in updateLCInformation.\n";
      edgeToChild->LCId = LCId;
      parent->updateMaxLoopChainSize();
      current = parent;
    }
    //Set the maxLooppChainSize for all loops that are part of the LC to zero
    for (auto i: node->getEdges()){
      if (i->LCId == LCId)
        zeroOutMaxLCForChildren(i->getTargetNode());
    }
    return true;
  }

  LoopInfo &LI;

  SmallVector<LCStorageNode> LCStorage;

};


raw_ostream &operator<<(raw_ostream &OS, const PreLCNode &N);
raw_ostream &operator<<(raw_ostream &OS, const PreLCEdge &E);
raw_ostream &operator<<(raw_ostream &OS, const PreLoopChainGraph &G);

class PreLoopChainPass : public PassInfoMixin<PreLoopChainPass> {

public:
  PreLoopChainOptions PreLoopChainOpts;
  unsigned MinimalLCSize = 2;
  bool UserMinimalLCSize = false;

  PreLoopChainPass(PreLoopChainOptions PreLCOpts = {}): PreLoopChainOpts(PreLCOpts) {}
  PreLoopChainPass(bool leapfrog =false): onlyLeapfrog(leapfrog) {}
  bool onlyLeapfrog = false;
  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
  void countBlocksInLoop(Loop *L, unsigned const &nesting, unsigned &maxLoopChainNumber);
  void countLoopChains(llvm::LoopInfo &LI);
  void detectLoopChains(Function &F, llvm::LoopInfo &LI, AliasAnalysis &AA, MemorySSA &MSSA);
  
/// @brief Splits the guarding basic block if it large than one instruction
/// @param guard - Pointer to the guarding basic block
/// @param guardingInstruction - Pointer to the guarding instruction
/// @param newBB - pointer the guarding instruction after splitting. If no splitting performed is equal to guard
/// @return - true if actual split is performed
  bool splitLoopGuard(BasicBlock *&guard, Instruction *guardingInstruction, BasicBlock *&newBB);
  bool splitLoopGuard(BasicBlock *&guard, Instruction *guardingInstruction, llvm::DominatorTree *DT,  llvm::LoopInfo *LI, llvm::MemorySSAUpdater *MSSAU, BasicBlock *&newBB);
  bool getFirstAndLastBasicBlockOfLC(LCStorageNode &Storage, llvm::BasicBlock *&firstBB, llvm::BasicBlock *&lastBB, Function &F, bool &guardSplitted);
  void outlineLoopChains(Function &F, llvm::LoopInfo &LI, AliasAnalysis &AA, MemorySSA &MSSA, SmallVector<LCStorageNode> &LCs);
  void printPipeline(raw_ostream &OS,
                     function_ref<StringRef(StringRef)> MapClassName2PassName);
};

//==
// GraphTraits specialization for the LC -- used for printing
//==
/// non-const versions of the graph trait specializations for LC
template <> struct GraphTraits<PreLCNode *> {
  using NodeRef = PreLCNode *;

  static PreLCNode *PreLCGetTargetNode(DGEdge<PreLCNode, PreLCEdge> *P) {
    return &P->getTargetNode();
  }

  // Provide a mapped iterator so that the GraphTrait-based implementations can
  // find the target nodes without having to explicitly go through the edges.
  using ChildIteratorType =
      mapped_iterator<PreLCNode::iterator, decltype(&PreLCGetTargetNode)>;
  using ChildEdgeIteratorType = PreLCNode::iterator;

  static NodeRef getEntryNode(NodeRef N) { return N; }
  static ChildIteratorType child_begin(NodeRef N) {
    //errs() <<  "In child_begin N = " << N << " N->begin() " << *N->begin() << "\n";
    return ChildIteratorType(N->begin(), &PreLCGetTargetNode);
  }
  static ChildIteratorType child_end(NodeRef N) {
    return ChildIteratorType(N->end(), &PreLCGetTargetNode);
  }

  static ChildEdgeIteratorType child_edge_begin(NodeRef N) {
    return N->begin();
  }
  static ChildEdgeIteratorType child_edge_end(NodeRef N) { return N->end(); }
};

template <> struct GraphTraits<PreLoopChainGraph *> : public GraphTraits<PreLCNode *>{
  using nodes_iterator = PreLoopChainGraph::iterator;
  static NodeRef getEntryNode(PreLoopChainGraph *DG) {
    return *DG->begin();
  }
  static nodes_iterator nodes_begin(PreLoopChainGraph *DG) {
    return DG->begin();
  }
  static nodes_iterator nodes_end(PreLoopChainGraph *DG) { return DG->end(); }
};

/// const versions of the graph trait specializations for LC
template <> struct GraphTraits<const PreLCNode *> {
  using NodeRef = const PreLCNode *;

  static const PreLCNode *PreLCGetTargetNode(const DGEdge<PreLCNode, PreLCEdge> *P) {
    return &P->getTargetNode();
  }

  // Provide a mapped iterator so that the GraphTrait-based implementations can
  // find the target nodes without having to explicitly go through the edges.
  using ChildIteratorType =
      mapped_iterator<PreLCNode::const_iterator, decltype(&PreLCGetTargetNode)>;
  using ChildEdgeIteratorType = PreLCNode::const_iterator;

  static NodeRef getEntryNode(NodeRef N) { return N; }
  static ChildIteratorType child_begin(NodeRef N) {
    //errs() <<  "In child_begin N = " << N << " N->begin() " << *N->begin() << "\n";
    return ChildIteratorType(N->begin(), &PreLCGetTargetNode);
  }
  static ChildIteratorType child_end(NodeRef N) {
    return ChildIteratorType(N->end(), &PreLCGetTargetNode);
  }

  static ChildEdgeIteratorType child_edge_begin(NodeRef N) {
    return N->begin();
  }
  static ChildEdgeIteratorType child_edge_end(NodeRef N) { return N->end(); }
};

template <> struct GraphTraits<const PreLoopChainGraph *>
    : public GraphTraits<const PreLCNode *> {
  using nodes_iterator = PreLoopChainGraph::const_iterator;
  static NodeRef getEntryNode(const PreLoopChainGraph *DG) {
    return *DG->begin();
  }
  static nodes_iterator nodes_begin(const PreLoopChainGraph *DG) {
    return DG->begin();
  }
  static nodes_iterator nodes_end(const PreLoopChainGraph *DG) {
    return DG->end();
  }
};

//====================================================================================


} // namespace llvm

#endif // LLVM_TRANSFORMS_SCALAR_LOOPCHAIN_PRE_H