#include <algorithm>
#include <vector>
#include <map>
#include <set>
#include <map>
#include <string>
#include <numeric>
#include <iostream>

#include "range_structures.h"
#include "region.h"
#include "SgUtils.h"

using namespace std;

static bool isParentStmt(SgStatement* stmt, SgStatement* parent)
{
    for (; stmt; stmt = stmt->controlParent())
        if (stmt == parent)
            return true;

    return false;
}

/*returns head block and loop*/
pair<SAPFOR::BasicBlock*, set<SAPFOR::BasicBlock*>> getBasicBlocksForLoop(const LoopGraph* loop, const vector<SAPFOR::BasicBlock*> blocks)
{
    set<SAPFOR::BasicBlock*> block_loop;
    SAPFOR::BasicBlock* head_block = nullptr;
    auto loop_operator = loop->loop->GetOriginal();
    for (const auto& block : blocks)
    {
        if (!block || (block->getInstructions().size() == 0))
            continue;

        SgStatement* first = block->getInstructions().front()->getInstruction()->getOperator();
        SgStatement* last = block->getInstructions().back()->getInstruction()->getOperator();
        if (isParentStmt(first, loop_operator) && isParentStmt(last, loop_operator))
        {
            block_loop.insert(block);

            if ((!head_block) && (first == loop_operator) && (last == loop_operator) &&
                (block->getInstructions().size() == 2) &&
                (block->getInstructions().back()->getInstruction()->getOperation() == SAPFOR::CFG_OP::JUMP_IF))
            {
                head_block = block;
            }

        }
    }
    return { head_block, block_loop };
}

static void buildLoopIndex(map<SgStatement*, LoopGraph*>& loopForIndex, LoopGraph* loop) {
    string index = loop->loopSymbol();
    loopForIndex[loop->loop->GetOriginal()] = loop;

    for (const auto& childLoop : loop->children)
        buildLoopIndex(loopForIndex, childLoop);
}

static string findIndexName(int pos, SAPFOR::BasicBlock* block, map<string, LoopGraph*>& loopForIndex) {
    set<SAPFOR::Argument*> args = { block->getInstructions()[pos]->getInstruction()->getArg1() };

    for (int i = pos - 1; i >= 0; i--)
    {
        SAPFOR::Argument* res = block->getInstructions()[i]->getInstruction()->getResult();
        if (res && args.find(res) != args.end())
        {
            SAPFOR::Argument* arg1 = block->getInstructions()[i]->getInstruction()->getArg1();
            SAPFOR::Argument* arg2 = block->getInstructions()[i]->getInstruction()->getArg2();
            if (arg1)
            {
                string name = arg1->getValue();
                int idx = name.find('%');
                if (idx != -1 && loopForIndex.find(name.substr(idx + 1)) != loopForIndex.end())
                    return name.substr(idx + 1);
                else
                    args.insert(arg1);
            }

            if (arg2)
            {
                string name = arg2->getValue();
                int idx = name.find('%');
                if (idx != -1 && loopForIndex.find(name.substr(idx + 1)) != loopForIndex.end())
                    return name.substr(idx + 1);
                else
                    args.insert(arg2);
            }
        }
    }
    return "";
}

