//
//  ModuleBasic.cpp
//  MNN
//
//  Created by MNN on 2021/10/15.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#include "MNN_generated.h"
#include <MNN/expr/Expr.hpp>
#include <MNN/expr/ExecutorScope.hpp>
#include <MNN/expr/Module.hpp>
#include <MNN/expr/ExprCreator.hpp>
#define MNN_OPEN_TIME_TRACE
#include <MNN/AutoTime.hpp>
#include "rapidjson/document.h"
#include "core/MemoryFormater.h"
#include <numeric>
#include <chrono>
#include <iostream>
#include <thread>
#include "ExprDebug.hpp"
#include "core/KVMeta.hpp"

using namespace MNN::Express;
using namespace MNN;
static bool compareOutput(VARP output, const std::string& directName, const std::string& name, Dimensionformat dataFormat, int order) {

    auto info = output->getInfo();
    auto ptr = output->readMap<float>();
    if (info && info->size <= 0) {
        MNN_PRINT("skip checking value for zero content tensor %s\n", name.c_str());
        return true;
    }

    if (nullptr == info || nullptr == ptr) {
        MNN_ERROR("TESTERROR name:%s, info:%p, ptr:%p. size:%zu\n", name.c_str(), info, ptr, info->size);
        return false;
    }

    std::ifstream outputOrigin;
    // First find key
    {
        std::ostringstream outputFileOs;
        outputFileOs << directName << "/" << name <<".txt";
        outputOrigin.open(outputFileOs.str().c_str());
    }
    // Second find order
    if (outputOrigin.fail()) {
        std::ostringstream outputFileOs;
        outputFileOs << directName << "/" << order <<".txt";
        outputOrigin.open(outputFileOs.str().c_str());
    }
    if (outputOrigin.fail()) {
        MNN_PRINT("Skip check %s\n", name.c_str());
        return true;
    }
    MNN_PRINT("before compare %s: (", name.c_str());
    for (int i=0; i<info->dim.size(); ++i) {
        MNN_PRINT("%d, ", info->dim[i]);
    }
    MNN_PRINT(")\n");
    auto outputPtr = output->readMap<float>();
    float diffAbsMaxV = 0.0f;
    float absMaxV = 0.0f;
#define MNN_IS_INF(x) (fabs(x) == INFINITY)
#define MNN_IS_NAN(x) ((x) != (x))
    for (int i=0; i<info->size; ++i) {
        double targetValue;
        outputOrigin >> targetValue;
        if (MNN_IS_INF(outputPtr[i]) || MNN_IS_NAN(outputPtr[i])) {
            MNN_ERROR("TESTERROR %s value error:%f\n", name.c_str(), outputPtr[i]);
            return false;
        }
        auto diff = fabsf((float)targetValue - outputPtr[i]);
        absMaxV = fmaxf(absMaxV, targetValue);
        diffAbsMaxV = fmaxf(diff, diffAbsMaxV);
    }

    MNN_PRINT("For %s, max = %f, diffmax = %f, diff rate = %f\n", name.c_str(), absMaxV, diffAbsMaxV, diffAbsMaxV / fmaxf(absMaxV, 1e-6));
    if (absMaxV * 0.01f < diffAbsMaxV || MNN_IS_NAN(absMaxV)) {
        MNN_ERROR("TESTERROR %s value error : absMaxV:%f - DiffMax %f\n", name.c_str(), absMaxV, diffAbsMaxV);
        return false;
    }
    return true;
}

