#pragma once

#include <string>
#include <vector>
#include <algorithm>
#include <set>
#include <map>

#include "../Distribution/DvmhDirective.h"
#include "../Distribution/GraphCSR.h"
#include "../Distribution/Distribution.h"
#include "AstWrapper.h"

#include "json.hpp"

#if __SPF
#include "SgUtils.h"
#endif

struct ParallelRegionLines
{
    ParallelRegionLines(double weight = 1.0) : weight(weight)
    {
        lines = std::make_pair(-1, -1);
        stats = std::make_pair<Statement*, Statement*>(NULL, NULL);
        intervalBefore = std::make_pair<Statement*, Statement*>(NULL, NULL);
        intervalAfter = std::make_pair<Statement*, Statement*>(NULL, NULL);
    }

    ParallelRegionLines(const std::pair<int, int> &lines, double weight = 1.0) : lines(lines), weight(weight)
    {
        stats = std::make_pair<Statement*, Statement*>(NULL, NULL);
        intervalBefore = std::make_pair<Statement*, Statement*>(NULL, NULL);
        intervalAfter = std::make_pair<Statement*, Statement*>(NULL, NULL);
    }

    ParallelRegionLines(const std::pair<int, int> &lines, const std::pair<Statement*, Statement*> stats, double weight = 1.0)
        : lines(lines), stats(stats), weight(weight)
    {
        intervalBefore = std::make_pair<Statement*, Statement*>(NULL, NULL);
        intervalAfter = std::make_pair<Statement*, Statement*>(NULL, NULL);
    }

    bool operator==(const ParallelRegionLines &regionLines) const { return lines == regionLines.lines && stats == regionLines.stats; }
    bool operator<(const ParallelRegionLines &otherLines) const { return lines.first < otherLines.lines.first; }

    void print(FILE *fileOut)
    {
        fprintf(fileOut, " [%d -- %d]", lines.first, lines.second);
        if (stats.first && stats.second)
            fprintf(fileOut, " explicit\n");
        else
            fprintf(fileOut, "\n");
    }

    bool isImplicit() const { return stats.first == NULL || stats.second == NULL; }

    // <start, end> lines
    std::pair<int, int> lines;
    // <start, end> stats
    std::pair<Statement*, Statement*> stats;

    // <start, end> interval
    std::pair<Statement*, Statement*> intervalBefore;
    std::pair<Statement*, Statement*> intervalAfter;

    double weight; // weight of the fragment among all fragments of this region
};

#if __SPF
struct ParallelRegionArray
{
private:
    std::string name;
    std::string fileName;
    SgSymbol *origSymbol;
    SgSymbol *copySymbol;
    std::vector<SgStatement*> declStatements;
    std::vector<ParallelRegionLines> allLines;
public:
    explicit ParallelRegionArray(const std::string &name, const std::string &fileName, SgSymbol *origSymbol, SgSymbol *copySymbol,
                                 const ParallelRegionLines &lines, std::vector<SgStatement*> &declStatements) :
                                 name(name), fileName(fileName), origSymbol(origSymbol), copySymbol(copySymbol), declStatements(declStatements)
    {
        allLines.push_back(lines);
    }

    const std::string& getName() const { return name; }
    const std::string& getFileName() const { return fileName; }
    SgSymbol* getOrigSymbol() const { return origSymbol; }
    SgSymbol* getCopySymbol() const { return copySymbol; }
    const std::vector<SgStatement*>& getDeclStatements() const { return declStatements; }
    const std::vector<ParallelRegionLines>& getAllLines() const { return allLines; }

    void addLines(const ParallelRegionLines &newLines)
    {
        for (auto &lines : allLines)
            if (lines == newLines)
                return;

        allLines.push_back(newLines);
    }

    void setCopySymbol(SgSymbol *copySymbol) { this->copySymbol = copySymbol; }
};
#endif

struct ParallelRegion
{
public:
    ParallelRegion(const uint64_t regionId, const std::string &originalName) : regionId(regionId), originalName(originalName) { }

