#include "leak_detector.h"

#include <limits>
#include <map>
#include <vector>
#include <string>
#include <tuple>

#include "dvm.h"
#include "PredictSchemeWithLibrary.h"
#include "../../projects/libpredictor/include/libpredict/predictor.h"
#include "../DirectiveProcessing/directive_parser.h"
#include "../Distribution/DvmhDirective.h"
#include "../ParallelizationRegions/ParRegions.h"
#include "../GraphLoop/graph_loops_func.h"
#include "../Utils/errors.h"
#include "../Utils/utils.h"

using std::map;
using std::string;
using std::vector;
using std::pair;
using std::tuple;

// МОЖЕТ КАК-ТО ВЫЧИСЛЯТЬ ДИРЕКТИВЫ, А ПОТОМ ДЕЛАТЬ ПРОГОНЫ?
double runLibpredictCalc(SgProject &project,
                      vector<size_t> topology,
                      string clusterConfStr,
                      vector<ParallelRegion*> &parallelRegions,
                      map<string, vector<LoopGraph*>> loopGraph,
                      map<string, vector<Messages>> &SPF_messages)
{
    libpredict::RetInit retInit = libpredict::Init(clusterConfStr, topology[0], topology[1], topology[2], topology[3]);

    if (retInit != libpredict::INIT_SUCCESS) {
        __spf_print(1, "ERROR: Failed to initialize libpredict with cluster config: %s, return code: %d\n", clusterConfStr.c_str(), (int)retInit);
    
        std::wstring messageR, messageE;
        __spf_printToLongBuf(messageE, L"Failed to initialize libpredict library with cluster config: %s, return code: %d",
                             to_wstring(clusterConfStr).c_str(), (int)retInit);
        __spf_printToLongBuf(messageR, R206);
        getObjectForFileFromMap(clusterConfStr.c_str(), SPF_messages).push_back(Messages(ERROR, 1, messageR, messageE, 1063));
        return -1;
    }

    // distribute и align из parallelRegions
    for (int z = 0; z < parallelRegions.size(); ++z) {
        const DataDirective &dataDirectives = parallelRegions[z]->GetDataDir();
        const vector<int> &currentVariant = parallelRegions[z]->GetCurrentVariant();
        DIST::Arrays<int> &allArrays = parallelRegions[z]->GetAllArraysToModify();

        auto &tmp = dataDirectives.distrRules;
        vector<pair<DIST::Array*, const DistrVariant*>> currentVar;
        for (int z1 = 0; z1 < currentVariant.size(); ++z1)
            currentVar.push_back(std::make_pair(tmp[z1].first, &tmp[z1].second[currentVariant[z1]]));

        // distribute
        for (const auto& distrRule : currentVar) {
            DIST::Array* array = distrRule.first;
            const DistrVariant* variant = distrRule.second;

            if (array && variant && !array->IsNotDistribute()) {
                size_t arrayId = array->GetId();
                size_t elemSize = array->GetTypeSize();

                vector<libpredict::DistributeAxisRule> axisDistributions;
                const auto& arraySizes = array->GetSizes();
                for (int dim = 0; dim < array->GetDimSize(); ++dim) {
                    size_t dimSize = arraySizes[dim].second - arraySizes[dim].first + 1;

                    if (dim < variant->distRule.size() && variant->distRule[dim] == dist::BLOCK) {
                        axisDistributions.emplace_back(dimSize, libpredict::TypeDistribute::BLOCK);
                    } else {
                        axisDistributions.emplace_back(dimSize, libpredict::TypeDistribute::NONE);
                    }
                }

                vector<pair<size_t, size_t>> shadowEdges;
                const auto& shadowSpec = array->GetShadowSpec();
                for (int dim = 0; dim < shadowSpec.size() && dim < array->GetDimSize(); ++dim) {
                    if (dim < variant->distRule.size() && variant->distRule[dim] == dist::BLOCK) {
                        shadowEdges.emplace_back(shadowSpec[dim].first, shadowSpec[dim].second);
                    }
                }

                libpredict::RetDistribute retDistribute = libpredict::Distribute(arrayId, elemSize, axisDistributions, shadowEdges);

                if (retDistribute != libpredict::DISTRIBUTE_SUCCESS) {
                    __spf_print(1, "ERROR: Failed to distribute array '%s' (id=%zu) with libpredict, return code: %d\n",
                               array->GetShortName().c_str(), arrayId, (int)retDistribute);

                    std::wstring messageR, messageE;
                    __spf_printToLongBuf(messageE, L"Failed to distribute array '%s' with libpredict, return code: %d",
                                         to_wstring(array->GetShortName()).c_str(), (int)retDistribute);
                    __spf_printToLongBuf(messageR, R207);
                    getObjectForFileFromMap(array->GetDeclInfo().begin()->first.c_str(), SPF_messages).push_back(
                        Messages(ERROR, array->GetDeclInfo().begin()->second, messageR, messageE, 1064));
                }
            }
        }

        // align
        for (const auto& alignRule : dataDirectives.alignRules) {
            DIST::Array* alignArray = alignRule.alignArray;
            DIST::Array* alignWithArray = alignRule.alignWith;

            if (alignArray && alignWithArray && !alignArray->IsNotDistribute()) {
                size_t arrayId = alignArray->GetId();
                size_t distributedArrayId = alignWithArray->GetId();
                size_t elemSize = alignArray->GetTypeSize();

                const auto& arraySizes = alignArray->GetSizes();
                vector<size_t> dimensions;
                for (int dim = 0; dim < alignArray->GetDimSize(); ++dim)
                    dimensions.push_back(arraySizes[dim].second - arraySizes[dim].first + 1);

                vector<libpredict::AlignDisplay> distributionExpressions;
                for (int dim = 0; dim < alignWithArray->GetDimSize(); ++dim) {
                    bool found = false;
                    for (int i = 0; i < alignRule.alignRuleWith.size(); ++i) {
                        const auto& ruleWith = alignRule.alignRuleWith[i];
                        if (ruleWith.first == dim) {
                            const auto& rule = ruleWith.second;
                            if (rule.first == 0) {
                                // Константа
                                distributionExpressions.emplace_back(rule.second);
                            } else {
                                // Линейное выражение a * I + b
                                distributionExpressions.emplace_back(i, rule.first, rule.second);
                            }
                            found = true;
                            break;
                        }
                    }
                    if (!found) {
                        // Нет правила для этого измерения
                        distributionExpressions.emplace_back();
                    }
                }

                vector<pair<size_t, size_t>> shadowEdges;
                const auto& shadowSpec = alignArray->GetShadowSpec();
                for (int dim = 0; dim < shadowSpec.size() && dim < alignArray->GetDimSize(); ++dim) {
                    shadowEdges.emplace_back(shadowSpec[dim].first, shadowSpec[dim].second);
                }

                libpredict::RetAlign retAlign = libpredict::Align(arrayId, distributedArrayId, elemSize, dimensions, distributionExpressions, shadowEdges);
                
                if (retAlign != libpredict::ALIGN_SUCCESS) {
                    __spf_print(1, "ERROR: Failed to align array '%s' (id=%zu) with array '%s' (id=%zu), return code: %d\n",
                               alignArray->GetShortName().c_str(), arrayId,
                               alignWithArray->GetShortName().c_str(), distributedArrayId, (int)retAlign);
                    
                    std::wstring messageR, messageE;
                    __spf_printToLongBuf(messageE, L"Failed to align array '%s' with array '%s' using libpredict, return code: %d",
                                         to_wstring(alignArray->GetShortName()).c_str(),
                                         to_wstring(alignWithArray->GetShortName()).c_str(), (int)retAlign);
                    __spf_printToLongBuf(messageR, R208);
                    getObjectForFileFromMap(alignArray->GetDeclInfo().begin()->first.c_str(), SPF_messages).push_back(
                        Messages(ERROR, alignArray->GetDeclInfo().begin()->second, messageR, messageE, 1065));
                }
            }
        }

        // shadow_renew
        map<LoopGraph*, ParallelDirective*> parallelDirs;
        for (int i = project.numberOfFiles() - 1; i >= 0; --i) {
            SgFile *file = &(project.file(i));
            auto fountInfo = findAllDirectives(file, getObjectForFileFromMap(file->filename(), loopGraph), parallelRegions[z]->GetId());
            parallelDirs.insert(fountInfo.begin(), fountInfo.end());
        }

        for (auto& dirPair : parallelDirs) {
            LoopGraph* loopPtr = dirPair.first;
            ParallelDirective* directive = dirPair.second;

            if (directive && !directive->shadowRenew.empty()) {
                for (size_t shadowIdx = 0; shadowIdx < directive->shadowRenew.size(); ++shadowIdx) {
                    const auto& shadowRenewItem = directive->shadowRenew[shadowIdx];
                    const string& arrayName = shadowRenewItem.first.second; // uniqName
                    const vector<pair<int, int>>& bounds = shadowRenewItem.second;

                    DIST::Array* shadowArray = allArrays.GetArrayByName(arrayName);
                    if (shadowArray == NULL)
                        continue;

                    if (shadowArray && !shadowArray->IsNotDistribute()) {
                        size_t arrayId = shadowArray->GetId();

                        vector<pair<size_t, size_t>> shadow_renew;
                        for (const auto& bound : bounds) {
                            shadow_renew.emplace_back(static_cast<size_t>(bound.first),
                                                    static_cast<size_t>(bound.second));
                        }

                        bool corner = directive->shadowRenewCorner[shadowIdx];

                        size_t number_loop_iterations = loopPtr ? static_cast<size_t>(loopPtr->countOfIters) : 1;

                        libpredict::RetShadowRenew retShadowRenew = libpredict::ShadowRenew(arrayId, shadow_renew, corner, number_loop_iterations);

                        if (retShadowRenew != libpredict::SHADOW_RENEW_SUCCESS) {
                            __spf_print(1, "ERROR: Failed to process shadow_renew for array '%s' (id=%zu), return code: %d\n",
                                       shadowArray->GetShortName().c_str(), arrayId, (int)retShadowRenew);
                            
                            std::wstring messageR, messageE;
                            __spf_printToLongBuf(messageE, L"Failed to process shadow_renew for array '%s' with libpredict, return code: %d",
                                                 to_wstring(shadowArray->GetShortName()).c_str(), (int)retShadowRenew);
                            __spf_printToLongBuf(messageR, R209);
                            getObjectForFileFromMap(shadowArray->GetDeclInfo().begin()->first.c_str(), SPF_messages).push_back(
                                Messages(ERROR, shadowArray->GetDeclInfo().begin()->second, messageR, messageE, 1066));
                        }
                    }
                }
            }
        }
    }

    return libpredict::GetTime();
}

