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

#include "CFGraph/CFGraph.h"
#include "Distribution/Array.h"
#include "errors.h"
#include "graph_loops.h"
#include "private_arrays_search.h"
#include "range_structures.h"
#include "region.h"
#include "SgUtils.h"
#include "Transformations/ArrayConstantPropagation/propagation.h"
#include "utils.h"
#include "Utils/AstWrapper.h"

using namespace std;

static set<Region*> collapsed;

static void removeEmptyPoints(arrayAccessingIndexes& container)
{
    arrayAccessingIndexes resultContainer;
    set<string> toRemove;

    for (auto& [arrayName, accessingSet] : container)
    {
        vector<vector<ArrayDimension>> points;
        for (auto& arrayPoint : accessingSet.GetElements())
        {
            if (!arrayPoint.empty())
                points.push_back(arrayPoint);
        }

        if (!points.empty())
            resultContainer[arrayName] = points;
        else
            toRemove.insert(arrayName);
    }

    for (const string& name : toRemove)
        container.erase(name);

    for (auto& [arrayName, accessingSet] : resultContainer)
        container[arrayName] = accessingSet;
}

static void collapse(Region* region)
{
    if (region->getBasickBlocks().empty())
        return;

    bool firstRegion = true;
    int blockCount = 0;
    for (Region* basickBlock : region->getBasickBlocks())
    {
        if (basickBlock->getNextRegions().empty())
        {
            if (firstRegion)
            {
                region->array_def = basickBlock->array_out;
                firstRegion = false;
            }
            else
            {
                set<string> toErease;
                for (auto& [arrayName, arrayRanges] : region->array_def)
                {
                    if (basickBlock->array_out.find(arrayName) != basickBlock->array_out.end())
                        arrayRanges = arrayRanges.Intersect(basickBlock->array_out[arrayName]);
                    else
                    {
                        arrayRanges = AccessingSet();
                        toErease.insert(arrayName);
                    }
                }
                for (string arrayName : toErease)
                    region->array_def.erase(arrayName);
            }
        }
    }

    RegionInstruction instruction;
    instruction.def = move(region->array_def);

    arrayAccessingIndexes recursivePriv;
    for (auto& byBlock : region->getBasickBlocks())
    {
        if (!byBlock->array_priv.empty())
            recursivePriv = byBlock->array_priv;
        for (auto& instruction : byBlock->instructions)
        {
            if (!instruction.def.empty() || !instruction.use.empty())
                blockCount++;
            for (auto& [arrayName, _] : instruction.use)
            {
                AccessingSet diff = instruction.use[arrayName].Diff(instruction.in[arrayName]);
                region->array_use[arrayName] = region->array_use[arrayName].Union(diff);
            }
        }
    }

    arrayAccessingIndexes useUnionB;
    for (auto& byBlock : region->getBasickBlocks())
        for (auto& instruction : byBlock->instructions)
            for (auto& [arrayName, _] : instruction.use)
                useUnionB[arrayName] = useUnionB[arrayName].Union(instruction.use[arrayName]);

    for (auto& [arrayName, _] : useUnionB)
        region->array_priv[arrayName] = useUnionB[arrayName].Diff(region->array_use[arrayName]);

    instruction.use = move(region->array_use);

    for (Region* prevBlock : region->getHeader()->getPrevRegions())
    {
        prevBlock->replaceInNextRegions(region, region->getHeader());
        region->addPrevRegion(prevBlock);
    }

    for (Region* nextBlock : region->getHeader()->getNextRegions())
    {
        nextBlock->replaceInPrevRegions(region, region->getHeader());
        region->addNextRegion(nextBlock);
    }
    region->instructions.push_back(instruction);
    if (blockCount == 1 && !recursivePriv.empty())
        region->array_priv = move(recursivePriv);
}

static void solveForBasickBlock(Region* block)
{
    arrayAccessingIndexes newIn;
    bool flagFirst = true;
    for (Region* prevBlock : block->getPrevRegions())
    {
        if (flagFirst)
        {
            newIn = prevBlock->array_out;
            flagFirst = false;
        }
        else
        {
            if (prevBlock->array_out.empty())
            {
                newIn.clear();
                break;
            }

            for (const auto& [arrayName, accessSet] : prevBlock->array_out)
            {
                if (newIn.find(arrayName) != newIn.end())
                    newIn[arrayName] = newIn[arrayName].Intersect(accessSet);
                else
                    newIn[arrayName] = AccessingSet();
            }
        }
    }

    if (block->instructions.empty())
        block->instructions.push_back(RegionInstruction());

    block->instructions[0].in = move(newIn);

    for (int i = 0; i < block->instructions.size(); i++)
    {
        auto& instruction = block->instructions[i];

        if (i > 0)
            instruction.in = block->instructions[i - 1].out;

        arrayAccessingIndexes newOut;
        if (instruction.def.empty())
            newOut = instruction.in;
        else if (instruction.in.empty())
            newOut = instruction.def;
        else
        {
            for (auto& [arrayName, accessSet] : instruction.def)
            {
                if (instruction.in.find(arrayName) != instruction.in.end())
                    newOut[arrayName] = instruction.def[arrayName].Union(instruction.in[arrayName]);
                else
                    newOut[arrayName] = accessSet;
            }
            for (auto& [arrayName, accessSet] : instruction.in)
            {
                if (newOut.find(arrayName) == newOut.end())
                    newOut[arrayName] = accessSet;
            }
        }

        instruction.out = move(newOut);
    }
    if (!block->instructions.empty())
        block->array_out = block->instructions.back().out;
}