    ParallelRegion(const ParallelRegion &copy) : allArrays(copy.allArrays), G(copy.G), reducedG(copy.reducedG), dataDirectives(copy.dataDirectives)
    {
        regionId = copy.regionId;
        originalName = copy.originalName;
        lines = copy.lines;
        functionsCall = copy.functionsCall;
        currentVariant = copy.currentVariant;
    }

    int AddLines(const std::pair<int, int> &linesToAdd, const std::string &file, const std::pair<Statement*, Statement*> *startEnd = NULL, double weight = 1.0)
    {
        if (linesToAdd.first > linesToAdd.second)
            return -1;

        auto it = lines.find(file);
        if (it == lines.end())
            it = lines.insert(it, make_pair(file, std::vector<ParallelRegionLines>()));

        if (startEnd)
            it->second.push_back(ParallelRegionLines(linesToAdd, *startEnd, weight));
        else
            it->second.push_back(ParallelRegionLines(linesToAdd, weight));

        return 0;
    }

    void AddFuncCalls(const std::string &func, const std::string &file, const int line)
    {
        auto *found_lines = GetLinesByLine(file, line);

        if (found_lines)
            functionsCall[func].insert(found_lines);
    }

#if __SPF
    void AddFuncCallsToAllCalls(FuncInfo *func, const std::string &file, const int line)
    {
        auto *found_lines = GetLinesByLine(file, line);

        if (found_lines)
            allFunctionsCall[func].insert(found_lines);
    }
#endif

    uint64_t GetId() const { return regionId; }
    const std::string& GetName() const { return originalName; }
    const std::map<std::string, std::vector<ParallelRegionLines>>& GetAllLines() const { return lines; }
    std::map<std::string, std::vector<ParallelRegionLines>>& GetAllLinesToModify() { return lines; }
    const std::vector<ParallelRegionLines>* GetLines(const std::string &file) const 
    {
        auto it = lines.find(file);
        if (it == lines.end())
            return NULL;
        else
            return &(it->second);
    }
    const ParallelRegionLines* GetLinesByLine(const std::string &file, const int line) const
    {
        auto fileLines = GetLines(file);
        if (fileLines)
            for (auto &lines : *fileLines)
                if (lines.lines.first <= line && line <= lines.lines.second)
                    return &lines;
        return NULL;
    }

    const DIST::GraphCSR<int, double, attrType>& GetGraph() const { return G; }
    DIST::GraphCSR<int, double, attrType>& GetGraphToModify() { return G; }

    const DIST::GraphCSR<int, double, attrType>& GetReducedGraph() const { return reducedG; }
    DIST::GraphCSR<int, double, attrType>& GetReducedGraphToModify() { return reducedG; }

    const DIST::Arrays<int>& GetAllArrays() const { return allArrays; }
    DIST::Arrays<int>& GetAllArraysToModify() { return allArrays; }

    void SetCurrentVariant(const std::vector<int> &newVariant) { currentVariant = newVariant; }
    const std::vector<int>& GetCurrentVariant() const { return currentVariant; }

    const DataDirective& GetDataDir() const { return dataDirectives; }
    DataDirective& GetDataDirToModify() { return dataDirectives; }

    const std::map<std::string, std::set<const ParallelRegionLines*>>& GetFuncCalls() const { return functionsCall; }

#if __SPF
    const std::map<FuncInfo*, std::set<const ParallelRegionLines*>>& GetAllFuncCalls() const { return allFunctionsCall; }
    const std::map<FuncInfo*, std::map<DIST::Array*, std::vector<ParallelRegionLines>>>& GetUsedLocalArrays() const { return usedLocalArrays; }
    const std::map<FuncInfo*, std::map<DIST::Array*, std::vector<ParallelRegionLines>>>& GetUsedCommonArrays() const { return usedCommonArrays; }

    void AddUsedLocalArray(FuncInfo *func, DIST::Array *array, const ParallelRegionLines &lines)
    {
        auto it = usedLocalArrays.find(func);
        if (it == usedLocalArrays.end())
            it = usedLocalArrays.insert(it, std::make_pair(func, std::map<DIST::Array*, std::vector<ParallelRegionLines>>()));

        auto itt = it->second.find(array);
        if (itt == it->second.end())
            itt = it->second.insert(itt, std::make_pair(array, std::vector<ParallelRegionLines>()));

        for (auto &curLines : itt->second)
            if (curLines == lines)
                return;

        itt->second.push_back(lines);
    }

