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

#include "private_arrays_search.h"
#include "../Utils/SgUtils.h"
#include "../GraphLoop/graph_loops.h"
#include "../CFGraph/CFGraph.h"

using namespace std;

void print_info(LoopGraph* loop)
{
    cout << "loopSymbol: " << loop->loopSymbol << endl;
    for (const auto& ops : loop->writeOpsForLoop)
    {
        cout << "Array name: " << ops.first->GetShortName() << endl;
        for (const auto i : ops.second)
        {
            i.printInfo();
        }
    }
    if (!loop->children.empty())
    {
        for (const auto child : loop->children)
        {
            print_info(child);
        }
    }
}

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

/*returns head block and loop*/
pair<SAPFOR::BasicBlock*, unordered_set<SAPFOR::BasicBlock*>> GetBasicBlocksForLoop(LoopGraph* loop, vector<SAPFOR::BasicBlock*> blocks)
{
    unordered_set<SAPFOR::BasicBlock*> block_loop;
    SAPFOR::BasicBlock* head_block = nullptr;
    auto loop_operator = loop->loop->GetOriginal();
    for (const auto& block : blocks)
    {
        if (!block || (block->getInstructions().size() == 0))
        {
            continue;
        }
        SgStatement* first = block->getInstructions().front()->getInstruction()->getOperator();
        SgStatement* last = block->getInstructions().back()->getInstruction()->getOperator();
        if (isParentStmt(first, loop_operator) && isParentStmt(last, loop_operator))
        {
            block_loop.insert(block);

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

        }
    }
    return { head_block, block_loop };
}


static void BuildLoopIndex(map<string, LoopGraph*>& loopForIndex, LoopGraph* loop) {
    string index = loop->loopSymbol;
    loopForIndex[index] = loop;
    for (const auto& childLoop : loop->children) {
        BuildLoopIndex(loopForIndex, childLoop);
    }
}

static string FindIndexName(int pos, SAPFOR::BasicBlock* block, map<string, LoopGraph*>& loopForIndex) {
    unordered_set<SAPFOR::Argument*> args = {block->getInstructions()[pos]->getInstruction()->getArg1()};

    for (int i = pos-1; i >= 0; i--) {
        SAPFOR::Argument* res = block->getInstructions()[i]->getInstruction()->getResult();
        if (res && args.find(res) != args.end()) {
            SAPFOR::Argument* arg1 = block->getInstructions()[i]->getInstruction()->getArg1();
            SAPFOR::Argument* arg2 = block->getInstructions()[i]->getInstruction()->getArg2();
            if (arg1) {
                string name = arg1->getValue();
                int idx = name.find('%');
                if (idx != -1 && loopForIndex.find(name.substr(idx + 1)) != loopForIndex.end())
                    return name.substr(idx + 1);
                else {
                    args.insert(arg1);
                }
            }
            if (arg2) {
                string name = arg2->getValue();
                int idx = name.find('%');
                if (idx != -1 && loopForIndex.find(name.substr(idx + 1)) != loopForIndex.end())
                    return name.substr(idx + 1);
                else {
                    args.insert(arg2);
                }
            }
        }
    }
    return "";
}

