#include "errors.h"
#include "private_variables_analysis.h"
#include "graph_loops.h"
#include "../LoopAnalyzer/loop_analyzer.h"
#include "../SageAnalysisTool/depGraph.h"
#include "../DirectiveProcessing/directive_parser.h"
#include "live_variable_analysis.h"
#include "RD_subst.h"


#include <string>
#include <vector>
#include <map>
#include <set>
#include <unordered_set>
#include <unordered_map>
#include <list>

using std::string;
using std::wstring;
using std::pair;
using std::vector;
using std::map;
using std::set;
using std::unordered_set;
using std::unordered_map;
using std::list;

static map<string, pair<set<int>, set<SAPFOR::Argument*>>> outForFunc;
static unordered_map<SgSymbol*, string> global_cache, local_cache;

static unordered_map <SAPFOR::Argument*, SgSymbol*> argumentToSymbol;

#define DEBUG_PRINT_PRIVATE 0

#define PRINT_PRIVATES 1
#define PRINT_WARNINGS 1
// print privates and lastprivates for loop
void printLoopInfo(const LoopGraph* loop)
{
    SgForStmt* loop_stmt = isSgForStmt(loop->loop);

    const set<string>& priv = loop->privateScalars;
    const set<string>& lastpriv = loop->lastprivateScalars;
    
    if(!loop_stmt)
        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

    __spf_print(PRINT_PRIVATES, "  loop in file '%s' at line %d\n", loop->fileName.c_str(), loop->lineNum);
    __spf_print(PRINT_PRIVATES, "    privates:");
    for(const auto& ident : priv)
        __spf_print(PRINT_PRIVATES, " %s", ident.c_str());
    __spf_print(PRINT_PRIVATES, "\n    lastprivates:");
    for (const auto& ident : lastpriv)
        __spf_print(PRINT_PRIVATES, " %s", ident.c_str());
    __spf_print(PRINT_PRIVATES, "\n");

    if (PRINT_WARNINGS)
    {
        set<string> added;
        for (auto& data : getAttributes<SgStatement*, SgStatement*>(loop_stmt, set<int>{ SPF_ANALYSIS_DIR }))
            fillPrivatesFromComment(new Statement(data), added);


        set<string> extra_old;
        std::set_difference(added.begin(), added.end(), priv.begin(), priv.end(), std::inserter(extra_old, extra_old.begin()));

        if (extra_old.size() != 0)
        {
            __spf_print(PRINT_WARNINGS, "  [WARNING] extra private variables:");
            for (const auto& ident : extra_old)
                __spf_print(PRINT_WARNINGS, " %s", ident.c_str());
            __spf_print(PRINT_WARNINGS, "\n");
        }
    }
}

// check if parent statement contains stmt
static bool isParentStmt(SgStatement* stmt, SgStatement* parent)
{
    for (; stmt; stmt = stmt->controlParent())
        if (stmt == parent)
            return true;
    
    return false;
}

// iterate over expr tree and find SgSymbol that matches arg (by name)
static SgSymbol* findSymbolByArgInExpression(const vector<pair<const Variable*, CommonBlock*>>& commonVars,
                                             const FuncInfo* func, const SAPFOR::Argument* arg, SgExpression* expr)
{
    if (!expr)
        return NULL;
    
    if (expr->symbol() && expr->symbol()->identifier() && getArgNameBySymbol(commonVars, func, expr->symbol(), global_cache, local_cache) == arg->getValue())
        return expr->symbol();
    
    SgSymbol* left;
    if (left = findSymbolByArgInExpression(commonVars, func, arg, expr->lhs()))
        return left;

    return findSymbolByArgInExpression(commonVars, func, arg, expr->rhs());
}

// find SgSymbol that matches arg (by name) in stmt
static SgSymbol* findSymbolByArgInStatement(const vector<pair<const Variable*, CommonBlock*>>& commonVars,
                                            const FuncInfo* func, const SAPFOR::Argument* arg, SgStatement* stmt)
{
    if (!stmt)
        return NULL;

    if (stmt->hasSymbol() && stmt->symbol() &&
        getArgNameBySymbol(commonVars, func, stmt->symbol(), global_cache, local_cache) == arg->getValue())
    {
        return stmt->symbol();
    }

    for (int i = 0; i < 3; i++) 
    {
        SgSymbol* found = findSymbolByArgInExpression(commonVars, func, arg, stmt->expr(i));
        if (found)
            return found;
    }

    return NULL;
}