    void AddUsedCommonArray(FuncInfo *func, DIST::Array *array, const ParallelRegionLines &lines)
    {
        auto it = usedCommonArrays.find(func);
        if (it == usedCommonArrays.end())
            it = usedCommonArrays.insert(it, std::make_pair(func, std::map<DIST::Array*, std::vector<ParallelRegionLines>>()));

        auto itt = it->second.find(array);
        if (itt == it->second.end())
            itt = it->second.insert(itt, std::make_pair(array, std::vector<ParallelRegionLines>()));

        for (auto &curLines : itt->second)
            if (curLines == lines)
                return;

        itt->second.push_back(lines);
    }
#endif

    bool HasThisLine(const int line, const std::string &file, const ParallelRegionLines** found = nullptr) const
    {
        bool retVal = false;
        auto it = lines.find(file);
        if (it != lines.end())
        {
            for (int i = 0; i < it->second.size(); ++i)
            {
                if (it->second[i].lines.first <= line && it->second[i].lines.second >= line)
                {
                    retVal = true;
                    if (found)
                        *found = &(it->second[i]);

                    break;
                }
            }
        }

        return retVal;
    }

    nlohmann::json toJson()
    {
        nlohmann::json retVal;

        retVal["packedRegionId"] = std::to_string(regionId);
        retVal["originalName"] = originalName;

        nlohmann::json arrays = nlohmann::json::array();
        for (auto& array : allArrays.GetArrays())
            arrays.push_back(array->toJson());
        retVal["packedArrays"] = arrays;

        nlohmann::json linesInfo = nlohmann::json::array();
        for (auto& [file, linesByFile] : lines)
        {
            nlohmann::json linesRegs;
            nlohmann::json lines = nlohmann::json::array();

            for (auto& elem : linesByFile)
            {
                JSON pair;
                pair["key"] = elem.lines.first;
                pair["value"] = elem.lines.second;
                lines.push_back(pair);
            }
            linesRegs["file"] = file;
            linesRegs["lines"] = lines;

            linesInfo.push_back(linesRegs);
        }
        retVal["regionsLines"] = linesInfo;

        nlohmann::json alignRules = nlohmann::json::array();
        for (auto& rule : dataDirectives.alignRules)
            alignRules.push_back(rule.toJson());
        retVal["alignRules"] = alignRules;

        return retVal;
    }

    void CleanData()
    {
        reducedG.ClearGraphCSR();
        currentVariant.clear();

        dataDirectives.distrRules.clear();
        dataDirectives.alignRules.clear();
    }

    void print(FILE *fileOut)
    {
        fprintf(fileOut, "  regionId %lld\n", regionId);
        fprintf(fileOut, "  originalName '%s'\n", originalName.c_str());
        fprintf(fileOut, "  functions call from %d:\n", (int)functionsCall.size());
        for (auto &func : functionsCall)
            fprintf(fileOut, "    '%s'\n", func.first.c_str());
        fprintf(fileOut, "  total lines %d:\n", (int)lines.size());
        for (auto &line : lines)
        {
            fprintf(fileOut, "    in file '%s':\n", line.first.c_str());
            for (auto &elem : line.second)
            {
                fprintf(fileOut, "     ");
                elem.print(fileOut);
            }
        }
    }

    void AddUserDirectives(const std::vector<Statement*> &dirs, const int type)
    {
        if (dirs.size() == 0)
            return;

        if (type == DVM_DISTRIBUTE_DIR || type == DVM_VAR_DECL)
            userDvmDistrDirs.insert(userDvmDistrDirs.end(), dirs.begin(), dirs.end());
        else if (type == DVM_ALIGN_DIR)
            userDvmAlignDirs.insert(userDvmAlignDirs.end(), dirs.begin(), dirs.end());
        else if (type == DVM_SHADOW_DIR)
            userDvmShadowDirs.insert(userDvmShadowDirs.end(), dirs.begin(), dirs.end());
        else if (type == DVM_REALIGN_DIR)
            userDvmRealignDirs.insert(userDvmRealignDirs.end(), dirs.begin(), dirs.end());
        else if (type == DVM_REDISTRIBUTE_DIR)
            userDvmRedistrDirs.insert(userDvmRedistrDirs.end(), dirs.begin(), dirs.end());
    }

