//
// Copyright 2024 NVIDIA Corporation. All rights reserved
//

// System headers
#include <iostream>
#include <iomanip>
#include <fstream>
#include <vector>
#include <unordered_map>
#include <sstream>
#include <algorithm>

#if defined(_WIN32)
    #include <process.h>
#else
    #include <unistd.h>
#endif

// CUPTI headers
#include "helper_cupti.h"
#include <cupti_target.h>
#include <cupti_pmsampling.h>
#include <cupti_profiler_target.h>
#include <cupti_profiler_host.h>

struct SamplerRange
{
    size_t rangeIndex;
    uint64_t startTimestamp;
    uint64_t endTimestamp;
    std::unordered_map<std::string, double> metricValues;
};

class CuptiProfilerHost
{
    std::string m_chipName;
    std::vector<SamplerRange> m_samplerRanges;
    CUpti_Profiler_Host_Object* m_pHostObject = nullptr;

public:
    CuptiProfilerHost() = default;
    ~CuptiProfilerHost() = default;

    void SetUp(std::string chipName, std::vector<uint8_t>& counterAvailibilityImage)
    {
        m_chipName = chipName;
        CUPTI_API_CALL(Initialize(counterAvailibilityImage));
    }

    void TearDown()
    {
        CUPTI_API_CALL(Deinitialize());
    }

    CUptiResult CreateConfigImage(std::vector<const char*> metricsList, std::vector<uint8_t>& configImage)
    {
        // Add metrics to config image
        {
            CUpti_Profiler_Host_ConfigAddMetrics_Params configAddMetricsParams {CUpti_Profiler_Host_ConfigAddMetrics_Params_STRUCT_SIZE};
            configAddMetricsParams.pHostObject = m_pHostObject;
            configAddMetricsParams.ppMetricNames = metricsList.data();
            configAddMetricsParams.numMetrics = metricsList.size();
            CUPTI_API_CALL(cuptiProfilerHostConfigAddMetrics(&configAddMetricsParams));
        }

        // Get Config image size and data
        {
            CUpti_Profiler_Host_GetConfigImageSize_Params getConfigImageSizeParams {CUpti_Profiler_Host_GetConfigImageSize_Params_STRUCT_SIZE};
            getConfigImageSizeParams.pHostObject = m_pHostObject;
            CUPTI_API_CALL(cuptiProfilerHostGetConfigImageSize(&getConfigImageSizeParams));
            configImage.resize(getConfigImageSizeParams.configImageSize);

            CUpti_Profiler_Host_GetConfigImage_Params getConfigImageParams = {CUpti_Profiler_Host_GetConfigImage_Params_STRUCT_SIZE};
            getConfigImageParams.pHostObject = m_pHostObject;
            getConfigImageParams.pConfigImage = configImage.data();
            getConfigImageParams.configImageSize = configImage.size();
            CUPTI_API_CALL(cuptiProfilerHostGetConfigImage(&getConfigImageParams));
        }

        // Get Num of Passes
        {
            CUpti_Profiler_Host_GetNumOfPasses_Params getNumOfPassesParam {CUpti_Profiler_Host_GetNumOfPasses_Params_STRUCT_SIZE};
            getNumOfPassesParam.pConfigImage = configImage.data();
            getNumOfPassesParam.configImageSize = configImage.size();
            CUPTI_API_CALL(cuptiProfilerHostGetNumOfPasses(&getNumOfPassesParam));
            std::cout << "Num of Passes: " << getNumOfPassesParam.numOfPasses << "\n";
        }

        return CUPTI_SUCCESS;
    }