static void solveDataFlowTopologically(Region* DFG)
{
    for (Region* b : DFG->getBasickBlocks())
    {
        collapsed.insert(b);
        solveForBasickBlock(b);
    }
}

static void solveDataFlow(Region* DFG)
{
    if (!DFG)
        return;

    for (Region* subRegion : DFG->getSubRegions())
    {
        solveDataFlow(subRegion);
        DFG->addBasickBlocks(subRegion);
    }
    vector<Region*>& blocks = DFG->getBasickBlocks();
    auto pos = remove_if(blocks.begin(), blocks.end(), [](Region* r) { return collapsed.find(r) != collapsed.end(); });
    blocks.erase(pos, blocks.end());
    topologySort(DFG->getBasickBlocks(), DFG->getHeader());
    solveDataFlowTopologically(DFG);
    collapse(DFG);
}

static bool getArrayDeclaredDimensions(SgArrayRefExp* arrayRef, vector<uint64_t>& declaredDims)
{
    if (!arrayRef || !arrayRef->symbol() || !isSgArrayType(arrayRef->symbol()->type()))
        return false;
    SgArrayType* arrayType = (SgArrayType*)arrayRef->symbol()->type();
    int dimCount = arrayType->dimension();
    for (int i = 0; i < dimCount; i++)
    {
        SgExpression* sizeExpr = arrayType->sizeInDim(i);
        SgConstantSymb* constValSymb = isSgConstantSymb(sizeExpr->symbol());
        SgSubscriptExp* subscriptExpr = isSgSubscriptExp(sizeExpr);
        uint64_t dimLength;
        if (sizeExpr && sizeExpr->variant() == INT_VAL)
            dimLength = stol(sizeExpr->unparse());
        else if (constValSymb)
            dimLength = stol(constValSymb->constantValue()->unparse());
        else if (subscriptExpr)
            dimLength = stol(subscriptExpr->rhs()->unparse()) - stol(subscriptExpr->lhs()->unparse());
        else
            return false;

        if (dimLength == 0)
            return false;
        declaredDims.push_back(dimLength);
    }
    return true;
}

static DIST::Array* getDistArrayBySymbol(SgSymbol* arrSym, const map<tuple<int, string, string>, pair<DIST::Array*, DIST::ArrayAccessInfo*>>& declaredArrays)
{
    if (!arrSym)
        return nullptr;
    for (auto& [key, val] : declaredArrays)
    {
        DIST::Array* distArr = val.first;
        if (!distArr)
            printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
        Symbol* declSym = distArr->GetDeclSymbol();
        if (!declSym)
            printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
        SgSymbol* sgDecl = declSym->GetOriginal();
        if (sgDecl && isEqSymbols(sgDecl, arrSym))
            return distArr;
    }
    return nullptr;
}

static bool checkDimensionLength(const AccessingSet& array, const map<tuple<int, string, string>, pair<DIST::Array*, DIST::ArrayAccessInfo*>>& declaredArrays)
{
    if (array.GetElements().empty())
        return false;
    size_t dimCount = array.GetElements()[0].size();
    SgArrayRefExp* arrayRef = array.GetElements()[0][0].array;
    if (!arrayRef || !arrayRef->symbol())
        return false;

    vector<uint64_t> declaredDims;
    declaredDims.reserve(dimCount);

    DIST::Array* distArr = getDistArrayBySymbol(arrayRef->symbol(), declaredArrays);
    if (distArr && distArr->GetDimSize() == (int)dimCount)
    {
        const auto& sizes = distArr->GetSizes();
        bool valid = true;
        for (size_t i = 0; i < dimCount && valid; ++i)
        {
            int lo = sizes[i].first;
            int hi = sizes[i].second;
            if (lo > hi)
                valid = false;
            else
                declaredDims.push_back((uint64_t)(hi - lo + 1));
        }
        if (valid && declaredDims.size() == dimCount)
        {
            vector<ArrayDimension> testArray(dimCount);
            for (size_t i = 0; i < dimCount; i++)
                testArray[i] = { 1, 1, declaredDims[i], nullptr };
            return AccessingSet({ testArray }).Diff(array).GetElements().empty();
        }
    }

    return false;
}

