#pragma once
#include "backward_data_flow.h"

#include <vector>
#include <set>
#include <algorithm>

#include "SgUtils.h"
#include "../CFGraph.h"
#include "../IR.h"
#include "../RD_subst.h"

/* Note: this file should be included in backward_data_flow.h to provide template definitions */

// minimizes the number of blocks beween the ends of back edges
template <class NodeType>
std::vector<SAPFOR::BasicBlock*>
BackwardDataFlowAnalysis<NodeType>::reorderSequence(const std::vector<SAPFOR::BasicBlock*>& blocks,
                                                    const std::set<SAPFOR::BasicBlock*> back_edge_sources)
{
    std::vector<SAPFOR::BasicBlock*> res = { };

    auto blocks_end = blocks.rend();
    for (auto it = blocks.rbegin(); it < blocks_end; it++)
    {
        SAPFOR::BasicBlock* curr = *it;
        auto res_end = res.end();
        auto inserter = res.begin();
        if (back_edge_sources.count(curr) == 0)
        {
            auto curr_next_begin = curr->getNext().begin();
            auto curr_next_end = curr->getNext().end();
            while (inserter < res_end && std::find(curr_next_begin, curr_next_end, *inserter) == curr_next_end)
                inserter++;
        }

        res.insert(inserter, curr);
    }

    return res;
}

template <class NodeType>
void BackwardDataFlowAnalysis<NodeType>::fit(const std::vector<SAPFOR::BasicBlock*>& blocks)
{
    std::set<std::pair<SAPFOR::BasicBlock*, SAPFOR::BasicBlock*>> back_edges = {};

    bool returned = false;
    std::map<SAPFOR::BasicBlock*, std::set<SAPFOR::BasicBlock*>> back_edges_by_src;

    auto blocks_sorted = sortCfgNodes(blocks, &back_edges);

    std::set<SAPFOR::BasicBlock*> back_edge_sources;

    for (auto& edge : back_edges)
    {
        back_edges_by_src[edge.first].insert(edge.second);
        back_edge_sources.insert(edge.first);
    }

    back_edges.clear();

    blocks_sorted = reorderSequence(blocks_sorted, back_edge_sources);
    back_edge_sources.clear();

    std::reverse(blocks_sorted.begin(), blocks_sorted.end());

    this->nodes.clear();
    std::map<SAPFOR::BasicBlock*, NodeType*> node_by_block;

    for (auto block : blocks_sorted)
    {
        NodeType* node = this->createNode(block);
        this->nodes.push_back(node);
        node_by_block[block] = node;
    }

    int nodes_size = this->nodes.size();

    for (int i = 0; i < nodes_size; i++)
    {
        NodeType* node = this->nodes[i];

        auto back_edges_by_src_it = back_edges_by_src.find(node->getBlock());
        if (back_edges_by_src_it != back_edges_by_src.end())
        {
            // This node is a source for back edge
            for (auto dest : back_edges_by_src_it->second)
            {
                auto node_by_block_it = node_by_block.find(dest);
                if (node_by_block_it != node_by_block.end())
                    node_by_block_it->second->getRollback().insert(i);
            }
        }

        for (auto next : node->getBlock()->getNext())
        {
            auto node_by_block_it = node_by_block.find(next);
            if (node_by_block_it != node_by_block.end())
                node->getPrevBlocks().insert(node_by_block_it->second);
        }
    }
}