    CUptiResult EvaluateCounterData(CUpti_PmSampling_Object* pSamplingObject, size_t rangeIndex, std::vector<const char*> metricsList, std::vector<uint8_t>& counterDataImage)
    {
        m_samplerRanges.push_back(SamplerRange{});
        SamplerRange& samplerRange = m_samplerRanges.back();

        CUpti_PmSampling_CounterData_GetSampleInfo_Params getSampleInfoParams = {CUpti_PmSampling_CounterData_GetSampleInfo_Params_STRUCT_SIZE};
        getSampleInfoParams.pPmSamplingObject = pSamplingObject;
        getSampleInfoParams.pCounterDataImage = counterDataImage.data();
        getSampleInfoParams.counterDataImageSize = counterDataImage.size();
        getSampleInfoParams.sampleIndex = rangeIndex;
        CUPTI_API_CALL(cuptiPmSamplingCounterDataGetSampleInfo(&getSampleInfoParams));

        samplerRange.startTimestamp = getSampleInfoParams.startTimestamp;
        samplerRange.endTimestamp = getSampleInfoParams.endTimestamp;

        std::vector<double> metricValues(metricsList.size());
        CUpti_Profiler_Host_EvaluateToGpuValues_Params evalauateToGpuValuesParams {CUpti_Profiler_Host_EvaluateToGpuValues_Params_STRUCT_SIZE};
        evalauateToGpuValuesParams.pHostObject = m_pHostObject;
        evalauateToGpuValuesParams.pCounterDataImage = counterDataImage.data();
        evalauateToGpuValuesParams.counterDataImageSize = counterDataImage.size();
        evalauateToGpuValuesParams.ppMetricNames = metricsList.data();
        evalauateToGpuValuesParams.numMetrics = metricsList.size();
        evalauateToGpuValuesParams.rangeIndex = rangeIndex;
        evalauateToGpuValuesParams.pMetricValues = metricValues.data();
        CUPTI_API_CALL(cuptiProfilerHostEvaluateToGpuValues(&evalauateToGpuValuesParams));

        for (size_t i = 0; i < metricsList.size(); ++i) {
            samplerRange.metricValues[metricsList[i]] = metricValues[i];
        }

        return CUPTI_SUCCESS;
    }

    void PrintSampleRanges()
    {
        if (m_samplerRanges.empty())
        {
            std::cout << "No samples to print\n";
            return;
        }

        std::cout << "Total num of Samples: " << m_samplerRanges.size() << "\n";
        const size_t numSamplesToPrint = std::min(static_cast<size_t>(50), m_samplerRanges.size());
        std::cout << "Printing first " << numSamplesToPrint << " samples:" << "\n";
        for(size_t sampleIndex = 0; sampleIndex < numSamplesToPrint; ++sampleIndex)
        {
            const auto& samplerRange = m_samplerRanges[sampleIndex];
            std::cout << "Sample Index: " << sampleIndex << "\n";
            std::cout << "Timestamps -> Start: [" << samplerRange.startTimestamp << "] \tEnd: [" << samplerRange.endTimestamp << "]" << "\n";
            std::cout << "-----------------------------------------------------------------------------------\n";
            for (const auto& metric : samplerRange.metricValues)
            {
                std::cout << std::fixed << std::setprecision(3);
                std::cout << std::setw(50) << std::left << metric.first;
                std::cout << std::setw(30) << std::right << metric.second << "\n";
            }
            std::cout << "-----------------------------------------------------------------------------------\n\n";
        }
    }

    CUptiResult GetSupportedSinglePassMetricsSets(std::vector<const char*>& metricsSetsList)
    {
        CUpti_Profiler_Host_GetSinglePassSets_Params getSinglePassSetsParams {CUpti_Profiler_Host_GetSinglePassSets_Params_STRUCT_SIZE};
        getSinglePassSetsParams.pChipName = m_chipName.c_str();
        CUPTI_API_CALL(cuptiProfilerHostGetSinglePassSets(&getSinglePassSetsParams));

        metricsSetsList.resize(getSinglePassSetsParams.numOfSinglePassSets);
        getSinglePassSetsParams.ppSinglePassSets = metricsSetsList.data();
        CUPTI_API_CALL(cuptiProfilerHostGetSinglePassSets(&getSinglePassSetsParams));
        return CUPTI_SUCCESS;
    }