static int GetDefUseArray(SAPFOR::BasicBlock* block, LoopGraph* loop, ArrayAccessingIndexes& def, ArrayAccessingIndexes& use) {
    auto instructions = block->getInstructions();
    map<string, LoopGraph*> loopForIndex;
    BuildLoopIndex(loopForIndex, loop);
    for(int i = 0; i < instructions.size(); i++)
    {
        auto instruction = instructions[i];
        if(!instruction->getInstruction()->getArg1()) {
            continue;
        }
        auto operation = instruction->getInstruction()->getOperation();
        auto type = instruction->getInstruction()->getArg1()->getType();
        if ((operation == SAPFOR::CFG_OP::STORE  || operation == SAPFOR::CFG_OP::LOAD) && type == SAPFOR::CFG_ARG_TYPE::ARRAY)
        {
            vector<SAPFOR::Argument*> index_vars;
            vector<int> refPos;
            string array_name;
            if (operation == SAPFOR::CFG_OP::STORE)
            {
                array_name = instruction->getInstruction()->getArg1()->getValue();
            }
            else
            {
                array_name = instruction->getInstruction()->getArg2()->getValue();
            }
            int j = i - 1;
            while (j >= 0 && instructions[j]->getInstruction()->getOperation() == SAPFOR::CFG_OP::REF)
            {
                index_vars.push_back(instructions[j]->getInstruction()->getArg1());
                refPos.push_back(j);
                j--;
            }
            /*to choose correct dimension*/
            int n = index_vars.size();
            vector<ArrayDimension> accessPoint(n);
            /*if (operation == SAPFOR::CFG_OP::STORE)
            {
                if (def[array_name].empty())
                {
                    def[array_name].resize(n);
                }
            }
            else
            {
                if (use[array_name].empty())
                {
                    use[array_name].resize(n);
                }
            }*/

            SgArrayRefExp* ref = (SgArrayRefExp*)instruction->getInstruction()->getExpression();
            vector<pair<int, int>> coefsForDims;
            for (int i = 0; i < ref->numberOfSubscripts(); ++i)
            {
                const vector<int*>& coefs = getAttributes<SgExpression*, int*>(ref->subscript(i), set<int>{ INT_VAL });
                if (coefs.size() == 1)
                {
                    const pair<int, int> coef(coefs[0][0], coefs[0][1]);
                    coefsForDims.push_back(coef);
                }
                
            }

            while (!index_vars.empty())
            {
                auto var = index_vars.back();
                int currentVarPos = refPos.back();
                pair currentCoefs = coefsForDims.back();
                ArrayDimension current_dim;
                if (var->getType() == SAPFOR::CFG_ARG_TYPE::CONST) {
                    current_dim = { stoul(var->getValue()), 1, 1 };
                }
                else
                {
                    string name, full_name = var->getValue();
                    int pos = full_name.find('%');
                    LoopGraph* currentLoop;
                    if (pos != -1) {
                        name = full_name.substr(pos+1);
                        if (loopForIndex.find(name) != loopForIndex.end()) {
                            currentLoop = loopForIndex[name];
                        }
                        else {
                            return -1;
                        }
                    }
                    else {
                        name = FindIndexName(currentVarPos, block, loopForIndex);
                        if (name == "") {
                            return -1;
                        }
                        if (loopForIndex.find(name) != loopForIndex.end()) {
                            currentLoop = loopForIndex[name];
                        }
                        else {
                            return -1;
                        }
                    }
                    uint64_t start = currentLoop->startVal * currentCoefs.first + currentCoefs.second;
                    uint64_t step = currentCoefs.first;
                    current_dim = { start, step, (uint64_t)currentLoop->calculatedCountOfIters };
                }
                /*if (operation == SAPFOR::CFG_OP::STORE)
                {
                    def[array_name][n - index_vars.size()].push_back(current_dim);
                }
                else
                {
                    use[array_name][n - index_vars.size()].push_back(current_dim);
                }*/
                accessPoint[n - index_vars.size()] = current_dim;
                index_vars.pop_back();
                refPos.pop_back();
                coefsForDims.pop_back();
            }
            if (operation == SAPFOR::CFG_OP::STORE)
            {
                def[array_name].Insert(accessPoint);
            }
            else
            {
                use[array_name].Insert(accessPoint);
            }
        }
    }
    return 0;

}

static vector<uint64_t> FindParticularSolution(const ArrayDimension& dim1, const ArrayDimension& dim2) 
{
    for (uint64_t i = 0; i < dim1.tripCount; i++) 
    {
        uint64_t leftPart = dim1.start + i * dim1.step;
        for (uint64_t j = 0; j < dim2.tripCount; j++) 
        {
            uint64_t rightPart = dim2.start + j * dim2.step;
            if (leftPart == rightPart) 
            {
                return {i, j};
            }
        }
    }
    return {};
}

/* dim1 /\ dim2 */
static ArrayDimension* DimensionIntersection(const ArrayDimension& dim1, const ArrayDimension& dim2) 
{
    vector<uint64_t> partSolution =  FindParticularSolution(dim1, dim2);
    if (partSolution.empty()) 
    {
        return NULL;
    }
    int64_t x0 = partSolution[0], y0 = partSolution[1];
    /* x = x_0 + c * t */
    /* y = y_0 + d * t */
    int64_t c = dim2.step / gcd(dim1.step, dim2.step);
    int64_t d = dim1.step / gcd(dim1.step, dim2.step);
    int64_t tXMin, tXMax, tYMin, tYMax;
    tXMin = -x0 / c;
    tXMax = (dim1.tripCount - 1 - x0) / c;
    tYMin = -y0 / d;
    tYMax = (dim2.tripCount - 1 - y0) / d;
    int64_t tMin = max(tXMin, tYMin);
    uint64_t tMax = min(tXMax, tYMax);
    if (tMin > tMax) 
    {
        return NULL;
    }
    uint64_t start3 = dim1.start + x0 * dim1.step;
    uint64_t step3 = c * dim1.step;
    ArrayDimension* result = new(ArrayDimension){ start3, step3, tMax + 1 };
    return result;
}