static int getDefUseArray(SAPFOR::BasicBlock* block, LoopGraph* loop, arrayAccessingIndexes& def, arrayAccessingIndexes& use, Region* region) {
    auto instructions = block->getInstructions();
    map<SgStatement*, LoopGraph*> loopForIndex;
    buildLoopIndex(loopForIndex, loop);
    for (int i = 0; i < instructions.size(); i++)
    {
        auto instruction = instructions[i];
        if (!instruction->getInstruction()->getArg1())
            continue;

        auto operation = instruction->getInstruction()->getOperation();
        auto type = instruction->getInstruction()->getArg1()->getType();
        if (operation == SAPFOR::CFG_OP::ASSIGN && instruction->getInstruction()->getResult()->getType() == SAPFOR::CFG_ARG_TYPE::ARRAY)
        {
            SgStatement* op = instruction->getInstruction()->getOperator();
            if (op && op->expr(0) && isArrayRef(op->expr(0)) && op->expr(0)->symbol() && op->expr(0)->type())
            {
                if (isSgArrayType(op->expr(0)->symbol()->type()))
                {
                    SgArrayType* arrayType = (SgArrayType*)op->expr(0)->symbol()->type();
                    int dimCount = ((SgArrayType*)op->expr(0)->symbol()->type())->dimension();
                    vector<ArrayDimension> point;
                    for (int i = 0; i < dimCount; i++)
                    {
                        string strDimLength = arrayType->sizeInDim(i)->unparse();
                        if (arrayType->sizeInDim(i)->variant() == INT_VAL && strDimLength != "0")
                            point.push_back({ 1ULL, 1ULL, (uint64_t)stoi(strDimLength), (SgArrayRefExp*)op->expr(0) });
                    }

                    if (point.size() == dimCount)
                    {
                        def[instruction->getInstruction()->getResult()->getValue()] = AccessingSet({ point });
                        RegionInstruction regionInstruction;
                        regionInstruction.def[instruction->getInstruction()->getResult()->getValue()] = AccessingSet({ point });
                    }
                }
            }
        }
        if ((operation == SAPFOR::CFG_OP::STORE || operation == SAPFOR::CFG_OP::LOAD) && type == SAPFOR::CFG_ARG_TYPE::ARRAY)
        {
            vector<SAPFOR::Argument*> index_vars;
            vector<int> refPos;
            string array_name = instruction->getInstruction()->getArg1()->getValue();
            int j = i - 1;
            while (j >= 0 && instructions[j]->getInstruction()->getOperation() == SAPFOR::CFG_OP::REF)
            {
                index_vars.push_back(instructions[j]->getInstruction()->getArg1());
                refPos.push_back(j);
                j--;
            }

            /*to choose correct dimension*/
            int n = index_vars.size();
            vector<ArrayDimension> accessPoint(n);

            auto* ref = isSgArrayRefExp(instruction->getInstruction()->getExpression());
            if (!ref)
                continue;
            int fillCount = 0;

            vector<pair<int, int>> coeffsForDims;
            int subs = ref->numberOfSubscripts();
            for (int i = 0; ref && i < ref->numberOfSubscripts(); ++i)
            {
                const vector<int*>& coeffs = getAttributes<SgExpression*, int*>(ref->subscript(i), set<int>{ INT_VAL });
                if (coeffs.size() == 1)
                {
                    const pair<int, int> coef(coeffs[0][0], coeffs[0][1]);
                    coeffsForDims.push_back(coef);
                }

            }
            coeffsForDims = { coeffsForDims.rbegin(), coeffsForDims.rend() };

            while (!index_vars.empty() && !refPos.empty() && !coeffsForDims.empty())
            {
                auto var = index_vars.back();
                int currentVarPos = refPos.back();
                ArrayDimension current_dim;
                if (var->getType() == SAPFOR::CFG_ARG_TYPE::CONST)
                    current_dim = { stoul(var->getValue()), 1, 1, ref };
                else
                {
                    string name, full_name = var->getValue();
                    int pos = full_name.find('%');
                    LoopGraph* currentLoop;

                    auto serachInstr = instruction->getInstruction()->getOperator();
                    while (serachInstr && serachInstr->variant() != FOR_NODE)
                        serachInstr = serachInstr->controlParent();

                    name = full_name.substr(pos + 1);
                    if (loopForIndex.find(serachInstr) != loopForIndex.end())
                        currentLoop = loopForIndex[serachInstr];
                    else
                        return -1;

                    uint64_t start = coeffsForDims.back().second * currentLoop->startVal + coeffsForDims.back().first;
                    uint64_t step = currentLoop->stepVal;
                    uint64_t iters = currentLoop->calculatedCountOfIters;
                    current_dim = { start, step, iters, ref };
                }

                if (current_dim.step != 0 && current_dim.tripCount != 0)
                {
                    accessPoint[n - index_vars.size()] = current_dim;
                    fillCount++;
                }
                index_vars.pop_back();
                refPos.pop_back();
                coeffsForDims.pop_back();
            }

            if (fillCount == accessPoint.size())
            {
                RegionInstruction instruction;
                if (operation == SAPFOR::CFG_OP::STORE)
                {
                    def[array_name].Insert(accessPoint);
                    instruction.def[array_name] = { { accessPoint } };
                }
                else
                {
                    instruction.use[array_name] = { { accessPoint } };
                    if (def.find(array_name) == def.end())
                        use[array_name].Insert(accessPoint);
                    else
                    {
                        AccessingSet element({ accessPoint });
                        use[array_name] = use[array_name].Union(element.Diff(def[array_name]));
                    }
                }
                region->instructions.push_back(instruction);
            }
        }
    }
    return 0;

}

static void removeHeaderConnection(SAPFOR::BasicBlock* header, const set<SAPFOR::BasicBlock*>& blockSet, map<SAPFOR::BasicBlock*, Region*>& bbToRegion)
{
    for (SAPFOR::BasicBlock* block : blockSet)
    {
        bool isCycleBlock = false;
        for (SAPFOR::BasicBlock* prevBlock : block->getPrev())
            isCycleBlock = isCycleBlock || (blockSet.find(prevBlock) != blockSet.end());

        if (isCycleBlock)
        {
            bbToRegion[block]->removeNextRegion(bbToRegion[header]);
            bbToRegion[header]->removePrevRegion(bbToRegion[block]);
        }
    }
}