    CUptiResult GetMetricsInSinglePassSet(const char* pSinglePassSet, std::vector<std::string>& metricsList)
    {
        std::vector<uint8_t> metricsBuffer;
        std::vector<size_t> metricsIndices;
        {
            CUpti_Profiler_Host_GetMetricsInSinglePassSet_Params params {CUpti_Profiler_Host_GetMetricsInSinglePassSet_Params_STRUCT_SIZE};
            params.pChipName = m_chipName.c_str();
            params.pSinglePassSetName = pSinglePassSet;
            CUPTI_API_CALL(cuptiProfilerHostGetMetricsInSinglePassSet(&params));

            metricsBuffer.resize(params.metricsBufferSize);
            metricsIndices.resize(params.numOfMetricsInSinglePassSet);
        }

        {
            CUpti_Profiler_Host_GetMetricsInSinglePassSet_Params params {CUpti_Profiler_Host_GetMetricsInSinglePassSet_Params_STRUCT_SIZE};
            params.pChipName = m_chipName.c_str();
            params.pSinglePassSetName = pSinglePassSet;
            params.numOfMetricsInSinglePassSet = metricsIndices.size();
            params.pMetricsIndicesBuffer = metricsIndices.data();
            params.metricsBufferSize = metricsBuffer.size();
            params.pMetricsBuffer = metricsBuffer.data();
            CUPTI_API_CALL(cuptiProfilerHostGetMetricsInSinglePassSet(&params));
        }

        metricsList.resize(metricsIndices.size());
        for (uint32_t i = 0; i < metricsIndices.size(); i++)
        {
            const auto metrixIndex = metricsIndices[i];
            metricsList[i] = std::string((const char*)&metricsBuffer[metrixIndex]);
        }
        return CUPTI_SUCCESS;
    }

    CUptiResult GetSupportedBaseMetrics(std::vector<std::string>& metricsList)
    {
        for (size_t metricTypeIndex = 0; metricTypeIndex < CUPTI_METRIC_TYPE__COUNT; ++metricTypeIndex)
        {
            CUpti_Profiler_Host_GetBaseMetrics_Params getBaseMetricsParams {CUpti_Profiler_Host_GetBaseMetrics_Params_STRUCT_SIZE};
            getBaseMetricsParams.pHostObject = m_pHostObject;
            getBaseMetricsParams.metricType = (CUpti_MetricType)metricTypeIndex;
            CUPTI_API_CALL(cuptiProfilerHostGetBaseMetrics(&getBaseMetricsParams));

            for (size_t metricIndex = 0; metricIndex < getBaseMetricsParams.numMetrics; ++metricIndex)
            {
                metricsList.push_back(getBaseMetricsParams.ppMetricNames[metricIndex]);
            }
        }
        return CUPTI_SUCCESS;
    }

    CUptiResult GetMetricProperties(const std::string& metricName, CUpti_MetricType& metricType, std::string& metricDescription)
    {
        CUpti_Profiler_Host_GetMetricProperties_Params getMetricPropertiesParams {CUpti_Profiler_Host_GetMetricProperties_Params_STRUCT_SIZE};
        getMetricPropertiesParams.pHostObject = m_pHostObject;
        getMetricPropertiesParams.pMetricName = metricName.c_str();
        CUPTI_API_CALL(cuptiProfilerHostGetMetricProperties(&getMetricPropertiesParams));
        metricType = getMetricPropertiesParams.metricType;
        metricDescription = getMetricPropertiesParams.pDescription;
        return CUPTI_SUCCESS;
    }