/* dim1 / dim2 */
static vector<ArrayDimension> DimensionDifference(const ArrayDimension& dim1, const ArrayDimension& dim2)
{
    ArrayDimension* intersection = DimensionIntersection(dim1, dim2);
    if (!intersection) 
    {
        return {dim1};
    }
    vector<ArrayDimension> result;
    /* add the part before intersection */
    if (dim1.start < intersection->start)
    {
        result.push_back({ dim1.start, dim1.step, (intersection->start - dim1.start) / dim1.step });
    }
    /* add the parts between intersection steps */
    uint64_t start = (intersection->start - dim1.start) / dim1.step;
    uint64_t interValue = intersection->start;
    for (int64_t i = start; dim1.start + i * dim1.step <= intersection->start + intersection->step * (intersection->tripCount - 1); i++)
    {
        uint64_t centerValue = dim1.start + i * dim1.step;
        if (centerValue == interValue)
        {
            if (i - start  > 1)
            {
                result.push_back({ dim1.start + (start + 1) * dim1.step, dim1.step, i - start - 1 });
                start = i;
            }
            interValue += intersection->step;
        }
    }
    /* add the part after intersection */
    if (intersection->start + intersection->step * (intersection->tripCount - 1) < dim1.start + dim1.step * (dim1.tripCount - 1))
    {
        /* first value after intersection */
        uint64_t right_start = intersection->start + intersection->step * (intersection->tripCount - 1) + dim1.step;
        uint64_t tripCount = (dim1.start + dim1.step * dim1.tripCount - right_start) / dim1.step;
        result.push_back({right_start, dim1.step, tripCount});
    }
    delete(intersection);
    return result;
}


static vector<ArrayDimension> DimensionUnion(const ArrayDimension& dim1, const ArrayDimension& dim2)
{
    vector<ArrayDimension> res;
    ArrayDimension* inter = DimensionIntersection(dim1, dim2);
    if(!inter)
    {
        return { dim1, dim2 };
    }
    res.push_back(*inter);
    delete(inter);
    vector<ArrayDimension> diff1, diff2;
    diff1 = DimensionDifference(dim1, dim2);
    diff2 = DimensionDifference(dim2, dim1);
    res.insert(res.end(), diff1.begin(), diff1.end());
    res.insert(res.end(), diff2.begin(), diff2.end());
    return res;
}

static vector<ArrayDimension> ElementsIntersection(const vector<ArrayDimension>& firstElement, const vector<ArrayDimension>& secondElement)
{
    if(firstElement.empty() || secondElement.empty()) {
        return {};
    }
    size_t dimAmount = firstElement.size();
    /* check if there is no intersecction */
    for(size_t i = 0; i < dimAmount; i++) 
    {
        if(FindParticularSolution(firstElement[i], secondElement[i]).empty()){
            return {};
        }
    }
    vector<ArrayDimension> result(dimAmount);
    for(size_t i = 0; i < dimAmount; i++) 
    {
        ArrayDimension* resPtr = DimensionIntersection(firstElement[i], secondElement[i]);
        if(resPtr) 
        {
            result[i] = *resPtr;
        }
        else
        {
            return {};
        }
    }
    return result;
}

static vector<vector<ArrayDimension>> ElementsDifference(const vector<ArrayDimension>& firstElement, 
                                                         const vector<ArrayDimension>& secondElement)
{
    if(firstElement.empty() || secondElement.empty()) {
        return {};
    }
    vector<ArrayDimension> intersection = ElementsIntersection(firstElement, secondElement);
    vector<vector<ArrayDimension>> result;
    if(intersection.empty()) 
    {
        return {firstElement};
    }
    for(int i = 0; i < firstElement.size(); i++)
    {
        auto dimDiff = DimensionDifference(firstElement[i], secondElement[i]);
        if(!dimDiff.empty()) 
        {
            vector<ArrayDimension> firstCopy = firstElement;
            for(const auto& range: dimDiff)
            {
                firstCopy[i] = range;
                result.push_back(firstCopy);
            }
        }
    }
    return result;
}