static bool DFS(Region* block, vector<Region*>& result, const set<Region*>& cycleBlocks, map<Region*, int>& color)
{
    auto it = color.find(block);
    if (it != color.end())
    {
        if (it->second == 0)
            return false;
        if (it->second == 1)
            return true;
    }
    color[block] = 0;
    for (Region* nextBlock : block->getNextRegions())
    {
        if (cycleBlocks.find(nextBlock) == cycleBlocks.end())
            continue;
        if (!DFS(nextBlock, result, cycleBlocks, color))
            return false;
    }
    color[block] = 1;
    result.push_back(block);
    return true;
}

bool hasCycle(Region* block, const std::set<Region*>& cycleBlocks, std::set<Region*>& visitedBlocks)
{
    return false;
    if (visitedBlocks.find(block) != visitedBlocks.end())
        return true;
    visitedBlocks.insert(block);
    for (Region* nextBlock : block->getNextRegions())
    {
        if (cycleBlocks.find(nextBlock) != cycleBlocks.end() && hasCycle(nextBlock, cycleBlocks, visitedBlocks))
            return true;
    }
    return false;
}

bool topologySort(std::vector<Region*>& basikBlocks, Region* header)
{
    set<Region*> cycleBlocks(basikBlocks.begin(), basikBlocks.end());
    vector<Region*> result;
    map<Region*, int> color;
    if (!DFS(header, result, cycleBlocks, color))
        return false;
    reverse(result.begin(), result.end());
    basikBlocks = move(result);
    return true;
}

static void setConnections(map<SAPFOR::BasicBlock*, Region*>& bbToRegion, const set<SAPFOR::BasicBlock*>& blockSet)
{
    for (SAPFOR::BasicBlock* block : blockSet)
    {
        for (SAPFOR::BasicBlock* nextBlock : block->getNext())
            if (bbToRegion.find(nextBlock) != bbToRegion.end())
                bbToRegion[block]->addNextRegion(bbToRegion[nextBlock]);

        for (SAPFOR::BasicBlock* prevBlock : block->getPrev())
            if (bbToRegion.find(prevBlock) != bbToRegion.end())
                bbToRegion[block]->addPrevRegion(bbToRegion[prevBlock]);
    }
}

static Region* createSubRegion(LoopGraph* loop, const vector<SAPFOR::BasicBlock*>& Blocks, map<SAPFOR::BasicBlock*, Region*>& bbToRegion)
{
    Region* region = new Region;
    auto [header, blockSet] = getBasicBlocksForLoop(loop, Blocks);
    removeHeaderConnection(header, blockSet, bbToRegion);
    if (bbToRegion.find(header) != bbToRegion.end())
        region->setHeader(bbToRegion.at(header));
    else
    {
        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
        return NULL;
    }

    for (SAPFOR::BasicBlock* block : blockSet)
        if (bbToRegion.find(block) != bbToRegion.end())
            region->addBasickBlocks(bbToRegion.at(block));

    for (LoopGraph* childLoop : loop->children)
    {
        if (!childLoop->isFor())
            continue;
        region->addSubRegions(createSubRegion(childLoop, Blocks, bbToRegion));
    }
    if (!topologySort(region->getBasickBlocks(), region->getHeader()))
        throw std::runtime_error("Unnoticed cycle");
    return region;
}

Region::Region(LoopGraph* loop, const vector<SAPFOR::BasicBlock*>& Blocks)
{
    auto [header, blockSet] = getBasicBlocksForLoop(loop, Blocks);
    map<SAPFOR::BasicBlock*, Region*> bbToRegion;
    for (auto poiner : blockSet)
    {
        bbToRegion[poiner] = new Region(*poiner);
        this->basickBlocks.push_back(bbToRegion[poiner]);
        getDefUseArray(poiner, loop, bbToRegion[poiner]->array_def, bbToRegion[poiner]->array_use, bbToRegion[poiner]);

    }
    this->header = bbToRegion[header];
    setConnections(bbToRegion, blockSet);
    removeHeaderConnection(header, blockSet, bbToRegion);
    //create subRegions
    for (LoopGraph* childLoop : loop->children)
    {
        if (!childLoop->isFor())
            continue;
        subRegions.insert(createSubRegion(childLoop, Blocks, bbToRegion));
        }
    if (!topologySort(basickBlocks, this->header))
        throw std::runtime_error("Unnoticed cycle");
}