static inline std::vector<int> parseIntList(const std::string& str, char delim) {
    std::vector<int> result;
    if (str.empty()) {
        return result;
    }
    std::ptrdiff_t p1 = 0, p2;
    while (1) {
        p2 = str.find(delim, p1);
        if (p2 != std::string::npos) {
            result.push_back(atoi(str.substr(p1, p2 - p1).c_str()));
            p1 = p2 + 1;
        } else {
            result.push_back(atoi(str.substr(p1).c_str()));
            break;
        }
    }
    return result;
}
int main(int argc, char *argv[]) {
    if (argc < 3) {
        MNN_PRINT("=======================================================================================================================================\n");
        MNN_ERROR("Usage: ./ModuleBasic.out ${test.mnn} ${Dir} [runMask] [forwardType] [runLoops] [numberThread] [precision | memory] [cacheFile] [cpuIds] [enableKleidiAI]\n");
        MNN_PRINT("=======================================================================================================================================\n");
        return 0;
    }
    BackendConfig backendConfigTmp;
    auto _executor = Executor::newExecutor(MNN_FORWARD_CPU, backendConfigTmp, 1);
    ExecutorScope _s(_executor);

    std::string modelName = argv[1];
    std::string directName = argv[2];
    MNN_PRINT("Test %s from input info: %s\n", modelName.c_str(), directName.c_str());
    std::map<std::string, float> inputInfo;
    std::map<std::string, std::vector<int>> inputShape;
    std::vector<std::string> inputNames;
    std::vector<std::string> outputNames;
    bool checkOutput = false;
    int runMask = 0;
    if (argc > 3) {
        runMask = atoi(argv[3]);
        if (runMask & 1) {
            _initDebug();
        }
        if (runMask & 2) {
            _initTensorStatic();
        }
    }
    int repeatNumber = 2;
    bool shapeMutable = true;
    std::vector<VARP> inputs;
    std::vector<VARP> outputs;
    if (runMask & 128) {
        MNN_PRINT("Use input.mnn and output.mnn for test\n");
        inputs = MNN::Express::Variable::load((directName + "/input.mnn").c_str());
        outputs = MNN::Express::Variable::load((directName + "/output.mnn").c_str());
        if (inputs.size() > 0 && outputs.size() > 0) {
            MNN_PRINT("Has input.mnn, use input.mnn and output.mnn instead of json\n");
        }
        for (auto v : inputs) {
            inputNames.emplace_back(v->name());
        }
        for (auto v : outputs) {
            outputNames.emplace_back(v->name());
        }
        checkOutput = outputs.size() > 0;
    }
    // Call Time / Per Second
    float freq = 0.0f;
    int cpuDecreaseRate = -1;
    int kvAdd = 0;
    if (inputNames.empty()) {
        rapidjson::Document document;
        std::ostringstream jsonNameOs;
        jsonNameOs << directName << "/input.json";
        std::ifstream fileNames(jsonNameOs.str().c_str());
        std::ostringstream output;
        output << fileNames.rdbuf();
        auto outputStr = output.str();
        document.Parse(outputStr.c_str());
        if (document.HasParseError()) {
            MNN_ERROR("Invalid json\n");
            return 0;
        }
        if (document.HasMember("inputs")) {
            auto inputsInfo = document["inputs"].GetArray();
            for (auto iter = inputsInfo.begin(); iter !=inputsInfo.end(); iter++) {
                auto obj = iter->GetObject();
                std::string name = obj["name"].GetString();
                inputNames.emplace_back(name);
                MNN_PRINT("%s\n", name.c_str());
                if (obj.HasMember("value")) {
                    float value = obj["value"].GetFloat();
                    inputInfo.insert(std::make_pair(name, value));
                }
                if (obj.HasMember("shape")) {
                    auto dims = obj["shape"].GetArray();
                    std::vector<int> shapes;
                    for (auto iter = dims.begin(); iter != dims.end(); iter++) {
                        shapes.emplace_back(iter->GetInt());
                    }
                    inputShape.insert(std::make_pair(name, shapes));
                }
            }
        }
        if (document.HasMember("outputs")) {
            checkOutput = true;
            auto array = document["outputs"].GetArray();
            for (auto iter = array.begin(); iter !=array.end(); iter++) {
                std::string name = iter->GetString();
                MNN_PRINT("output: %s\n", name.c_str());
                outputNames.emplace_back(name);
            }
        }
        if (document.HasMember("shapeMutable")) {
            shapeMutable = document["shapeMutable"].GetBool();
        }
        if (document.HasMember("repeat")) {
            repeatNumber = document["repeat"].GetInt();
        }
        if (document.HasMember("freq")) {
            freq = document["freq"].GetFloat();
        }
        if (document.HasMember("kv_add")) {
            kvAdd = document["kv_add"].GetInt();
        }
        if (document.HasMember("cpu_decrease_rate")) {
            cpuDecreaseRate = document["cpu_decrease_rate"].GetInt();
        }
    }
    auto type = MNN_FORWARD_CPU;
    if (argc > 4) {
        type = (MNNForwardType)atoi(argv[4]);
        MNN_PRINT("Use extra forward type: %d\n", type);
    }

    // Default single thread
    int modeNum = 1;
    if (argc > 6) {
        modeNum = ::atoi(argv[6]);
    }

    int power = BackendConfig::Power_Normal;
    int precision = BackendConfig::Precision_Normal;
    int memory = BackendConfig::Memory_Normal;
    if (argc > 7) {
        int mask = atoi(argv[7]);
        precision = mask % 4;
        memory = (mask / 4) % 4;
        power = (mask / 16) % 4;
    }
    const char* cacheFileName = ".tempcache";
    if (argc > 8) {
        cacheFileName = argv[8];
    }
    // CPU IDs
    std::vector<int> cpuIds;
    if (argc > 9) {
        cpuIds = parseIntList(argv[9], ',');
    }
    MNN_PRINT("cpuIds: ");
    for (auto id : cpuIds) {
        MNN_PRINT("%d ", id);
    }
    bool enableKleidiAI = false;
    if (argc > 10) {
        enableKleidiAI = atoi(argv[10]) > 0 ? true : false;
    }
    int mixedRatio = 17;
    if (argc > 11) {
        mixedRatio = atoi(argv[11]);
    }
    MNN_PRINT("\n");
    FUNC_PRINT(precision);
    FUNC_PRINT(memory);
    FUNC_PRINT(power);
    FUNC_PRINT_ALL(cacheFileName, s);
    FUNC_PRINT(enableKleidiAI);
    FUNC_PRINT(mixedRatio);
    // create session
    MNN::ScheduleConfig config;
    config.type      = type;
    /*modeNum means gpuMode for GPU usage, Or means numThread for CPU usage.*/
    config.numThread = modeNum;
    // If type not fount, let it failed
    config.backupType = type;
    BackendConfig backendConfig;
    // config.path.outputs.push_back("ResizeBilinear_2");
    backendConfig.power = (BackendConfig::PowerMode)power;
    backendConfig.precision = static_cast<MNN::BackendConfig::PrecisionMode>(precision);
    backendConfig.memory = static_cast<MNN::BackendConfig::MemoryMode>(memory);
    config.backendConfig     = &backendConfig;

    MNN::Express::Module::Config mConfig;
    if (runMask & 256) {
        mConfig.dynamic = true;
    }
    mConfig.shapeMutable = shapeMutable;
    std::shared_ptr<Executor::RuntimeManager> rtmgr(Executor::RuntimeManager::createRuntimeManager(config));
    rtmgr->setCache(cacheFileName);
    rtmgr->setHint(MNN::Interpreter::INIT_THREAD_NUMBER, 4);
    rtmgr->setHint(MNN::Interpreter::HintMode::CPU_CORE_IDS, cpuIds.data(), cpuIds.size());

    if (cpuDecreaseRate > 0 && cpuDecreaseRate <= 100) {
        rtmgr->setHint(Interpreter::CPU_LITTLECORE_DECREASE_RATE, cpuDecreaseRate);
    }
    if (runMask & 1) {
        // Need dump tensor, open debug
        rtmgr->setMode(Interpreter::Session_Debug);
    }
    if (runMask & 2) {
        // Need tensor static for each op, open debug
        rtmgr->setMode(Interpreter::Session_Debug);
    }
    // For Debug
    if (false) {
        int geometryMask = Interpreter::GeometryComputeMask::GEOMETRCOMPUTEMASK_ALL;
        geometryMask -= Interpreter::GeometryComputeMask::GEOMETRCOMPUTEMASK_FUSEREGION;
        geometryMask -= Interpreter::GeometryComputeMask::GEOMETRCOMPUTEMASK_OPENCACHE;
        rtmgr->setHint(Interpreter::GEOMETRY_COMPUTE_MASK, geometryMask);
    }
    if (runMask & 4) {
        // Need time trace for each op, open debug
        rtmgr->setMode(Interpreter::Session_Debug);
    }
    if (runMask & 8) {
        rtmgr->setMode(Interpreter::Session_Input_Inside);
    }
    if (runMask & 16) {
        rtmgr->setMode(Interpreter::Session_Backend_Auto);
        rtmgr->setHint(Interpreter::MAX_TUNING_NUMBER, 50);
    }
    if (runMask & 32) {
        mConfig.rearrange = true;
    }
    if (runMask & 512) {
        rtmgr->setHint(Interpreter::WINOGRAD_MEMORY_LEVEL, 0);
    }
    if (runMask & 1024) {
        /*
        2: INPUT_BLOCK_QUANT
        1: INPUT_SHARE_ONE_SCALE
        0: INPUT_CHANNEL_QUANT
        */
        rtmgr->setHint(Interpreter::DYNAMIC_QUANT_OPTIONS, 2);
    }

    if (enableKleidiAI) {
        rtmgr->setHint(Interpreter::CPU_ENABLE_KLEIDIAI, true);
    }
    KVMeta kvMeta;
    if (kvAdd > 0) {
        kvMeta.add = kvAdd;
        rtmgr->setHintPtr(Interpreter::KVCACHE_INFO, &kvMeta);
    }

    // rtmgr->setHint(Interpreter::CPU_SME2_INSTRUCTIONS, false);

    if (runMask & 2048) {
        rtmgr->setExternalPath("tmp", Interpreter::EXTERNAL_FEATUREMAP_DIR);
    }
    
    rtmgr->setHint(Interpreter::CPU_SME2_NEON_DIVISION_RATIO, mixedRatio);
    // set npu model dir, npu model and mnn model in same path
    size_t pos = modelName.find_last_of("/\\");
    std::string modelPath;
    if (pos == std::string::npos) {
        // current path
        modelPath = "./";
    } else {
        modelPath = modelName.substr(0, pos);
    }
    rtmgr->setExternalPath(modelPath, 3);
    std::shared_ptr<Module> net;
    {
        AUTOTIME;
        net.reset(Module::load(inputNames, outputNames, modelName.c_str(), rtmgr, &mConfig));
        if (net == nullptr) {
            MNN_PRINT("Error: can't load module\n");
            return 0;
        }
        if (runMask & 64) {
            net.reset(Module::clone(net.get()));
        }
        if (net == nullptr) {
            return 0;
        }
    }
    auto mInfo = net->getInfo();

#define LOAD_DATA(TYPE)\
    if (inputInfo.find(inputName) != inputInfo.end()) {\
        auto value = inputInfo[inputName];\
        for (int i=0; i<info->size; ++i) {\
            ptr[i] = value;\
        }\
    } else {\
        std::ostringstream fileNameOs;\
        fileNameOs << directName << "/" << inputName << ".txt";\
        auto fileName = fileNameOs.str();\
        std::ifstream inputOs(fileName.c_str());\
        if (inputOs.fail()) {\
            MNN_ERROR("TESTERROR Can't open %s\n", fileName.c_str());\
            continue;\
        }\
        for (int i=0; i<info->size; ++i) {\
            double tempValue;\
            inputOs >> tempValue;\
            ptr[i] = tempValue;\
        }\
    }

    if (inputs.empty()) {
        inputs.resize(mInfo->inputs.size());
        for (int i=0; i<inputs.size(); ++i) {
            inputs[i] = _Input(mInfo->inputs[i].dim, mInfo->inputs[i].order, mInfo->inputs[i].type);
        }
        // Load inputs
        for (int i=0; i<inputs.size(); ++i) {
            auto inputName = inputNames[i];
            // Resize
            auto shapeIter = inputShape.find(inputName);
            auto order = mInfo->inputs[i].order;
            if (MNN::Express::Dimensionformat::NC4HW4 == mInfo->inputs[i].order) {
                order = MNN::Express::Dimensionformat::NCHW;
            }

            if (shapeIter != inputShape.end()) {
                auto s = shapeIter->second;
                inputs[i] = _Input(s, order, mInfo->inputs[i].type);
            }
            auto info = inputs[i]->getInfo();
            if (info->type == halide_type_of<float>()){
                auto ptr = inputs[i]->writeMap<float>();
                LOAD_DATA(float)
            } else {
                auto floatVar = _Input(info->dim, info->order, halide_type_of<float>());
                auto ptr = floatVar->writeMap<float>();
                LOAD_DATA(float)
                auto temp = _Cast(floatVar, info->type);
                inputs[i]->input(temp);
            }
            if (MNN::Express::Dimensionformat::NC4HW4 == mInfo->inputs[i].order) {
                inputs[i] = _Convert(inputs[i], MNN::Express::Dimensionformat::NC4HW4);
            }
        }
    }
#undef LOAD_DATA

    bool modelError = false;
    for (int repeat = 0; repeat < repeatNumber; ++repeat) {
        MNN_PRINT("Run for %d time\n", repeat);
        std::vector<VARP> subInputs = inputs;
        if (repeat % 2 == 1) {
            for (int i=0; i<inputs.size(); ++i) {
                subInputs[i] = _Clone(inputs[i], true);
            }
        }
        kvMeta.add = kvAdd;
        auto outputs = net->onForward(inputs);
        kvMeta.sync();
        if (outputs.empty()) {
            MNN_ERROR("Error in forward\n");
            return 0;
        }
        for (int i=0; i<outputNames.size(); ++i) {
            auto name = outputNames[i];
            auto v = outputs[i];
            auto info = v->getInfo();
            if (nullptr == info) {
                continue;
            }
            if (info->order == NC4HW4 && info->dim.size() > 1) {
                v = _Convert(v, mInfo->defaultFormat);
            }
            if (info->type.code != halide_type_float) {
                v = _Cast<float>(v);
            }
            v.fix(VARP::CONSTANT);
            outputs[i] = v;
        }
        if (checkOutput) {
            for (int i=0; i<outputNames.size(); ++i) {
                auto output = outputs[i];
                bool success = compareOutput(output, directName, outputNames[i], mInfo->defaultFormat, i);
                if (!success) {
                    modelError = true;
                    MNN_ERROR("%d run Error for output %s\n", repeat, outputNames[i].c_str());
                }
            }
        }
        if (0 == repeat) {
            for (int i=0; i<inputNames.size(); ++i) {
                inputs[i].fix(VARP::CONSTANT);
                inputs[i]->setName(inputNames[i]);
            }
            for (int i=0; i<outputNames.size(); ++i) {
                outputs[i].fix(VARP::CONSTANT);
                outputs[i]->setName(outputNames[i]);
            }
            Variable::save(inputs, "output/input.mnn");
            Variable::save(outputs, "output/output.mnn");
        }
        for (int i=0; i<outputNames.size(); ++i) {
            auto name = outputNames[i];
            auto v = outputs[i];
            auto info = v->getInfo();
            std::ostringstream fileNameOs;
            fileNameOs << "output/" << repeat <<"_"<< i << ".txt";
            auto fileName = fileNameOs.str();
            MNN_PRINT("Write %s output to %s\n", name.c_str(), fileName.c_str());
            std::ofstream _output(fileName.c_str());
            auto ptr = v->readMap<float>();
            for (int v=0; v<info->size; ++v) {
                _output << ptr[v] << "\n";
            }
        }
        // Print module's memory
        float memoryInMB = 0.0f;
        rtmgr->getInfo(Interpreter::MEMORY, &memoryInMB);
        FUNC_PRINT_ALL(memoryInMB, f);
    }

    // benchmark. for CPU, op time means calc duration; for others, op time means schedule duration.
    int runTime = 0;
    if (argc > 5) {
        runTime = ::atoi(argv[5]);
    }

    if (runTime > 0) {
        kvMeta.remove = kvMeta.previous;
        int t = runTime;
        if (runMask & 4) {
            _initTimeTrace();
        }
        float minTime = std::numeric_limits<float>::max();
        float maxTime = 0.0f;
        float sum    = 0.0f;

        for (int i = 0; i < t; ++i) {
            Timer _l;
            kvMeta.add = kvAdd;
            auto out = net->onForward(inputs);
            kvMeta.sync();
            Variable::compute(out);
            for (auto o : out) {
                ((MNN::Tensor*)o->getTensor())->wait(MNN::Tensor::MAP_TENSOR_READ, true);
            }
            auto time = _l.durationInUs() / 1000.0f;
            if (freq > 0.0f) {
                float remainMs = (1000.0f / freq) - time;
                if (remainMs > 0.0f) {
                    std::this_thread::sleep_for(std::chrono::milliseconds((int)remainMs));
                }
            }
            if (maxTime < time) {
                maxTime = time;
            }
            if (minTime > time) {
                minTime = time;
            }
            sum += time;
        }
        if (nullptr != gTimeTraceInfo) {
            MNN_PRINT("Per Op Trace: \n");
            gTimeTraceInfo->dump(true);
            MNN_PRINT("Per Type Trace: \n");
            gTimeTraceInfo->dump(false);
        }
        MNN_PRINT("Avg= %f ms, min= %f ms, max= %f ms\n", sum / (float)t, minTime, maxTime);
    }
    rtmgr->updateCache();
    return 0;
}