    CUptiResult GetSubMetrics(const std::string& metricName, std::vector<std::string>& subMetricsList)
    {
        CUpti_MetricType metricType;
        std::string metricDescription;
        CUPTI_API_CALL(GetMetricProperties(metricName, metricType, metricDescription));

        CUpti_Profiler_Host_GetSubMetrics_Params getSubMetricsParams {CUpti_Profiler_Host_GetSubMetrics_Params_STRUCT_SIZE};
        getSubMetricsParams.pHostObject = m_pHostObject;
        getSubMetricsParams.pMetricName = metricName.c_str();
        getSubMetricsParams.metricType = metricType;
        CUPTI_API_CALL(cuptiProfilerHostGetSubMetrics(&getSubMetricsParams));

        for (size_t subMetricIndex = 0; subMetricIndex < getSubMetricsParams.numOfSubmetrics; ++subMetricIndex)
        {
            subMetricsList.push_back(getSubMetricsParams.ppSubMetrics[subMetricIndex]);
        }
        return CUPTI_SUCCESS;
    }

    void ExportPerfettoTrace(const char* fileName)
    {
        if (m_samplerRanges.size() <= 1)
        {
            std::cerr << "Warning: No valid samples to export to Perfetto trace." << std::endl;
            return;
        }

        auto escapeJsonString = [](const std::string& input) {
            std::string output;
            output.reserve(input.size() + input.size() / 8); // reserve extra space

            for (char c : input) {
                switch (c) {
                    case '"': output += "\\\""; break;
                    case '\\': output += "\\\\"; break;
                    case '\b': output += "\\b"; break;
                    case '\f': output += "\\f"; break;
                    case '\n': output += "\\n"; break;
                    case '\r': output += "\\r"; break;
                    case '\t': output += "\\t"; break;
                    default:
                        if (c < 0x20) {
                            output += "\\u";
                            char buf[5];
                            snprintf(buf, sizeof(buf), "%04x", (unsigned char)c);
                            output += buf;
                        } else {
                            output += c;
                        }
                        break;
                }
            }
            return output;
        };

        std::ofstream outputFile(fileName);
        if (!outputFile.is_open()) {
            std::cerr << "Error: Failed to open file for writing: " << fileName << std::endl;
            return;
        }

        try {
            std::unordered_map<std::string, int> metricTidMap;
            int nextTid = 1;

#if defined(_WIN32)
            int pid = static_cast<int>(_getpid());
#else
            int pid = static_cast<int>(getpid());
#endif

            // Start JSON structure
            outputFile << "{\n  \"traceEvents\": [\n";

            bool firstEvent = true;

            // Skip the first sample as it sometimes contains anomalous values
            auto startIdx = m_samplerRanges.empty() ? 0 : 1;
            uint64_t baseTimestamp = m_samplerRanges[startIdx].startTimestamp;

            for (size_t i = startIdx; i < m_samplerRanges.size(); ++i)
            {
                const auto& samplerRange = m_samplerRanges[i];

                // Skip invalid samples
                if (samplerRange.startTimestamp == UINT64_MAX ||
                    samplerRange.endTimestamp == UINT64_MAX ||
                    samplerRange.endTimestamp < samplerRange.startTimestamp) {
                    continue;
                }

                // Add comma separator between events
                if (!firstEvent) {
                    outputFile << ",\n";
                }
                firstEvent = false;

                // Adjust sample index for display
                size_t displaySampleIdx = i - startIdx + 1;

                // Main duration event for the sample
                outputFile << "    {\n";
                outputFile << "      \"name\": \"Sample " << displaySampleIdx << "\",\n";
                outputFile << "      \"ph\": \"X\",\n";
                outputFile << "      \"ts\": " << (samplerRange.startTimestamp - baseTimestamp) << ",\n";
                outputFile << "      \"dur\": " << (samplerRange.endTimestamp - samplerRange.startTimestamp) << ",\n";
                outputFile << "      \"pid\": " << pid << ",\n";
                outputFile << "      \"tid\": 0,\n";
                outputFile << "      \"args\": {\n";
                outputFile << "        \"sample_index\": " << displaySampleIdx << ",\n";
                outputFile << "        \"start_timestamp\": " << samplerRange.startTimestamp << ",\n";
                outputFile << "        \"end_timestamp\": " << samplerRange.endTimestamp;

                // Add metrics to args
                if (!samplerRange.metricValues.empty()) {
                    outputFile << ",\n        \"metrics\": {\n";
                    bool firstMetric = true;
                    for (const auto& metric : samplerRange.metricValues) {
                        if (!firstMetric) {
                            outputFile << ",\n";
                        }
                        firstMetric = false;

                        std::string metricName = metric.first;
                        std::replace(metricName.begin(), metricName.end(), '.', '(');
                        metricName += ")";

                        outputFile << "          \"" << escapeJsonString(metricName) << "\": " << metric.second;
                    }
                    outputFile << "\n        }";
                }

                outputFile << "\n      }\n    }";

                // Add counter events for each metric
                for (const auto& metric : samplerRange.metricValues) {
                    int metricTid = 0;
                    if (metricTidMap.count(metric.first) == 0) {
                        metricTidMap[metric.first] = nextTid++;
                    }
                    metricTid = metricTidMap[metric.first];

                    outputFile << ",\n    {\n";
                    outputFile << "      \"name\": \"" << escapeJsonString(metric.first) << "\",\n";
                    outputFile << "      \"ph\": \"C\",\n";
                    outputFile << "      \"ts\": " << (samplerRange.startTimestamp - baseTimestamp) << ",\n";
                    outputFile << "      \"pid\": " << pid << ",\n";
                    outputFile << "      \"tid\": " << metricTid << ",\n";
                    outputFile << "      \"args\": {\n";
                    outputFile << "        \"value\": " << metric.second << "\n";
                    outputFile << "      }\n    }";
                }
            }

            // Close JSON structure
            outputFile << "\n  ]\n}";

            std::cout << "Perfetto trace JSON written to: " << fileName << std::endl;
            std::cout << "Upload this file to https://ui.perfetto.dev/ for visualization" << std::endl;

        } catch (const std::exception& e) {
            std::cerr << "Error writing Perfetto JSON: " << e.what() << std::endl;
        }

        outputFile.close();
    }

private:
    CUptiResult Initialize(std::vector<uint8_t>& counterAvailibilityImage)
    {
        CUpti_Profiler_Host_Initialize_Params hostInitializeParams = {CUpti_Profiler_Host_Initialize_Params_STRUCT_SIZE};
        hostInitializeParams.profilerType = CUPTI_PROFILER_TYPE_PM_SAMPLING;
        hostInitializeParams.pChipName = m_chipName.c_str();
        hostInitializeParams.pCounterAvailabilityImage = counterAvailibilityImage.data();
        CUPTI_API_CALL(cuptiProfilerHostInitialize(&hostInitializeParams));
        m_pHostObject = hostInitializeParams.pHostObject;
        return CUPTI_SUCCESS;
    }