    const std::vector<Statement*>* GetUsersDirecites(const int type) const
    {
        if (type == DVM_DISTRIBUTE_DIR || type == DVM_VAR_DECL)
            return &userDvmDistrDirs;
        else if (type == DVM_ALIGN_DIR)
            return &userDvmAlignDirs;
        else if (type == DVM_SHADOW_DIR)
            return &userDvmShadowDirs;
        else if (type == DVM_REALIGN_DIR)
            return &userDvmRealignDirs;
        else if (type == DVM_REDISTRIBUTE_DIR)
            return &userDvmRedistrDirs;
        else
            return NULL;
    }

    bool HasUserDvmDirs() const 
    {
        return userDvmDistrDirs.size()   != 0 ||
               userDvmAlignDirs.size()   != 0 ||
               userDvmShadowDirs.size()  != 0 ||
               userDvmRealignDirs.size() != 0 ||
               userDvmRedistrDirs.size() != 0;
    }
#if __SPF
    void ClearUserDirs()
    {
        ClearVector(userDvmDistrDirs);
        ClearVector(userDvmAlignDirs);
        ClearVector(userDvmShadowDirs);
        ClearVector(userDvmRealignDirs);
        ClearVector(userDvmRedistrDirs);
    }
#endif
private:
    uint64_t regionId;
    //name in program
    std::string originalName;
    // file -> lines info
    std::map<std::string, std::vector<ParallelRegionLines>> lines;
    std::map<std::string, std::set<const ParallelRegionLines*>> functionsCall; // func name -> fragments with calls

#if __SPF
    // for RESOLVE_PAR_REGIONS
    std::map<FuncInfo*, std::set<const ParallelRegionLines*>> allFunctionsCall; // function -> fragments with calls
    std::map<FuncInfo*, std::map<DIST::Array*, std::vector<ParallelRegionLines>>> usedLocalArrays;  // func -> array -> lines
    std::map<FuncInfo*, std::map<DIST::Array*, std::vector<ParallelRegionLines>>> usedCommonArrays; // func -> array -> lines
    //
#endif

    // for LOOP_ANALYZER_DATA_DIST
    DIST::GraphCSR<int, double, attrType> G;
    DIST::Arrays<int> allArrays;

    DIST::GraphCSR<int, double, attrType> reducedG;
    //

    //for directive creating
    DataDirective dataDirectives;
    std::vector<int> currentVariant;
    //

    std::vector<Statement*> userDvmDistrDirs;
    std::vector<Statement*> userDvmAlignDirs;
    std::vector<Statement*> userDvmShadowDirs;
    std::vector<Statement*> userDvmRealignDirs;
    std::vector<Statement*> userDvmRedistrDirs;

#if __SPF
    void ClearVector(std::vector<Statement*> &toRem)
    {
        for (auto& elem : toRem)
        {
            if (SgFile::switchToFile(elem->fileName()))
                ;// printInternalError(convertFileName(__FILE__).c_str(), __LINE__);
            else
                elem->deleteStmt();
        }
        toRem.clear();
    }
#endif
};

ParallelRegion* getRegionById(const std::vector<ParallelRegion*>& regions, const uint64_t regionId);
ParallelRegion* getRegionByName(const std::vector<ParallelRegion*>& regions, const std::string& regionName);
ParallelRegion* getRegionByLine(const std::vector<ParallelRegion*>& regions, const std::string& file, const int line);
std::pair<ParallelRegion*, const ParallelRegionLines*> getRegionAndLinesByLine(const std::vector<ParallelRegion*>& regions, const std::string& file, const int line);
std::set<ParallelRegion*> getAllRegionsByLine(const std::vector<ParallelRegion*>& regions, const std::string& file, const int line);