#include "select_array_conf.h"
#include "../Utils/utils.h"

using std::map;
using std::string;
using std::vector;
using std::set;
using std::pair;
using std::wstring;

using std::inserter;
using std::copy;

bool IsSetsIntersect(const set<DIST::Array*>& lhs, const set<DIST::Array*>& rhs) {
    if (lhs.empty() || rhs.empty())
        return false;
    
    if (lhs.size() * 100 < rhs.size()) 
    {
        for (const auto& x : lhs)
            if (rhs.find(x) != rhs.end())
                return true;
        return false;
    }

    if (rhs.size() * 100 < lhs.size()) 
    {
        for (const auto& x : rhs)
            if (lhs.find(x) != lhs.end())
                return true;
        return false;
    }

    auto l_it = lhs.begin(), l_end = lhs.end();
    auto r_it = rhs.begin(), r_end = rhs.end();

    while (l_it != l_end && r_it != r_end) 
    {
        if (*l_it < *r_it) 
        {
            l_it = lhs.lower_bound(*r_it);
            continue;
        }
        if (*r_it < *l_it) 
        {
            r_it = rhs.lower_bound(*l_it);
            continue;
        }

        return true;
    }

    return false;
}

static void findUsedArraysInParallelLoops(LoopGraph* loop, set<DIST::Array*>& res) {
    if(loop->directive)
        copy(loop->usedArraysAll.begin(), loop->usedArraysAll.end(), inserter(res, res.end()));
    else
        for(LoopGraph* child : loop->children)
            findUsedArraysInParallelLoops(child, res);
}

static void preventLoopsFromParallelizations(LoopGraph* loop, const set<DIST::Array*>& prevent, 
                                             vector<Directive*>& createdDirectives) {
    if (loop->directive)
    {
        if (IsSetsIntersect(prevent, loop->usedArraysAll))
        {
            // prevent this loop
            int loopLine = loop->lineNum;

            for (auto dir_it = createdDirectives.begin(); dir_it != createdDirectives.end(); dir_it++)
            {
                if ((*dir_it)->line == loopLine)
                {
                    delete *dir_it;
                    dir_it = createdDirectives.erase(dir_it);

                    break;
                }
            }

            delete loop->directive;
            loop->directive = NULL;
        }
    }
    else
    {
        for (LoopGraph* child : loop->children)
            preventLoopsFromParallelizations(child, prevent, createdDirectives);
    }
}

struct DimConf
{
    vector<pair<int, int>> dims;

    DimConf(DIST::Array* a) : dims(a->GetSizes()) { }

    friend bool operator<(const DimConf& l, const DimConf& r) { return l.dims < r.dims; }
};

static map<DimConf, map<FuncInfo*, set<DIST::Array*>>>::const_iterator
pickBest(const map<DimConf, map<FuncInfo*, set<DIST::Array*>>>& arrs)
{
    const int undefined = -1;
    int max_elems = undefined;
    auto best_it = arrs.begin();

    for (auto it = arrs.begin(); it != arrs.end(); it++)
    {
        int elems = 1;
        for (const auto& p : it->first.dims)
        {
            if(p.first >= 0 && p.second >= p.first)
            {
                elems *= p.second - p.first;
            }
            else
            {
                elems = undefined;
                break;
            }
        }

        if (elems > max_elems)
        {
            max_elems = elems;
            best_it = it;
        }
    }

    return best_it;
}