// find SgSymbol for arg in instr's operator and fill argumentToSymbol 
// do nothing if arg already has found symbol
static void addPlaceWithDef(const vector<pair<const Variable*, CommonBlock*>>& commonVars,
                            const FuncInfo* func, SAPFOR::Argument* arg, SAPFOR::Instruction* instr) {
    if (argumentToSymbol.find(arg) != argumentToSymbol.end())
        return;

    SgStatement* stmt_with_decl = instr->getOperator();
    if (!stmt_with_decl)
        return;

    SgSymbol* found = findSymbolByArgInStatement(commonVars, func, arg, stmt_with_decl);

    if (found)
        argumentToSymbol[arg] = found;

}

static void getDefined(const string& func, const vector<SAPFOR::Argument*>& params,
                       set<SAPFOR::Argument*>& result)
{
    if (!joinGlobalsWithParameters(params, outForFunc, func, result))
    {
        if (func == string("_WRITE") || isIntrinsic(func.c_str()))
            return;
        else //READ or else
            insertIfVar(params.begin(), params.end(), result);
    }
}

static void fillOutForFunc(const FuncInfo* func, const vector<SAPFOR::BasicBlock*>& blocks)
{
    if (blocks.size() == 0)
        return;

    set<SAPFOR::BasicBlock*> exits = {};

    for (auto block : blocks)
    {
        if (block->getNext().size() == 0)
            exits.insert(block);
    }

    set<int> defined;
    set<SAPFOR::Argument*> common_defined;

    for (auto byExit : exits)
    {
        for (const auto& byRd : byExit->getRD_Out())
        {
            auto out_arg = byRd.first;
            if (!(byRd.second.size() == 1 && *(byRd.second.begin()) < 0))
            {
                switch (out_arg->getMemType())
                {
                case SAPFOR::CFG_MEM_TYPE::COMMON_:
                case SAPFOR::CFG_MEM_TYPE::MODULE_:
                    common_defined.insert(out_arg);
                    break;
                case SAPFOR::CFG_MEM_TYPE::FUNC_PARAM_:
                    {
                        int num = getParamIndex(out_arg, func->funcParams.countOfPars);
                        if(func->funcParams.isArgOut(num))
                            defined.insert(num);
                        break;
                    }
                }
            }
        }
    }


    outForFunc[func->funcName] = { defined, common_defined };
}