static void ElementsUnion(const vector<ArrayDimension>& firstElement, const vector<ArrayDimension>& secondElement, 
                          vector<vector<ArrayDimension>>& lc, vector<vector<ArrayDimension>>& rc, 
                          vector<ArrayDimension>& intersection)
{
    /* lc(rc) is a set of ranges, which only exist in first(second) element*/
    intersection = ElementsIntersection(firstElement, secondElement);
    lc = ElementsDifference(firstElement, intersection);
    rc = ElementsDifference(secondElement, intersection);
}

void AccessingSet::FindUncovered(const vector<ArrayDimension>& element, vector<vector<ArrayDimension>>& result) const{
    vector<vector<ArrayDimension>> newTails;
    result.push_back(element);
    for(const auto& currentElement: allElements) 
    {
        for(const auto& tailLoc: result)
        {
            auto intersection = ElementsIntersection(tailLoc, currentElement);
            auto diff = ElementsDifference(tailLoc, intersection);
            if(!diff.empty()) {
                newTails.insert(newTails.end(), diff.begin(), diff.end());
            }
        }
        result = move(newTails);
    }
}

bool AccessingSet::ContainsElement(const vector<ArrayDimension>& element) const
{
    vector<vector<ArrayDimension>> tails;
    FindUncovered(element, tails);
    return !tails.empty();
}

void AccessingSet::FindCoveredBy(const vector<ArrayDimension>& element, vector<vector<ArrayDimension>>& result) const
{
    for(const auto& currentElement: allElements) 
    {
        auto intersection = ElementsIntersection(element, currentElement);
        if(!intersection.empty()) {
            result.push_back(intersection);
        }
    }
}

vector<vector<ArrayDimension>> AccessingSet::GetElements() const
{
    return allElements;
}

void AccessingSet::Insert(const vector<ArrayDimension>& element)
{
    vector<vector<ArrayDimension>> tails;
    FindUncovered(element, tails);
    allElements.insert(allElements.end(), tails.begin(), tails.end());
}

AccessingSet AccessingSet::Union(const AccessingSet& source) {
    AccessingSet result;
    for(auto& element: source.GetElements()) {
        result.Insert(element);
    }
    return result;
}

AccessingSet AccessingSet::Intersect(const AccessingSet& secondSet) const
{
    vector<vector<ArrayDimension>> result;
    for(const auto& element: allElements) 
    {
        if(secondSet.ContainsElement(element)) 
        {
            result.push_back(element);
        }
        else
        {
            vector<vector<ArrayDimension>> coveredBy;
            secondSet.FindCoveredBy(element, coveredBy);
            if(!coveredBy.empty()) 
            {
                result.insert(result.end(), coveredBy.begin(), coveredBy.end());
            }
        }
    }
    return AccessingSet(result);
}

AccessingSet AccessingSet::Diff(const AccessingSet& secondSet) const
{
    AccessingSet intersection = this->Intersect(secondSet);
    AccessingSet uncovered = *this;
    vector<vector<ArrayDimension>> result;
    for (const auto& element : intersection.GetElements())
    {
        vector<vector<ArrayDimension>> current_uncovered;
        uncovered.FindUncovered(element, current_uncovered);
        uncovered = AccessingSet(current_uncovered);
    }
    return uncovered;
}

bool operator==(const ArrayDimension& lhs, const ArrayDimension& rhs)
{
    return lhs.start == rhs.start && lhs.step == rhs.step && lhs.tripCount == rhs.tripCount;
}


bool operator==(const AccessingSet& lhs, const AccessingSet& rhs) 
{
    for (size_t i = 0; i < lhs.allElements.size(); i++) 
    {
        for (size_t j = 0; j < lhs.allElements[i].size(); j++) 
        {
            if (lhs.allElements[i][j] != rhs.allElements[i][j])
            {
                return false;
            }
        }
    }
    return true;
}

void Collapse(Region* region)
{
    //Region* newBlock = new Region();  
    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);
        }
    }

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

static void SetConnections(unordered_map<SAPFOR::BasicBlock*, Region*>& bbToRegion, const unordered_set<SAPFOR::BasicBlock*>& blockSet)
{
    for (SAPFOR::BasicBlock* block : blockSet)
    {
        for (SAPFOR::BasicBlock* nextBlock : block->getNext())
        {
            if (bbToRegion.find(nextBlock) != bbToRegion.end())
            {
                bbToRegion[block]->addNextRegion(bbToRegion[nextBlock]);
            }
        }
        for (SAPFOR::BasicBlock* prevBlock : block->getPrev())
        {
            if (bbToRegion.find(prevBlock) != bbToRegion.end())
            {
                bbToRegion[block]->addPrevRegion(bbToRegion[prevBlock]);
            }
        }
    }
}