void SelectArrayConfForParallelization(SgProject* proj, map<string, vector<FuncInfo*>>& funcByFile,
                                       const map<string, vector<LoopGraph*>>& loopGraph, map<std::string, vector<Directive*>>& createdDirectives,
                                       map<string, vector<Messages>>& allMessages, const map<DIST::Array*, set<DIST::Array*>>& arrayLinksByFuncCalls, 
                                       const vector<ParallelRegion*>& regions)
{
    map<string, FuncInfo*> funcByName;
    for (const auto& byFile : funcByFile)
        for(const auto& byFunc : byFile.second)
            funcByName[byFunc->funcName] = byFunc;
    
    // array(real ref)    dims     func           array in func       
    map<DIST::Array*, map<DimConf, map<FuncInfo*, set<DIST::Array*>>>> usedArrays;

    for (const auto& byFile : loopGraph)
    {
        SgFile::switchToFile(byFile.first);

        auto& loops = byFile.second;

        auto file_funcs_it = funcByFile.find(byFile.first);

        if(file_funcs_it == funcByFile.end())
            printInternalError(convertFileName(__FILE__).c_str(), __LINE__); // no such file in funcByFile

        map<FuncInfo*, set<DIST::Array*>> usedInLoops;

        for (const auto& loop : loops)
        {
            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 : file_funcs_it->second)
            {
                if (byFunc->funcPointer->GetOriginal() == search_func)
                {
                    if(!loop->usedArraysAll.empty())
                        findUsedArraysInParallelLoops(loop, usedInLoops[byFunc]);

                    loop_analyzed = true;
                    break;
                }
            }

            if (!loop_analyzed)
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__); //no func found for loop
        }
        
        for (const auto& byFunc : usedInLoops)
        {
            for (DIST::Array* arr : byFunc.second)
            {
                set<DIST::Array*> realRefs;
                getRealArrayRefs(arr, arr, realRefs, arrayLinksByFuncCalls);

                int num_of_dims = arr->GetDimSize();
                for(DIST::Array* ref : realRefs)
                    usedArrays[ref][DimConf(arr)][byFunc.first].insert(arr);
            }
        }
    }

    map<FuncInfo*, set<DIST::Array*>> preventFromParallelization;

    for (const auto& byRealRef : usedArrays)
    {
        DIST::Array* realRef = byRealRef.first;
        auto& refferedDims = byRealRef.second;

        // prevent from parallelization all configurations except best one

        auto best_conf_it = pickBest(refferedDims);

        for (auto by_worse_dim_conf_it = refferedDims.begin(); by_worse_dim_conf_it != refferedDims.end(); by_worse_dim_conf_it++)
        {
            if(by_worse_dim_conf_it != best_conf_it)
            {
                for (const auto& byFunc : by_worse_dim_conf_it->second)
                {
                    FuncInfo* f = byFunc.first;
                    auto& to_prevent = byFunc.second;

                    auto& destSet = preventFromParallelization[f];
                    copy(to_prevent.begin(), to_prevent.end(), inserter(destSet, destSet.end()));
                }
            }
        }
    }

    for (const auto& byFile : loopGraph)
    {
        auto dirs_it = createdDirectives.find(byFile.first);
        if (dirs_it != createdDirectives.end()) 
        {
            SgFile::switchToFile(byFile.first);

            auto& loops = byFile.second;

            auto file_funcs_it = funcByFile.find(byFile.first);

            if (file_funcs_it == funcByFile.end())
                printInternalError(convertFileName(__FILE__).c_str(), __LINE__); // no such file in funcByFile

            map<FuncInfo*, set<DIST::Array*>> usedInLoops;

            for (const auto& loop : loops)
            {
                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 : file_funcs_it->second)
                {
                    if (byFunc->funcPointer->GetOriginal() == search_func)
                    {
                        auto prevent_it = preventFromParallelization.find(byFunc);
                        if (prevent_it != preventFromParallelization.end())
                            preventLoopsFromParallelizations(loop, prevent_it->second, dirs_it->second);

                        loop_analyzed = true;
                        break;
                    }
                }

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

void removeRegionsWithoutDirs(const map<string, vector<Directive*>>& createdDirectives,
                             vector<ParallelRegion*>& parallelRegions,
                             const map<string, vector<FuncInfo*>>& allFuncInfo,
                             map<string, vector<Messages>>& SPF_messages)
{
    set<ParallelRegion*> empty_regs(parallelRegions.begin(), parallelRegions.end());

    for (const auto& byFile : createdDirectives)
        for (const auto& dir : byFile.second)
            empty_regs.erase(getRegionByLine(parallelRegions, byFile.first, dir->line));

    set<int> idxToDel;
    for (int z = 0; z < parallelRegions.size(); ++z)
    {
        if (empty_regs.find(parallelRegions[z]) != empty_regs.end())
        {
            __spf_print(1, "  no parallel directives for parallel region '%s'\n", parallelRegions[z]->GetName().c_str());

            if (parallelRegions[z]->GetId() == 0) // DEFAULT
            {
                wstring bufE, bufR;
                __spf_printToLongBuf(bufE, L"Can not find arrays or free loops for distribution in this project");
                __spf_printToLongBuf(bufR, R130);

                for (auto& funcByFile : allFuncInfo)
                {
                    vector<Messages>& fileM = getObjectForFileFromMap(funcByFile.first.c_str(), SPF_messages);
                    for (auto& func : funcByFile.second)
                    {
                        auto stat = func->funcPointer->GetOriginal();
                        if (stat->variant() == PROG_HEDR)
                            fileM.push_back(Messages(ERROR, stat->lineNumber(), bufR, bufE, 3010));
                    }
                }
            }
            else
            {
                wstring bufE, bufR;
                __spf_printToLongBuf(bufE, L"Can not find arrays or free loops for distribution in this region");
                __spf_printToLongBuf(bufR, R131);

                for (auto& linesByFile : parallelRegions[z]->GetAllLines())
                {
                    vector<Messages>& fileM = getObjectForFileFromMap(linesByFile.first.c_str(), SPF_messages);
                    for (auto& lines : linesByFile.second)
                        if (!lines.isImplicit())
                            fileM.push_back(Messages(ERROR, lines.lines.first, bufR, bufE, 3010));
                }
            }
            idxToDel.insert(z);
        }
    }

    vector<ParallelRegion*> newParReg;
    for (int z = 0; z < parallelRegions.size(); ++z)
    {
        if (idxToDel.find(z) != idxToDel.end())
        {
            ParallelRegion* regToDel = parallelRegions[z];
#ifdef _WIN32
            removeFromCollection(parallelRegions[z]);
#endif
            delete parallelRegions[z];
        }
        else
            newParReg.push_back(parallelRegions[z]);
    }
    parallelRegions.clear();
    parallelRegions = newParReg;

    if (parallelRegions.size() == 0)
        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
}