static void getDefsFromBlock(SAPFOR::BasicBlock* block, set<SAPFOR::Argument*>& res, 
                             const vector<pair<const Variable*, CommonBlock*>>& commonVars,
                             const FuncInfo* func)
{
    vector<SAPFOR::Argument*> lastParamRef;

    for (auto ir_block : block->getInstructions())
    {
        SAPFOR::Instruction* instr = ir_block->getInstruction();
        SAPFOR::CFG_OP instr_operation = instr->getOperation();
        if (instr_operation == SAPFOR::CFG_OP::PARAM)
        {
            SAPFOR::Argument* arg = instr->getArg1();
            if(arg->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
                addPlaceWithDef(commonVars, func, arg, instr);

            lastParamRef.push_back(arg);
        }
        else if (instr_operation == SAPFOR::CFG_OP::F_CALL)
        {
            int count = stoi(instr->getArg2()->getValue());
            if (lastParamRef.size() != count)
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

            const string& fName = instr->getArg1()->getValue();

            getDefined(fName, lastParamRef, res);

            lastParamRef.clear();
        }
    }

    int first_instr_num = block->getInstructions().front()->getNumber();
    int last_instr_num = block->getInstructions().back()->getNumber();

    for (const auto& def : block->getRD_Out())
        for (int place : def.second)
            if (place >= first_instr_num && place <= last_instr_num && def.first->getType() == SAPFOR::CFG_ARG_TYPE::VAR)
            {
                res.insert(def.first);
                addPlaceWithDef(commonVars, func, def.first, block->getInstructions()[place - first_instr_num]->getInstruction());
            }
}

// recursively analyze FOR loops 
static set<SAPFOR::BasicBlock*> analyzeLoop(LoopGraph* loop, const set<SAPFOR::BasicBlock*>& blocks,
                                            const vector<pair<const Variable*, CommonBlock*>>& commonVars,
                                            const map<string, SgSymbol*>& commonArgs, FuncInfo* func,
                                            map<string, vector<Messages>>& messages)
{
    if (!loop->isFor())
        printInternalError(convertFileName(__FILE__).c_str(), __LINE__); //should be called only with FOR loops

    SgStatement* loop_operator = loop->loop->GetOriginal();

    // find blocks related to the current loop
    set<SAPFOR::BasicBlock*> currentLoop = { };

    /*
    service block that looks like
        _reg = _reg < _reg
        IF_FALSE _reg then goto ...
    */
    SAPFOR::BasicBlock* head_block = NULL;

    int loop_start = loop->lineNum, loop_end = loop->lineNumAfterLoop;
    for (auto bb : blocks)
    {
        if (!bb || (bb->getInstructions().size() == 0))
            continue;

        SgStatement* first = bb->getInstructions().front()->getInstruction()->getOperator();
        SgStatement* last = bb->getInstructions().back()->getInstruction()->getOperator();

        if (isParentStmt(first, loop_operator) && isParentStmt(last, loop_operator))
        {
            currentLoop.insert(bb);

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

    if (!head_block)
        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
    
    // find blocks inside loop wich can lead to head_block
    list<SAPFOR::BasicBlock*> search_stack = { head_block };
    set<SAPFOR::BasicBlock*> cycleBody = { };

    while (search_stack.size() > 0)
    {
        SAPFOR::BasicBlock* curr = search_stack.front();
        search_stack.pop_front();

        if (cycleBody.insert(curr).second)
            for (auto next : curr->getPrev())
                if (currentLoop.count(next))
                    search_stack.push_back(next);
    }

    set<SAPFOR::Argument*> useValueInLoop = {};

    for (const auto& use : head_block->getLiveIn())
    {
        for (SAPFOR::BasicBlock* place : use.second)
        {
            if (currentLoop.count(place))
            {
                useValueInLoop.insert(use.first);
                break;
            }
        }
    }

    set<SAPFOR::Argument*> changeValueInLoop = { };
    set<SAPFOR::Argument*> changeValueOnExit = { };
    set<SAPFOR::Argument*> LiveWhenLoopEnds = { };

    auto loop_it = currentLoop.begin();
    auto loop_end_it = currentLoop.end();

    auto body_it = cycleBody.begin();
    auto body_end = cycleBody.end();

    while ((loop_it != loop_end_it) && (body_it != body_end))
    {
        if (*loop_it == *body_it)
        {
            getDefsFromBlock(*loop_it, changeValueInLoop, commonVars, func);
            body_it++;
        }
        else
            getDefsFromBlock(*loop_it, changeValueOnExit, commonVars, func);
        loop_it++;
    }

    if (body_it != body_end)
        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

    for (; loop_it != loop_end_it; loop_it++)
        getDefsFromBlock(*loop_it, changeValueOnExit, commonVars, func);


    for (auto bb : currentLoop)
    {
        //fill LiveWhenLoopEnds
        bool has_next_outside_body = false;
        for (const auto& next : bb->getNext())
        {
            if (cycleBody.find(next) == cycleBody.end())
            {
                has_next_outside_body = true;
                break;
            }
        }

        if (has_next_outside_body)
        {
            for (const auto& use : bb->getLiveOut())
                for (SAPFOR::BasicBlock* place : use.second)
                    if (currentLoop.count(place) == 0)
                    {
                        LiveWhenLoopEnds.insert(use.first);
                        break;
                    }
        }
    }

    set<SAPFOR::Argument*> arg_private = {}, arg_lastprivate = {};

    std::set_difference(changeValueInLoop.begin(), changeValueInLoop.end(), 
                        useValueInLoop.begin(), useValueInLoop.end(),
                        std::inserter(arg_private, arg_private.begin()));

    arg_private.insert(changeValueOnExit.begin(), changeValueOnExit.end());

    std::set_intersection(arg_private.begin(), arg_private.end(), 
                          LiveWhenLoopEnds.begin(), LiveWhenLoopEnds.end(),
                          std::inserter(arg_lastprivate, arg_lastprivate.begin()));

    changeValueInLoop.clear();
    LiveWhenLoopEnds.clear();
    useValueInLoop.clear();

    vector<const depNode*> sym_private;
    for (auto e : arg_private)
    {
        SgSymbol* arg_sym = NULL;

        auto it = argumentToSymbol.find(e);
        if (it != argumentToSymbol.end())
            arg_sym = it->second;
        else if (e->getMemType() == SAPFOR::CFG_MEM_TYPE::COMMON_ || e->getMemType() == SAPFOR::CFG_MEM_TYPE::MODULE_) 
            continue;
        
        if (arg_sym && arg_sym->identifier())
        {
            if (arg_lastprivate.find(e) != arg_lastprivate.end())
            {
                loop->lastprivateScalars.insert(arg_sym->identifier());

                const char* identifier = arg_sym->identifier();
                wstring messageE, messageR;
                __spf_printToLongBuf(messageE, L"add lastprivate scalar '%s' to loop on line %d", to_wstring(identifier).c_str(), loop->lineNum);
                __spf_printToLongBuf(messageR, R200, to_wstring(identifier).c_str(), loop->lineNum);

                messages[loop->fileName].push_back(Messages(NOTE, loop->lineNum, messageR, messageE, 3002));
            }
            else
            {
                SgStatement* s = new SgStatement(0);
                s->setlineNumber(loop_start);

                SgExpression* exp = new SgExpression(VAR_REF);
                exp->setSymbol(arg_sym);

                sym_private.push_back(new depNode(s, s, exp, exp, 0, 0, NULL, NULL, 0));
                loop->privateScalars.insert(arg_sym->identifier());
            }
        }
        else
            printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
    }

    addPrivatesToLoops(loop, sym_private, { {loop_start, isSgForStmt(loop_operator)} }, messages[loop->fileName]);

    for (const depNode* node : sym_private)
    {
        delete node->stmtin;
        delete node;
    }

    arg_private.clear();
    arg_lastprivate.clear();

    printLoopInfo(loop);
    return currentLoop;
}


static void recAnalyzeLoop(LoopGraph* loop, const set<SAPFOR::BasicBlock*>& blocks,
                           const vector<pair<const Variable*, CommonBlock*>>& commonVars,
                           const map<string, SgSymbol*>& commonArgs,
                           FuncInfo* func, map<string, vector<Messages>>& messages)
{
    const auto& loop_body = loop->isFor() ? analyzeLoop(loop, blocks, commonVars, commonArgs, func, messages) : blocks;
    
    for (const auto& inner_loop : loop->children)
        recAnalyzeLoop(inner_loop, loop_body, commonVars, commonArgs, func, messages);
}

void runPrivateVariableAnalysis(const map<string, vector<LoopGraph*>>& loopGraph,
                                const map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& CFGraph_for_project,
                                const map<string, CommonBlock*>& commonBlocks,
                                map<string, vector<Messages>>& messages)
{
    map<string, FuncInfo*> funcByName;
    for (const auto& byFunc : CFGraph_for_project)
        if (byFunc.first)
            funcByName[byFunc.first->funcName] = byFunc.first;


    //fill outForFunc
    for (const auto& byLoopFunc : loopGraph)
    {
        for (const auto& byLoop : byLoopFunc.second)
        {
            for (const auto& byFuncCall : byLoop->calls)
            {
                const string& fname = byFuncCall.first;
                if (outForFunc.find(fname) == outForFunc.end())
                {
                    auto func_it = funcByName.find(fname);
                    if (func_it != funcByName.end())
                    {
                        FuncInfo* func = func_it->second;
                        auto cfg_it = CFGraph_for_project.find(func);
                        if (cfg_it == CFGraph_for_project.end())
                            printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

                        fillOutForFunc(cfg_it->first, cfg_it->second);
                    }
                }
            }
        }
    }

    map<FuncInfo*, map<string, SgSymbol*>> commonArgsByFunc = { };
    for (const auto& byFile : loopGraph)
    {
        SgFile::switchToFile(byFile.first);

        for (auto loop : byFile.second)
        {
            int loop_start = loop->lineNum, loop_end = loop->lineNumAfterLoop;

            SgStatement* search_func = loop->loop->GetOriginal();
            
            while (search_func && (!isSgProgHedrStmt(search_func)))
                search_func = search_func->controlParent();
                
            if(!search_func)
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__); //loop statement outside any function statement

            bool loop_analyzed = false;
            for (const auto& byFunc : CFGraph_for_project)
            {
                if (byFunc.first->fileName == byFile.first && byFunc.first->funcPointer->GetOriginal() == search_func)
                {
                    const auto& commonVars = getCommonsByFunction(current_file, byFunc.first->funcPointer, commonBlocks);
                    const auto& commonArgs = getCommonArgsByFunc(byFunc.first, commonArgsByFunc, commonVars, global_cache, local_cache);

                    set<SAPFOR::BasicBlock*> loop_ir;
                    loop_ir.insert(byFunc.second.begin(), byFunc.second.end());
                    recAnalyzeLoop(loop, loop_ir, commonVars, commonArgs, byFunc.first, messages);
                    argumentToSymbol.clear(); //clear cache
                    loop_analyzed = true;
                    break;
                }
            }

            if(!loop_analyzed)
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__); //no func found for loop
        }
        
        local_cache.clear();
        commonArgsByFunc.clear();
    }

    outForFunc.clear();
    global_cache.clear();
}