void runPredictScheme(SgProject &project,
                      vector<vector<size_t>> &topologies, // такой способ передачи разве хочу?
                      vector<ParallelRegion*> &parallelRegions,
                      map<string, vector<LoopGraph*>> loopGraph,
                      map<string, vector<Messages>> &SPF_messages)
{
    // calculating maximum dimension of distribution
    int maxSizeDist = 0;
    for (int z = 0; z < parallelRegions.size(); ++z) {
        const DataDirective &dataDirectives = parallelRegions[z]->GetDataDir();
        const vector<int> &currentVariant = parallelRegions[z]->GetCurrentVariant();

        auto &tmp = dataDirectives.distrRules;
        vector<const DistrVariant*> currentVar;
        for (int z1 = 0; z1 < currentVariant.size(); ++z1)
            currentVar.push_back(&tmp[z1].second[currentVariant[z1]]);

        for (auto var : currentVar) {
            int countBlock = 0;
            for (int z = 0; z < var->distRule.size(); ++z)
                if (var->distRule[z] == dist::BLOCK)
                    ++countBlock;
            maxSizeDist = countBlock;
        }
    }

    // calculating name of a cluster configuration file
    string clusterConfStr;
    if (project.numberOfFiles() > 0) {
        string firstFilePath = project.fileName(0);

        size_t lastSlash = firstFilePath.find_last_of("/\\");
        clusterConfStr = firstFilePath.substr(0, lastSlash + 1) + "cluster.conf";
    }

    // iterating through topologies to find most optimal one
    topologies = vector<vector<size_t>>();
    if (maxSizeDist) {

        if (maxSizeDist > 4) maxSizeDist = 4;

        // TODO: look at cluster configuration
        size_t n1max = 10;
        size_t n2max = (maxSizeDist >= 2) ? 10 : 1;
        size_t n3max = (maxSizeDist >= 3) ? 10 : 1;
        size_t n4max = (maxSizeDist >= 4) ? 10 : 1;

        for (size_t n1 = 1; n1 <= n1max; ++n1) {
            for (size_t n2 = 1; n2 <= n2max; ++n2) {
                for (size_t n3 = 1; n3 <= n3max; ++n3) {
                    for (size_t n4 = 1; n4 <= n4max; ++n4) {
                        topologies.push_back(vector<size_t>{n1, n2, n3, n4});
                    }
                }
            }
        }

        vector<size_t> best;
        double bestTime = std::numeric_limits<double>::max();
        for (auto &topology : topologies) {
            // if (DEBUG) {
                // string outStr = "";
                // for (const auto &elem : top)
                //     outStr += std::to_string(elem) + " ";
                // __spf_print(1, "topology %s has time %f\n", outStr.c_str(), currTime);
            // }
            double currTime = runLibpredictCalc(project, topology, clusterConfStr, parallelRegions, loopGraph, SPF_messages);

            if (currTime == -1)
                return;
            
            if (currTime < bestTime) {
                bestTime = currTime;
                best = topology;
            }
        }
        string outStr;
        for (const auto &elem : best)
            outStr += std::to_string(elem) + " ";

        __spf_print(1, "best topology %s with time %f\n", outStr.c_str(), bestTime);
    } else {
        __spf_print(1, "impossible to calculate best topology: project does not contain distribution directives\n");
    }
}