#include <map>
#include <unordered_set>
#include <vector>
#include <queue>
#include <iostream>

#include "../Utils/errors.h"
#include "../Utils/SgUtils.h"
#include "../GraphCall/graph_calls.h"
#include "../GraphCall/graph_calls_func.h"
#include "../CFGraph/CFGraph.h"
#include "../CFGraph/IR.h"
#include "../GraphLoop/graph_loops.h"
#include "swapOperators.h"

using namespace std;


unordered_set<int> loop_tags = {FOR_NODE/*, FORALL_NODE, WHILE_NODE, DO_WHILE_NODE*/};
unordered_set<int> importantDepsTags = {FOR_NODE, IF_NODE};
unordered_set<int> importantUpdDepsTags = {ELSEIF_NODE};
unordered_set<int> importantEndTags = {CONTROL_END};


vector<SAPFOR::IR_Block*> findInstructionsFromOperator(SgStatement* st, vector<SAPFOR::BasicBlock*> Blocks)
{
    vector<SAPFOR::IR_Block*> result;
    string filename = st -> fileName();
    for (auto& block: Blocks)
    {
        vector<SAPFOR::IR_Block*> instructionsInBlock = block -> getInstructions();
        for (auto& instruction: instructionsInBlock)
        {
            SgStatement* curOperator = instruction -> getInstruction() -> getOperator();
            if (curOperator -> lineNumber() == st -> lineNumber())
                result.push_back(instruction);
        }
    }
    return result;
}

vector<SAPFOR::BasicBlock*> findFuncBlocksByFuncStatement(SgStatement *st, map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR)
{
    vector<SAPFOR::BasicBlock*> result;
    Statement* forSt = (Statement*)st;
    for (auto& func: FullIR)
    {
        if (func.first -> funcPointer -> getCurrProcessFile() == forSt -> getCurrProcessFile()
            && func.first -> funcPointer -> lineNumber() == forSt -> lineNumber())
            result = func.second;
    }
    return result;
}

map<SgForStmt*, vector<SAPFOR::BasicBlock*>> findAndAnalyzeLoops(SgStatement *st, vector<SAPFOR::BasicBlock*> blocks)
{
    map<SgForStmt*, vector<SAPFOR::BasicBlock*>> result;
    SgStatement *lastNode = st->lastNodeOfStmt();
    while (st && st != lastNode)
    {
        if (loop_tags.find(st -> variant()) != loop_tags.end())
        {
            // part with find statements of loop
            SgForStmt *forSt = (SgForStmt*)st;
            SgStatement *loopBody = forSt -> body();
            SgStatement *lastLoopNode = st->lastNodeOfStmt();
            // part with find blocks and instructions of loops
            unordered_set<int> blocks_nums;
            while (loopBody && loopBody != lastLoopNode)
            {
                SAPFOR::IR_Block* IR = findInstructionsFromOperator(loopBody, blocks).front();
                if (blocks_nums.find(IR -> getBasicBlock() -> getNumber()) == blocks_nums.end())
                {
                    result[forSt].push_back(IR -> getBasicBlock());
                    blocks_nums.insert(IR -> getBasicBlock() -> getNumber());
                }
                loopBody = loopBody -> lexNext();
            }
            std::sort(result[forSt].begin(), result[forSt].end());
        }
        st = st -> lexNext();
    }
    return result;
}