static int addPrivateArraysToLoop(LoopGraph* loop, const arrayAccessingIndexes& privates, set<SgStatement*>& insertedPrivates,
                                   map<string, vector<Messages>>& SPF_messages,
                                   const map<tuple<int, string, string>, pair<DIST::Array*, DIST::ArrayAccessInfo*>>& declaredArrays)
{
    SgStatement* spfStat = new SgStatement(SPF_ANALYSIS_DIR);
    spfStat->setlineNumber(loop->loop->lineNumber());
    spfStat->setFileName(loop->loop->fileName());
    SgExpression* toAdd = new SgExpression(EXPR_LIST, new SgExpression(ACC_PRIVATE_OP), NULL, NULL);
    set<SgSymbol*> arraysToInsert;
    for (const auto& [arrayName, accessingSet] : privates)
    {
        int idx = arrayName.find('%');
        string name = (idx != -1 ? arrayName.substr(idx+1) : arrayName);
        if (!checkDimensionLength(accessingSet, declaredArrays))
        {
            wstring messageE, messageR;
            __spf_printToLongBuf(messageE, L"Private array '%s' was skipped because dimension lengths are inconsistent", to_wstring(name).c_str());
            __spf_printToLongBuf(messageR, R159, to_wstring("array " + name + " has inconsistent dimension lengths").c_str());
            SPF_messages[loop->loop->fileName()].push_back(Messages(WARR, loop->loop->lineNumber(), messageR, messageE, 1029));
            continue;
        }

        for (const auto& arrayElement : accessingSet.GetElements())
        {
            if (arrayElement.empty())
                continue;
            arraysToInsert.insert(arrayElement[0].array->symbol());
        }
    }

    spfStat->setExpression(0, *toAdd);
    toAdd = toAdd->lhs();
    bool first = true;
    for (auto& elem : arraysToInsert)
    {
        if (first)
        {
            toAdd->setLhs(new SgExpression(EXPR_LIST));
            toAdd = toAdd->lhs();
            first = false;
        }
        else
        {
            toAdd->setRhs(new SgExpression(EXPR_LIST));
            toAdd = toAdd->rhs();
        }
        toAdd->setLhs(new SgVarRefExp(elem));
    }

    if (arraysToInsert.size() != 0)
    {
        loop->loop->insertStmtBefore(*spfStat, *loop->loop->controlParent());
        insertedPrivates.insert(spfStat);
    }
}

int findPrivateArrays(map<string, vector<LoopGraph*>>& loopGraph, map<FuncInfo*, 
                       vector<SAPFOR::BasicBlock*>>& FullIR, 
                       set<SgStatement*>& insertedPrivates,
                       map<string, vector<Messages>>& SPF_messages,
                       const map<tuple<int, string, string>, pair<DIST::Array*, DIST::ArrayAccessInfo*>>& declaredArrays)
{
    map<LoopGraph*, arrayAccessingIndexes> result;
    for (const auto& [fileName, loops] : loopGraph)
    {
        if (SgFile::switchToFile(fileName) == -1)
            printInternalError(convertFileName(__FILE__).c_str(), __LINE__);

        for (const auto& loop : loops)
        {
            if (!loop->isFor())
                continue;
            SgStatement* search_func = loop->loop->GetOriginal();

            while (search_func && (!isSgProgHedrStmt(search_func)))
                search_func = search_func->controlParent();

            for (const auto& [funcInfo, blocks] : FullIR)
            {
                if (funcInfo->fileName == fileName && funcInfo->funcPointer->GetOriginal() == search_func)
                {
                    Region* loopRegion = nullptr;
                    try
                    {
                        loopRegion = new Region(loop, blocks);
                    }
                    catch (...)
                    {
                        printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
                    }

                    if (loopRegion->getBasickBlocks().size() <= 1)
                    {
                        delete loopRegion;
                        continue;
                    }

                    solveDataFlow(loopRegion);
                    removeEmptyPoints(loopRegion->array_priv);
                    result[loop] = loopRegion->array_priv;
                    delete loopRegion;
                }
            }

            if (result.find(loop) != result.end() && !result[loop].empty())
                addPrivateArraysToLoop(loop, result[loop], insertedPrivates, SPF_messages, declaredArrays);
        }
    }

    for (const auto& [loop, accesing] : result)
    {
        if (accesing.size())
        {
            __spf_print(1, "found for loop on line %d in file %s\n", loop->lineNum, loop->fileName.c_str());
            for (const auto& [name, accesingSet] : accesing)
            {
                const auto& byDimention = accesingSet.GetElements();
                __spf_print(1, " for array %s with dimention %d\n", name.c_str(), byDimention.size());

                for (int z = 0; z < byDimention.size(); ++z)
                {
                    __spf_print(1, "  dim %d (start, step, tripCount):\n", z);
                    for (auto& elem : byDimention[z])
                        __spf_print(1, "   [%ld %ld %ld]\n", elem.start, elem.step, elem.tripCount);
                }
            }
        }
    }

    return insertedPrivates.size();
}