    CUptiResult Deinitialize()
    {
        CUpti_Profiler_Host_Deinitialize_Params deinitializeParams = {CUpti_Profiler_Host_Deinitialize_Params_STRUCT_SIZE};
        deinitializeParams.pHostObject = m_pHostObject;
        CUPTI_API_CALL(cuptiProfilerHostDeinitialize(&deinitializeParams));
        return CUPTI_SUCCESS;
    }
};


class CuptiPmSampling
{
    CUpti_PmSampling_Object* m_pmSamplerObject = nullptr;

public:
    CuptiPmSampling() = default;
    ~CuptiPmSampling() = default;

    void SetUp(int deviceIndex)
    {
        std::cout << "CUDA Device Number: " << deviceIndex << "\n";
        CUdevice cuDevice;
        DRIVER_API_CALL(cuDeviceGet(&cuDevice, deviceIndex));

        {
            int computeCapabilityMajor = 0, computeCapabilityMinor = 0;
            DRIVER_API_CALL(cuDeviceGetAttribute(&computeCapabilityMajor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, cuDevice));
            DRIVER_API_CALL(cuDeviceGetAttribute(&computeCapabilityMinor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, cuDevice));
            std::cout << "Compute Capability of Device: " << computeCapabilityMajor << "." << computeCapabilityMinor << "\n";
            if (computeCapabilityMajor * 10 + computeCapabilityMinor < 75)
            {
                std::cerr << "Sample not supported as it requires compute capability >= 7.5\n";
                exit(0);
            }
        }

        CUPTI_API_CALL(InitializeProfiler());
    }