map<SgStatement*, set<SgStatement*>> AnalyzeLoopAndFindDeps(SgForStmt* forStatement, vector<SAPFOR::BasicBlock*> loopBlocks, map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR)
{
    map<SgStatement*, set<SgStatement*>> result;
    for (SAPFOR::BasicBlock* bb: loopBlocks)
    {
        map<SAPFOR::Argument*, set<int>> blockReachingDefinitions = bb -> getRD_In();
        vector<SAPFOR::IR_Block*> instructions = bb -> getInstructions();
        for (SAPFOR::IR_Block* irBlock: instructions)
        {
            // TODO: Think about what to do with function calls and array references. Because there are also dependencies there that are not reflected in RD, but they must be taken into account
            SAPFOR::Instruction* instr = irBlock -> getInstruction();
            result[instr -> getOperator()];
            // take Argument 1 and it's RD and push operators to final set
            if (instr -> getArg1() != NULL)
            {
                SAPFOR::Argument* arg = instr -> getArg1();
                set<int> prevInstructionsNumbers = blockReachingDefinitions[arg];
                for (int i: prevInstructionsNumbers)
                {
                    SAPFOR::Instruction* foundInstruction = getInstructionAndBlockByNumber(FullIR, i).first;
                    if (foundInstruction != NULL)
                    {
                        SgStatement* prevOp = foundInstruction -> getOperator();
                        if (prevOp != forStatement && instr -> getOperator() != forStatement && instr -> getOperator() -> lineNumber() > prevOp -> lineNumber()
                            && prevOp -> lineNumber() > forStatement -> lineNumber())
                            result[instr -> getOperator()].insert(prevOp);
                    }
                }
            }
            // take Argument 2 (if exists) and it's RD and push operators to final set
            if (instr -> getArg2() != NULL)
            {
                SAPFOR::Argument* arg = instr -> getArg2();
                set<int> prevInstructionsNumbers = blockReachingDefinitions[arg];
                for (int i: prevInstructionsNumbers)
                {
                    SAPFOR::Instruction* foundInstruction = getInstructionAndBlockByNumber(FullIR, i).first;
                    if (foundInstruction != NULL)
                    {
                        SgStatement* prevOp = foundInstruction -> getOperator();
                        if (prevOp != forStatement && instr -> getOperator() != forStatement&& instr -> getOperator() -> lineNumber() > prevOp -> lineNumber()
                            && prevOp -> lineNumber() > forStatement -> lineNumber())
                            result[instr -> getOperator()].insert(prevOp);
                    }
                }
            }
            // update RD
            if (instr -> getResult() != NULL)
                blockReachingDefinitions[instr -> getResult()] = {instr -> getNumber()};
        }
    }
    return result;
}

void buildAdditionalDeps(SgForStmt* forStatement, map<SgStatement*, set<SgStatement*>>& dependencies)
{
    SgStatement* lastNode = forStatement->lastNodeOfStmt();
    vector<SgStatement*> importantDeps;
    SgStatement* st = (SgStatement*) forStatement;
    st = st -> lexNext();
    SgStatement* logIfOp = NULL;
    while (st && st != lastNode)
    {
        if(importantDeps.size() != 0)
        {
            if (st != importantDeps.back())
            {
                dependencies[st].insert(importantDeps.back());
            }
        }
        if (logIfOp != NULL)
        {
            dependencies[st].insert(logIfOp);
            logIfOp = NULL;
        }
        if (st -> variant() == LOGIF_NODE)
        {
            logIfOp = st;
        }
        if (importantDepsTags.find(st -> variant()) != importantDepsTags.end())
        {
            importantDeps.push_back(st);
        }
        if (importantUpdDepsTags.find(st -> variant()) != importantUpdDepsTags.end())
        {
            importantDeps.pop_back();
            importantDeps.push_back(st);
        }
        if (importantEndTags.find(st -> variant()) != importantEndTags.end())
        {
            if(importantDeps.size() != 0)
            {
                importantDeps.pop_back();
            }
        }
        st = st -> lexNext();
    }
}

struct ReadyOp {
    SgStatement* stmt;
    int degree;
    size_t arrival;
    ReadyOp(SgStatement* s, int d, size_t a): stmt(s), degree(d), arrival(a) {}
};

struct ReadyOpCompare {
    bool operator()(const ReadyOp& a, const ReadyOp& b) const {
        if (a.degree != b.degree)
            return a.degree > b.degree;
        else
            return a.arrival > b.arrival;
    }
};

