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

#include "private_arrays_search.h"
#include "range_structures.h"
#include "region.h"
#include "SgUtils.h"
#include "graph_loops.h"
#include "CFGraph/CFGraph.h"

using namespace std;

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

    for (auto& [arrayName, arrayRanges] : region->getHeader()->array_out)
    {
        for (Region* byBlock : region->getBasickBlocks())
        {
            AccessingSet intersection = byBlock->array_def[arrayName].Intersect(arrayRanges);
            region->array_def[arrayName] = region->array_def[arrayName].Union(intersection);
        }
    }

    for (auto& byBlock : region->getBasickBlocks()) 
    {
        for (auto& [arrayName, arrayRanges] : byBlock->array_use)
        {
            AccessingSet diff = byBlock->array_use[arrayName].Diff(byBlock->array_in[arrayName]);
            region->array_use[arrayName] = region->array_use[arrayName].Union(diff);
        }
    }
    ArrayAccessingIndexes useUnion;
    for (auto& byBlock : region->getBasickBlocks())
        for (auto& [arrayName, arrayRanges] : byBlock->array_use)
            useUnion[arrayName] = useUnion[arrayName].Union(byBlock->array_use[arrayName]);

    for (auto& [arrayName, arrayRanges] : useUnion)
        region->array_priv[arrayName] = useUnion[arrayName].Diff(region->array_use[arrayName]);
    
    for (Region* prevBlock : region->getHeader()->getPrevRegions()) 
        prevBlock->replaceInNextRegions(region, region->getHeader());
    
    for (Region* nextBlock : region->getHeader()->getNextRegions())
        nextBlock->replaceInPrevRegions(region, region->getHeader());
}

static void SolveDataFlowIteratively(Region* DFG) 
{
    unordered_set<Region*> worklist(DFG->getBasickBlocks());
    do
    {
        Region* b = *worklist.begin();
        ArrayAccessingIndexes newIn;
        bool flagFirst = true;
        for (Region* prevBlock : b->getPrevRegions())
        {
            if (flagFirst)
            {
                newIn = prevBlock->array_out;
                flagFirst = false;
            }
            else
            {
                if (prevBlock->array_out.empty())
                {
                    newIn.clear();
                    continue;
                }
                for (const auto& [arrayName, accessSet] : prevBlock->array_out)
                {
                    if (newIn.find(arrayName) != newIn.end())
                        newIn[arrayName] = newIn[arrayName].Intersect(accessSet);                    
                    else
                        newIn[arrayName] = AccessingSet();
                }
            }
        }

        b->array_in = move(newIn);
        ArrayAccessingIndexes newOut;
        if (b->array_def.empty())        
            newOut = b->array_in;
        else if (b->array_in.empty())
            newOut = b->array_def;
        else
        {
            for (auto& [arrayName, accessSet] : b->array_def)
            {
                if (newOut.find(arrayName) != newOut.end())
                    newOut[arrayName] = b->array_def[arrayName].Union(b->array_in[arrayName]);
                else
                    newOut[arrayName] = accessSet;
            }
        }

        /* can not differ */
        if (newOut != b->array_out) 
            b->array_out = newOut;
        else
            worklist.erase(b);
    } 
    while (!worklist.empty());
}

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

    SolveDataFlowIteratively(DFG);
    for (Region* subRegion : DFG->getSubRegions())
        SolveDataFlow(subRegion);
    Collapse(DFG);
}

unsigned long long CalculateLength(const AccessingSet& array) 
{
    if (array.GetElements().empty())
        return 0;
    unsigned long long result = 1;
    for (const auto& range : array.GetElements()) 
    {
        for (const auto& dim : range)
        {
            result *= (dim.step * dim.tripCount);
        }
    }
    return result;
}

void AddPrivateArraysToLoop(LoopGraph* loop, ArrayAccessingIndexes privates)
{
    SgStatement* spfStat = new SgStatement(SPF_ANALYSIS_DIR);
    spfStat->setlineNumber(loop->lineNum);
    spfStat->setFileName((char*)loop->loop->fileName());
    SgExpression* toAdd = new SgExpression(EXPR_LIST, new SgExpression(ACC_PRIVATE_OP), NULL, NULL);
    set<SgSymbol*> arraysToInsert;
    for (const auto& [_, accessingSet] : privates)
    {
        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));
    }

    loop->loop->addAttribute(SPF_ANALYSIS_DIR, spfStat, sizeof(SgStatement));
}

map<LoopGraph*, ArrayAccessingIndexes>  FindPrivateArrays(map<string, vector<LoopGraph*>> &loopGraph, map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR)
{  
    map<LoopGraph*, ArrayAccessingIndexes> result;
    for (const auto& [fileName, loops] : loopGraph)
    {
        SgFile::switchToFile(fileName);
        for (const auto& loop : loops)
        {
            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 = new Region(loop, blocks);
                    if (loopRegion->getBasickBlocks().size() <= 1)
                    {
                        delete(loopRegion);
                        continue;
                    }
                    SolveDataFlow(loopRegion);
                    result[loop] = loopRegion->array_priv;
                    delete(loopRegion);
                }
            }
            if (result.find(loop) != result.end() && !result[loop].empty())
            {
                AddPrivateArraysToLoop(loop, result[loop]);
            }
        }
    }
    return result;
}