    void TearDown()
    {
        CUPTI_API_CALL(DeInitializeProfiler());
    }

    CUptiResult CreateCounterDataImage( uint64_t maxSamples, std::vector<const char*> metricsList, std::vector<uint8_t>& counterDataImage)
    {
        CUpti_PmSampling_GetCounterDataSize_Params getCounterDataSizeParams = {CUpti_PmSampling_GetCounterDataSize_Params_STRUCT_SIZE};
        getCounterDataSizeParams.pPmSamplingObject = m_pmSamplerObject;
        getCounterDataSizeParams.numMetrics = metricsList.size();
        getCounterDataSizeParams.pMetricNames = metricsList.data();
        getCounterDataSizeParams.maxSamples  = maxSamples;
        CUPTI_API_CALL(cuptiPmSamplingGetCounterDataSize(&getCounterDataSizeParams));

        counterDataImage.resize(getCounterDataSizeParams.counterDataSize);
        CUpti_PmSampling_CounterDataImage_Initialize_Params initializeParams {CUpti_PmSampling_CounterDataImage_Initialize_Params_STRUCT_SIZE};
        initializeParams.pPmSamplingObject = m_pmSamplerObject;
        initializeParams.counterDataSize = counterDataImage.size();
        initializeParams.pCounterData = counterDataImage.data();
        CUPTI_API_CALL(cuptiPmSamplingCounterDataImageInitialize(&initializeParams));
        return CUPTI_SUCCESS;
    }

    CUptiResult ResetCounterDataImage(std::vector<uint8_t>& counterDataImage)
    {
        CUpti_PmSampling_CounterDataImage_Initialize_Params initializeParams {CUpti_PmSampling_CounterDataImage_Initialize_Params_STRUCT_SIZE};
        initializeParams.pPmSamplingObject = m_pmSamplerObject;
        initializeParams.counterDataSize = counterDataImage.size();
        initializeParams.pCounterData = counterDataImage.data();
        CUPTI_API_CALL(cuptiPmSamplingCounterDataImageInitialize(&initializeParams));
        return CUPTI_SUCCESS;
    }

    CUptiResult EnablePmSampling(size_t devIndex)
    {
        CUpti_PmSampling_Enable_Params enableParams {CUpti_PmSampling_Enable_Params_STRUCT_SIZE};
        enableParams.deviceIndex = devIndex;
        CUPTI_API_CALL(cuptiPmSamplingEnable(&enableParams));
        m_pmSamplerObject = enableParams.pPmSamplingObject;
        return CUPTI_SUCCESS;
    }

    CUptiResult DisablePmSampling()
    {
        CUpti_PmSampling_Disable_Params disableParams {CUpti_PmSampling_Disable_Params_STRUCT_SIZE};
        disableParams.pPmSamplingObject = m_pmSamplerObject;
        CUPTI_API_CALL(cuptiPmSamplingDisable(&disableParams));
        return CUPTI_SUCCESS;
    }

    CUptiResult SetConfig(std::vector<uint8_t>& configImage, size_t hardwareBufferSize, uint64_t samplingInterval)
    {
        CUpti_PmSampling_SetConfig_Params setConfigParams = {CUpti_PmSampling_SetConfig_Params_STRUCT_SIZE};
        setConfigParams.pPmSamplingObject = m_pmSamplerObject;

        setConfigParams.configSize = configImage.size();
        setConfigParams.pConfig = configImage.data();

        setConfigParams.hardwareBufferSize = hardwareBufferSize;
        setConfigParams.samplingInterval = samplingInterval;

        setConfigParams.triggerMode = CUpti_PmSampling_TriggerMode::CUPTI_PM_SAMPLING_TRIGGER_MODE_GPU_SYSCLK_INTERVAL;
        CUPTI_API_CALL(cuptiPmSamplingSetConfig(&setConfigParams));
        return CUPTI_SUCCESS;
    }