vector<SgStatement*> scheduleOperations(const map<SgStatement*, set<SgStatement*>>& dependencies)
{
    // get all statements
    unordered_set<SgStatement*> allStmtsSet;
    for (const auto& pair : dependencies)
    {
        allStmtsSet.insert(pair.first);
        for (SgStatement* dep : pair.second)
        {
            allStmtsSet.insert(dep);
        }
    }
    vector<SgStatement*> allStmts(allStmtsSet.begin(), allStmtsSet.end());
    // count deps and build reversed graph
    unordered_map<SgStatement*, vector<SgStatement*>> graph;
    unordered_map<SgStatement*, int> inDegree;
    unordered_map<SgStatement*, int> degree;
    for (auto op : allStmts)
        inDegree[op] = 0;
    // find and remember initial dependencies
    unordered_set<SgStatement*> dependentStmts;
    for (const auto& pair : dependencies)
    {
        SgStatement* op = pair.first;
        const auto& deps = pair.second;
        degree[op] = deps.size();
        inDegree[op] = deps.size();
        if (!deps.empty())
            dependentStmts.insert(op);
        for (auto dep : deps)
            graph[dep].push_back(op);
    }
    for (SgStatement* op : allStmts)
    {
        if (!degree.count(op))
        {
            degree[op] = 0;
        }
    }
    // build queues
    using PQ = priority_queue<ReadyOp, vector<ReadyOp>, ReadyOpCompare>;
    PQ readyDependent;
    queue<SgStatement*> readyIndependent;
    size_t arrivalCounter = 0;
    for (auto op : allStmts)
    {
        if (inDegree[op] == 0)
        {
            if (dependentStmts.count(op))
            {
                readyDependent.emplace(op, degree[op], arrivalCounter++);
            }
            else
            {
                readyIndependent.push(op);
            }
        }
    }
    // main sort algorythm
    vector<SgStatement*> executionOrder;
    while (!readyDependent.empty() || !readyIndependent.empty())
    {
        SgStatement* current = nullptr;
        if (!readyDependent.empty())
        {
            current = readyDependent.top().stmt;
            readyDependent.pop();
        }
        else
        {
            current = readyIndependent.front();
            readyIndependent.pop();
        }
        executionOrder.push_back(current);
        for (SgStatement* neighbor : graph[current])
        {
            inDegree[neighbor]--;
            if (inDegree[neighbor] == 0) {
                if (dependentStmts.count(neighbor))
                {
                    readyDependent.emplace(neighbor, degree[neighbor], arrivalCounter++);
                }
                else
                {
                    readyIndependent.push(neighbor);
                }
            }
        }
    }
    return executionOrder;
}

void buildNewAST(SgStatement* loop, vector<SgStatement*>& newBody)
{
    SgStatement* endDo = loop->lastNodeOfStmt();
    SgStatement* st = loop;
    int lineNum = loop -> lineNumber() + 1;
    for (int i = 0; i < newBody.size(); i++)
    {
        st -> setLexNext(*newBody[i]);
        st = st -> lexNext();
        st -> setlineNumber(lineNum);
        lineNum++;
    }
    st -> setLexNext(*endDo);
}



void runSwapOperators(SgFile *file, std::map<std::string, std::vector<LoopGraph*>>& loopGraph, std::map<FuncInfo*, std::vector<SAPFOR::BasicBlock*>>& FullIR, int& countOfTransform)
{
    std::cout << "SWAP_OPERATORS Pass" << std::endl;    // to remove
    countOfTransform += 1;                              // to remove

    const int funcNum = file -> numberOfFunctions();
    for (int i = 0; i < funcNum; ++i)
    {
        SgStatement *st = file -> functions(i);
        vector<SAPFOR::BasicBlock*> blocks = findFuncBlocksByFuncStatement(st, FullIR);
        map<SgForStmt*, vector<SAPFOR::BasicBlock*>> loopsMapping = findAndAnalyzeLoops(st, blocks);
        for (pair<SgForStmt*, vector<SAPFOR::BasicBlock*>> loopForAnalyze: loopsMapping)
        {
            map<SgStatement*, set<SgStatement*>> dependencyGraph = AnalyzeLoopAndFindDeps(loopForAnalyze.first, loopForAnalyze.second, FullIR);
            // TODO: Write a function that will go through the operators and update all dependencies so that there are no mix-ups and splits inside the semantic blocks (for if, do and may be some other cases)
            buildAdditionalDeps(loopForAnalyze.first, dependencyGraph);
            cout << endl;
            int firstLine = loopForAnalyze.first -> lineNumber();
            int lastLine = loopForAnalyze.first -> lastNodeOfStmt() -> lineNumber();
            cout << "LOOP ANALYZE FROM " << firstLine << " TO " << lastLine << " RES" << endl;
            // for (auto &v: dependencyGraph) {
            //     cout << "OPERATOR: " << v.first -> lineNumber()  << "    " << v.first -> variant() << "\nDEPENDS ON:" << endl;
            //     if (v.second.size() != 0)
            //         for (auto vv: v.second)
            //             cout << vv -> lineNumber() << "  ";
            //     cout << endl;
            // }
            vector<SgStatement*> new_order = scheduleOperations(dependencyGraph);
            cout << "RESULT ORDER:" << endl;
            for (auto v: new_order) 
                if (v -> lineNumber() > firstLine) 
                    cout << v -> lineNumber() << endl;
            buildNewAST(loopForAnalyze.first, new_order);
            st = loopForAnalyze.first -> lexNext();
            while (st != loopForAnalyze.first -> lastNodeOfStmt())
            {
                cout << st -> lineNumber() << "    " << st -> sunparse() << endl;
                st = st -> lexNext();
            }
        }
    }

    return;
};