/*
 * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef TRT_SAMPLE_OPTIONS_H
#define TRT_SAMPLE_OPTIONS_H

#include <array>
#include <iostream>
#include <optional>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#include "NvInfer.h"

#if ENABLE_UNIFIED_BUILDER
#include "safeCommon.h"
#endif

namespace sample
{

// Build default params
constexpr int32_t defaultAvgTiming{8};
constexpr int32_t defaultMaxAuxStreams{-1};
constexpr int32_t defaultBuilderOptimizationLevel{-1};
constexpr int32_t defaultTilingOptimizationLevel{static_cast<int32_t>(nvinfer1::TilingOptimizationLevel::kNONE)};
constexpr int32_t defaultMaxTactics{-1};

// System default params
constexpr int32_t defaultDevice{0};

// Inference default params
constexpr int32_t defaultBatch{1};
constexpr int32_t batchNotProvided{0};
constexpr int32_t defaultStreams{1};
constexpr int32_t defaultIterations{10};
constexpr int32_t defaultOptProfileIndex{0};
constexpr float defaultWarmUp{200.F};
constexpr float defaultDuration{3.F};
constexpr float defaultSleep{};
constexpr float defaultIdle{};
constexpr float defaultPersistentCacheRatio{0};

// Reporting default params
constexpr int32_t defaultAvgRuns{10};
constexpr std::array<float, 3> defaultPercentiles{90, 95, 99};


enum class ModelFormat
{
    kANY,
    kONNX
};

enum class SparsityFlag
{
    kDISABLE,
    kENABLE,
    kFORCE
};

enum class TimingCacheMode
{
    kDISABLE,
    kLOCAL,
    kGLOBAL
};

enum class MemoryAllocationStrategy
{
    kSTATIC,  //< Allocate device memory based on max size across all profiles.
    kPROFILE, //< Allocate device memory based on max size of the current profile.
    kRUNTIME, //< Allocate device memory based on the current input shapes.
};

enum class AccuracyValidationAlgorithm
{
    kL0,               //< L0 accuracy algorithm
    kL1,               //< L1 accuracy algorithm
    kL2,               //< L2 accuracy algorithm
    kLInf,             //< LInf accuracy algorithm
    kCosineSimilarity, //< Cosine similarity accuracy algorithm
};

enum class TuningSearchAlgorithm
{
    kFAST,       //< Fast searching algorithm: baseline + one-off variations (linear in #knobs)
    kEXHAUSTIVE, //< Exhaustive: all combinations enumerated (product of knob value-counts)
    kMIXED,      //< Two-phase: fast scan, then exhaustive over knobs that improved performance
};

//! \brief Convert a TuningSearchAlgorithm enum to its CLI / cache-file string.
//!   kFAST       -> "fast"
//!   kEXHAUSTIVE -> "full"
//!   kMIXED      -> "mixed"
inline std::string toString(TuningSearchAlgorithm algo)
{
    switch (algo)
    {
    case TuningSearchAlgorithm::kFAST: return "fast";
    case TuningSearchAlgorithm::kEXHAUSTIVE: return "full";
    case TuningSearchAlgorithm::kMIXED: return "mixed";
    }
    return "unknown";
}

//!
//! \enum RuntimeMode
//!
//! \brief Used to dictate which TensorRT runtime library to dynamically load.
//!
enum class RuntimeMode
{
    //! Maps to libnvinfer.so or nvinfer.dll
    kFULL,

    //! Maps to libnvinfer_dispatch.so or nvinfer_dispatch.dll
    kDISPATCH,

    //! Maps to libnvinfer_lean.so or nvinfer_lean.dll
    kLEAN,

    //! Maps to libnvinfer_safe.so or nvinfer_safe.dll
    kSAFE,
};

inline std::ostream& operator<<(std::ostream& os, RuntimeMode const mode)
{
    switch (mode)
    {
    case RuntimeMode::kFULL:
    {
        os << "full";
        break;
    }
    case RuntimeMode::kDISPATCH:
    {
        os << "dispatch";
        break;
    }
    case RuntimeMode::kLEAN:
    {
        os << "lean";
        break;
    }
    case RuntimeMode::kSAFE:
    {
        os << "safe";
        break;
    }
    }

    return os;
}

using Arguments = std::unordered_multimap<std::string, std::pair<std::string, int32_t>>;

//! An IO format specification.
struct IOFormat
{
    nvinfer1::TensorFormats formats{};
};

using ShapeRange = std::array<std::vector<int64_t>, nvinfer1::EnumMax<nvinfer1::OptProfileSelector>()>;

using LayerDeviceTypes = std::unordered_map<std::string, nvinfer1::DeviceType>;
using DecomposableAttentions = std::unordered_map<std::string, bool>;

using StringSet = std::unordered_set<std::string>;

class WeightStreamingBudget
{
public:
    static constexpr int64_t kDISABLE{-2};
    static constexpr int64_t kAUTOMATIC{-1};
    int64_t bytes{kDISABLE};
    double percent{static_cast<double>(100.0)};

    bool isDisabled()
    {
        return bytes == kDISABLE && percent == kDISABLE;
    }
};

class Options
{
public:
    virtual ~Options() = default;
    virtual void parse(Arguments& arguments) = 0;
};

class BaseModelOptions : public Options
{
public:
    ModelFormat format{ModelFormat::kANY};
    std::string model;

    void parse(Arguments& arguments) override;

    static void help(std::ostream& out);
};

class ModelOptions : public Options
{
public:
    BaseModelOptions baseModel;
    std::string prototxt;
    std::vector<std::string> outputs;

    void parse(Arguments& arguments) override;

    static void help(std::ostream& out);
};

constexpr nvinfer1::TempfileControlFlags getTempfileControlDefaults()
{
    using F = nvinfer1::TempfileControlFlag;
    return (1U << static_cast<uint32_t>(F::kALLOW_TEMPORARY_FILES))
        | (1U << static_cast<uint32_t>(F::kALLOW_IN_MEMORY_FILES));
}

class BuildOptions : public Options
{
public:
    // Unit in MB.
    double workspace{-1.0};
    // Unit in MB.
    double dlaSRAM{-1.0};
    // Unit in MB.
    double dlaLocalDRAM{-1.0};
    // Unit in MB.
    double dlaGlobalDRAM{-1.0};
    // Unit in KB.
    double tacticSharedMem{-1.0};
    int32_t avgTiming{defaultAvgTiming};
    bool tf32{true};
    bool stronglyTyped{true};
    bool directIO{false};
    LayerDeviceTypes layerDeviceTypes;
    DecomposableAttentions decomposableAttentions;
    StringSet debugTensors;
    bool markUnfusedTensorsAsDebugTensors{false};
    StringSet debugTensorStates;
    bool safe{false};
    bool consistency{false};
    bool dumpKernelText{false};
    bool buildDLAStandalone{false};
    bool allowGPUFallback{false};
    bool skipInference{false};
    bool save{false};
    bool saveAllEngines{false}; //!< Save per-iteration engines as <engine>.iter<N> during tuning
    bool load{false};
    bool asyncFileReader{false};
    bool refittable{false};
    bool stripWeights{false};
    bool versionCompatible{false};
    bool pluginInstanceNorm{false};
    bool enableUInt8AsymmetricQuantizationDLA{false};
    bool reportCapabilityDLA{false};
    bool adjustForDLA{false};
    bool enablePluginOverride{false};
    bool excludeLeanRuntime{false};
    bool disableCompilationCache{false};
    bool enableMonitorMemory{false};
    bool cpuOnly{false};
    int32_t builderOptimizationLevel{defaultBuilderOptimizationLevel};
    int32_t maxTactics{defaultMaxTactics};
    SparsityFlag sparsity{SparsityFlag::kDISABLE};
    nvinfer1::ProfilingVerbosity profilingVerbosity{nvinfer1::ProfilingVerbosity::kLAYER_NAMES_ONLY};
    std::string engine;
    using ShapeProfile = std::unordered_map<std::string, ShapeRange>;
    std::vector<ShapeProfile> optProfiles;
    std::vector<IOFormat> inputFormats;
    std::vector<IOFormat> outputFormats;
    nvinfer1::TacticSources enabledTactics{0};
    nvinfer1::TacticSources disabledTactics{0};
    TimingCacheMode timingCacheMode{TimingCacheMode::kLOCAL};
    std::string timingCacheFile{};
    bool errorOnTimingCacheMiss{false};
    // C++11 does not automatically generate hash function for enum class.
    // Use int32_t to support C++11 compilers.
    std::unordered_map<int32_t, bool> previewFeatures;
    nvinfer1::HardwareCompatibilityLevel hardwareCompatibilityLevel{nvinfer1::HardwareCompatibilityLevel::kNONE};
    nvinfer1::RuntimePlatform runtimePlatform{nvinfer1::RuntimePlatform::kSAME_AS_BUILD};
    std::string tempdir{};
    nvinfer1::TempfileControlFlags tempfileControls{getTempfileControlDefaults()};
    RuntimeMode useRuntime{RuntimeMode::kFULL};
    std::string leanDLLPath{};
    std::string buildRoute{}; //!< --setBuildRoute=<route> string; passed to IBuilderConfig::setBuildRoute()
    int32_t maxAuxStreams{defaultMaxAuxStreams};
    bool getPlanVersionOnly{false};

    bool allowWeightStreaming{false};

    int32_t tilingOptimizationLevel{defaultTilingOptimizationLevel};
    int64_t l2LimitForTiling{-1};
    bool distributiveIndependence{false};
    std::string remoteAutoTuningConfig{};

    void parse(Arguments& arguments) override;

    static void help(std::ostream& out);
};

class SystemOptions : public Options
{
public:
    int32_t device{defaultDevice};
    int32_t DLACore{-1};
    bool enableStaticPlugins{true};
    bool ignoreParsedPluginLibs{false};
    std::vector<std::string> plugins;
    std::vector<std::string> setPluginsToSerialize;
    std::vector<std::string> dynamicPlugins;
#if ENABLE_UNIFIED_BUILDER
    std::vector<samplesSafeCommon::SafetyPluginLibraryArgument> safetyPlugins;
#endif

    void parse(Arguments& arguments) override;

    static void help(std::ostream& out, bool enableStaticPlugins = true);
};

class InferenceOptions : public Options
{
public:
    int32_t batch{batchNotProvided};
    int32_t iterations{defaultIterations};
    int32_t infStreams{defaultStreams};
    int32_t optProfileIndex{defaultOptProfileIndex};
    float warmup{defaultWarmUp};
    float duration{defaultDuration};
    float sleep{defaultSleep};
    float idle{defaultIdle};
    float persistentCacheRatio{defaultPersistentCacheRatio};
    float atol{1e-5};                  // Element-wise accuracy threshold absolute tolerance
    float rtol{1e-5};                  // Element-wise accuracy threshold relative tolerance
    float accuracyThresholdEndToEnd{}; // End-to-end accuracy threshold should not have default value
                                       // because it depends on the model and accuracy validation algorithm
    bool overlap{true};
    bool includeTransfers{false};
    bool useManaged{false};
    bool spin{true};
    bool threads{false};
    bool graph{true};
    bool timeDeserialize{false};
    bool timeRefit{false};
    bool setOptProfile{false};
    // Reference pairs for accuracy validation with multiple test cases.
    // Each pair contains (inputMap, refOutputMap) indexed by pair number.
    // refPairs[0] is used as the primary pair (for --loadInputs/--loadRefOutputs without --refPair).
    using RefPair = std::pair<std::unordered_map<std::string, std::string>,
        std::unordered_map<std::string, std::string>>; // pair of (inputMap, refOutputMap)
    std::vector<RefPair> refPairs{1};                  // Initialize with one empty pair
    AccuracyValidationAlgorithm accuracyValidationAlgorithm{AccuracyValidationAlgorithm::kL0};
    using ShapeProfile = std::unordered_map<std::string, std::vector<int64_t>>;
    ShapeProfile shapes;
    nvinfer1::ProfilingVerbosity nvtxVerbosity{nvinfer1::ProfilingVerbosity::kLAYER_NAMES_ONLY};
    MemoryAllocationStrategy memoryAllocationStrategy{MemoryAllocationStrategy::kSTATIC};
    std::unordered_map<std::string, std::string> debugTensorFileNames;
    std::vector<std::string> dumpAlldebugTensorFormats;
    WeightStreamingBudget weightStreamingBudget;
    std::string refitOnnxModel;

    void parse(Arguments& arguments) override;

    static void help(std::ostream& out);
};

class ReportingOptions : public Options
{
public:
    bool verbose{false};
    int32_t avgs{defaultAvgRuns};
    std::vector<float> percentiles{defaultPercentiles.begin(), defaultPercentiles.end()};
    bool refit{false};
    bool output{false};
    bool dumpRawBindings{false};
    bool profile{false};
    bool layerInfo{false};
    bool optProfileInfo{false};
    std::string exportTimes;
    std::string exportOutput;
    std::string exportProfile;
    std::string exportLayerInfo;

    void parse(Arguments& arguments) override;

    static void help(std::ostream& out);
};

class SafeBuilderOptions : public Options
{
public:
    std::string serialized{};
    std::string onnxModelFile{};
    bool help{false};
    bool verbose{false};
    std::vector<IOFormat> inputFormats;
    std::vector<IOFormat> outputFormats;
    bool int8{false};
    bool fp8{false};
    bool int4{false};
    std::vector<std::string> plugins;
    bool consistency{false};
    bool standard{false};
    TimingCacheMode timingCacheMode{TimingCacheMode::kLOCAL};
    std::string timingCacheFile{};
    SparsityFlag sparsity{SparsityFlag::kDISABLE};
    int32_t avgTiming{defaultAvgTiming};

    void parse(Arguments& arguments) override;

    static void printHelp(std::ostream& out);
};

//! \brief Options for `--tuneBuildRoutes`-driven autotuning of build routes.
//!
//! The tuning loop is implemented in samples/trtexec/trtexec.cpp by forking
//! the trtexec process for each route; each child runs `runOnceBuildAndInfer`
//! with `--setBuildRoute=<route>` injected by the parent.
class TuningOptions : public Options
{
public:
    std::string tuningCacheFile{"best_config.json"};
    std::string tuningExpr{};                                          //!< --tuneBuildRoutes
    std::string tuningExprFile{};                                      //!< --tuneBuildRouteFile
    TuningSearchAlgorithm tuningSearchAlgorithm{TuningSearchAlgorithm::kFAST};
    int64_t timeout{-1};                                               //!< --tuningTimeOut (s); -1 = no timeout
    bool helpBuildRoute{false};                                        //!< --helpBuildRoute (short-circuit)
    std::string helpBuildRouteKnob{};                                  //!< --helpBuildRoute=<knob> filter
    bool continueFromCache{false};                                     //!< --continue
    bool dryRun{false};                                                //!< --dryRun (enumerate, don't build)
    //! \brief Hidden parent->child IPC channel.
    //!
    //! When set, runOnceBuildAndInfer writes a small JSON to this path containing
    //! gpu_time_ms, accuracy_failed, and per-tensor accuracy_loss before returning.
    //! Injected into the child's argv by the tuning loop; never shown in --help.
    std::string tuningResultFile{};                                    //!< --tuningResultFile=<path>

    void parse(Arguments& arguments) override;
    static void help(std::ostream& out);
};

class AllOptions : public Options
{
public:
    ModelOptions model;
    BuildOptions build;
    SystemOptions system;
    InferenceOptions inference;
    ReportingOptions reporting;
    TuningOptions tuning;
    bool helps{false};

    void parse(Arguments& arguments) override;

    static void help(std::ostream& out, bool enableStaticPlugins = true);
};

class TaskInferenceOptions : public Options
{
public:
    std::string engine;
    int32_t device{defaultDevice};
    int32_t DLACore{-1};
    int32_t batch{batchNotProvided};
    bool graph{true};
    float persistentCacheRatio{defaultPersistentCacheRatio};
    void parse(Arguments& arguments) override;
    static void help(std::ostream& out);
};

Arguments argsToArgumentsMap(int32_t argc, char* argv[]);

bool parseHelp(Arguments& arguments);

void helpHelp(std::ostream& out);

// Functions to print options

std::ostream& operator<<(std::ostream& os, BaseModelOptions const& options);

std::ostream& operator<<(std::ostream& os, IOFormat const& format);

std::ostream& operator<<(std::ostream& os, ShapeRange const& dims);

std::ostream& operator<<(std::ostream& os, ModelOptions const& options);

std::ostream& operator<<(std::ostream& os, BuildOptions const& options);

std::ostream& operator<<(std::ostream& os, SystemOptions const& options);

std::ostream& operator<<(std::ostream& os, InferenceOptions const& options);

std::ostream& operator<<(std::ostream& os, ReportingOptions const& options);

std::ostream& operator<<(std::ostream& os, AllOptions const& options);

std::ostream& operator<<(std::ostream& os, SafeBuilderOptions const& options);

std::ostream& operator<<(std::ostream& os, nvinfer1::DataType dtype);

std::ostream& operator<<(std::ostream& os, nvinfer1::DeviceType devType);


inline std::ostream& operator<<(std::ostream& os, nvinfer1::Dims const& dims)
{
    for (int32_t i = 0; i < dims.nbDims; ++i)
    {
        os << (i ? "x" : "") << dims.d[i];
    }
    return os;
}
inline std::ostream& operator<<(std::ostream& os, const nvinfer1::WeightsRole role)
{
    switch (role)
    {
    case nvinfer1::WeightsRole::kKERNEL:
    {
        os << "Kernel";
        break;
    }
    case nvinfer1::WeightsRole::kBIAS:
    {
        os << "Bias";
        break;
    }
    case nvinfer1::WeightsRole::kSHIFT:
    {
        os << "Shift";
        break;
    }
    case nvinfer1::WeightsRole::kSCALE:
    {
        os << "Scale";
        break;
    }
    case nvinfer1::WeightsRole::kCONSTANT:
    {
        os << "Constant";
        break;
    }
    case nvinfer1::WeightsRole::kANY:
    {
        os << "Any";
        break;
    }
    }

    return os;
}

inline std::ostream& operator<<(std::ostream& os, std::vector<int64_t> const& vec)
{
    for (int32_t i = 0, e = static_cast<int32_t>(vec.size()); i < e; ++i)
    {
        os << (i ? "x" : "") << vec[i];
    }
    return os;
}

} // namespace sample

#endif // TRT_SAMPLES_OPTIONS_H