    CUptiResult StartPmSampling()
    {
        CUpti_PmSampling_Start_Params startProfilingParams = {CUpti_PmSampling_Start_Params_STRUCT_SIZE};
        startProfilingParams.pPmSamplingObject = m_pmSamplerObject;
        CUPTI_API_CALL(cuptiPmSamplingStart(&startProfilingParams));
        return CUPTI_SUCCESS;
    }

    CUptiResult StopPmSampling()
    {
        CUpti_PmSampling_Stop_Params stopProfilingParams = {CUpti_PmSampling_Stop_Params_STRUCT_SIZE};
        stopProfilingParams.pPmSamplingObject = m_pmSamplerObject;
        CUPTI_API_CALL(cuptiPmSamplingStop(&stopProfilingParams));
        return CUPTI_SUCCESS;
    }

    CUptiResult DecodePmSamplingData(std::vector<uint8_t>& counterDataImage)
    {
        CUpti_PmSampling_DecodeData_Params decodeDataParams = {CUpti_PmSampling_DecodeData_Params_STRUCT_SIZE};
        decodeDataParams.pPmSamplingObject = m_pmSamplerObject;
        decodeDataParams.pCounterDataImage = counterDataImage.data();
        decodeDataParams.counterDataImageSize = counterDataImage.size();
        CUPTI_API_CALL(cuptiPmSamplingDecodeData(&decodeDataParams));
        return CUPTI_SUCCESS;
    }

    CUpti_PmSampling_Object* GetPmSamplerObject()
    {
        return m_pmSamplerObject;
    }

    static CUptiResult GetChipName(size_t deviceIndex, std::string& chipName)
    {
        CUpti_Profiler_Initialize_Params profilerInitializeParams = { CUpti_Profiler_Initialize_Params_STRUCT_SIZE };
        CUPTI_API_CALL(cuptiProfilerInitialize(&profilerInitializeParams));

        CUpti_Device_GetChipName_Params getChipNameParams = { CUpti_Device_GetChipName_Params_STRUCT_SIZE };
        getChipNameParams.deviceIndex = deviceIndex;
        CUPTI_API_CALL(cuptiDeviceGetChipName(&getChipNameParams));
        chipName = getChipNameParams.pChipName;
        return CUPTI_SUCCESS;
    }

    static CUptiResult GetCounterAvailabilityImage(size_t deviceIndex, std::vector<uint8_t>& counterAvailibilityImage)
    {
        CUpti_PmSampling_GetCounterAvailability_Params getCounterAvailabilityParams { CUpti_PmSampling_GetCounterAvailability_Params_STRUCT_SIZE };
        getCounterAvailabilityParams.deviceIndex = deviceIndex;
        CUPTI_API_CALL(cuptiPmSamplingGetCounterAvailability(&getCounterAvailabilityParams));

        counterAvailibilityImage.clear();
        counterAvailibilityImage.resize(getCounterAvailabilityParams.counterAvailabilityImageSize);
        getCounterAvailabilityParams.pCounterAvailabilityImage = counterAvailibilityImage.data();
        CUPTI_API_CALL(cuptiPmSamplingGetCounterAvailability(&getCounterAvailabilityParams));
        return CUPTI_SUCCESS;
    }

private:
    CUptiResult InitializeProfiler()
    {
        CUpti_Profiler_Initialize_Params profilerInitializeParams = { CUpti_Profiler_Initialize_Params_STRUCT_SIZE };
        CUPTI_API_CALL(cuptiProfilerInitialize(&profilerInitializeParams));
        return CUPTI_SUCCESS;
    }

    CUptiResult DeInitializeProfiler()
    {
        CUpti_Profiler_DeInitialize_Params profilerDeInitializeParams = { CUpti_Profiler_DeInitialize_Params_STRUCT_SIZE };
        CUPTI_API_CALL(cuptiProfilerDeInitialize(&profilerDeInitializeParams));
        return CUPTI_SUCCESS;
    }
};