static Region* CreateSubRegion(LoopGraph* loop, const vector<SAPFOR::BasicBlock*>& Blocks, const unordered_map<SAPFOR::BasicBlock*, Region*>& bbToRegion)
{
    Region* region = new Region;
    auto [header, blockSet] = GetBasicBlocksForLoop(loop, Blocks);
    for (SAPFOR::BasicBlock* block : Blocks)
    {
        region->addBasickBlocks(bbToRegion.at(block));
    }
    for (LoopGraph* childLoop : loop->children)
    {
        region->addSubRegions(CreateSubRegion(childLoop, Blocks, bbToRegion));
    }
    return region;
}

Region::Region(LoopGraph* loop, vector<SAPFOR::BasicBlock*>& Blocks)
{
    auto [header, blockSet] = GetBasicBlocksForLoop(loop, Blocks);
    unordered_map<SAPFOR::BasicBlock*, Region*> bbToRegion;
    for (auto poiner : blockSet)
    {
        bbToRegion[poiner] = new Region(*poiner);
    }
    this->header = bbToRegion[header];
    SetConnections(bbToRegion, blockSet);
    //create subRegions
    for (LoopGraph* childLoop : loop->children)
    {
        subRegions.insert(CreateSubRegion(childLoop, Blocks, bbToRegion));
    }
}

void SolveDataFlowIteratively(Region* DFG) 
{
    unordered_set<Region*> worklist(DFG->getBasickBlocks());
    do
    {
        Region* b = *worklist.begin();
        ArrayAccessingIndexes newIn;
        for (Region* prevBlock : b->getPrevRegions())
        {
            for (const auto& [arrayName, accessSet] : prevBlock->array_out)
            {
                newIn[arrayName] = newIn[arrayName].Intersect(accessSet);
            }
        }
        b->array_in = newIn;
        ArrayAccessingIndexes newOut;
        for (auto& [arrayName, accessSet] : b->array_in)
        {
            newOut[arrayName] = b->array_in[arrayName].Union(b->array_def[arrayName]);
        }
        /* can not differ */
        if (newOut != b->array_out) 
        {
            b->array_out = newOut;
        }
        else
        {
            worklist.erase(b);
        }
    } 
    while (!worklist.empty());
}

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

void FindPrivateArrays(map<string, vector<LoopGraph*>> &loopGraph, map<FuncInfo*, vector<SAPFOR::BasicBlock*>>& FullIR)
{  
    for (const auto& curr_graph_pair: loopGraph)
    {
        for (const auto& curr_loop : curr_graph_pair.second)
        {
            auto block_loop = GetBasicBlocksForLoop(curr_loop, (*FullIR.begin()).second);
            for (const auto& bb : block_loop.second) {
                ArrayAccessingIndexes def, use;
                //GetDefUseArray(bb, curr_loop, def, use);
            }
            ArrayAccessingIndexes loopDimensionsInfo;
            //GetDimensionInfo(curr_loop, loopDimensionsInfo, 0);
            //print_info(curr_loop);
        }
    }

}

void GetDimensionInfo(LoopGraph* loop, map<DIST::Array*, vector<vector<ArrayDimension>>>& loopDimensionsInfo, int level)
{
    cout << "line_num: " << loop->lineNum << endl;
    for (const auto& writeOpPairs : loop->writeOpsForLoop)
    {
        vector<vector<ArrayDimension>> arrayDimensions(writeOpPairs.first->GetDimSize());
        loopDimensionsInfo[writeOpPairs.first] = arrayDimensions;
        for (const auto& writeOp : writeOpPairs.second)
        {
            for (const auto& coeficient_pair : writeOp.coefficients)
            {
                uint64_t start, step, tripCount;
                start = loop->startVal * coeficient_pair.first.first + coeficient_pair.first.second;
                step = loop->stepVal * coeficient_pair.first.first;
                tripCount = (loop->endVal - coeficient_pair.first.second) / step;
                if (start <= loop->endVal)
                {
                    loopDimensionsInfo[writeOpPairs.first][level].push_back({start, step, tripCount});
                    cout << "level: " << level << endl;
                    cout << "start: " << start << endl;
                    cout << "step: " << step << endl;
                    cout << "trip_count: " << tripCount << endl;
                    cout << endl;
                }
                
                
            }
        }
    }
    cout << "line_num_after: " << loop->lineNumAfterLoop << endl;
    if (!loop->children.empty())
    {
        for (const auto& childLoop : loop->children)
        {
            GetDimensionInfo(childLoop, loopDimensionsInfo, level+1);
        }
    }
}