#include "ClipApp.h"

NAMESPACE_MAS_BEGIN

NAMESPACE_CLIP_BEGIN

ClipApp* ClipApp::_instance = nullptr;
tzc::Mutex ClipApp::_mutex;

ClipApp* ClipApp::Instance()
{
    if (!_instance)
	{
        _mutex.Lock();
        if (!_instance)
		{
            _instance = new ClipApp();
            TZLogInfo("ClipApp Instance created~~~");
        }
        _mutex.Unlock();
    }

    ++(_instance->m_usecnt);
    return _instance;
}

void ClipApp::DestroyInstance()
{
    tzc::ScopedLock lock(_mutex);
    if ((_instance->m_usecnt) < 0)
	{
        TZLogWarn("ClipApp Instance does not exist!!!");
		_instance->m_usecnt = 0;
    }
    else if ((_instance->m_usecnt) == 0)
	{
        TZ_delete(_instance);
        TZLogInfo("ClipApp Instance destroyed~~~");
    }
    else
	{
        TZLogWarn("ClipApp still in use by %d objects!", _instance->m_usecnt);
		TZ_delete(_instance);
		TZLogInfo("ClipApp Instance destroyed~~~");
    }
}

TZ_INT ClipApp::Initialize(const std::string& initParam)
{
    if (m_inited)
	{
        TZLogInfo("ClipApp has been initialized~~~");
        return MEC_OK;
    }

    ClipBuild buildParam;
	ClipBuild::fromJson(initParam, buildParam);

    // Create ONNX Runtime environment
    m_env = new Ort::Env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, "ClipApp");

    // Set session options
    m_sessionOptions.SetIntraOpNumThreads(1);

    // Configure CUDA execution provider if available
    auto providers = Ort::GetAvailableProviders();
    auto cudaAvailable = std::find(providers.begin(), providers.end(), "CUDAExecutionProvider");
    if (cudaAvailable != providers.end())
    {
        OrtCUDAProviderOptions cudaOptions;
        cudaOptions.device_id = buildParam.gpu_id;
        m_sessionOptions.AppendExecutionProvider_CUDA(cudaOptions);
    }

    // Create ONNX session
    m_session = new Ort::Session(*m_env, buildParam.model_path.c_str(), m_sessionOptions);

    m_inited = true;
    return MEC_OK;
}

TZ_INT ClipApp::Dispose()
{
    if (!m_inited)
	{
        return MEC_NOT_INITED;
    }

    _mutex.Lock();

    --(m_usecnt);
    TZLogInfo("ClipApp usecnt -1, now = %d~~~", m_usecnt);
    TZ_BOOL empty = (m_usecnt == 0);

    _mutex.Unlock();

    if (empty)
	{
        m_inited = FALSE;
        TZLogInfo("ClipApp Dispose~~~");
        this->DestroyInstance();
    }

    return MEC_OK;
}

// Preprocess the input image
TZ_INT ClipApp::PreprocessImage(cv::Mat& input, std::vector<TZ_FLOAT>& outputTensor)
{
    if (input.empty())
    {
        TZLogError("Input image is empty!!!");
        return MEC_FAILED;
    }

    // Resize and normalize the image
    cv::Mat resized;
    cv::resize(input, resized, cv::Size(224, 224));
    resized.convertTo(resized, CV_32F, 1.0 / 255);
    cv::cvtColor(resized, resized, cv::COLOR_BGR2RGB);

    // Convert to CHW format
    outputTensor.clear();
    for (int c = 0; c < 3; ++c)
    {
        for (int h = 0; h < resized.rows; ++h)
        {
            for (int w = 0; w < resized.cols; ++w)
            {
                outputTensor.push_back(resized.at<cv::Vec3f>(h, w)[c]);
            }
        }
    }

    return MEC_OK;
}

TZ_INT ClipApp::DoDetect(cv::Mat& input, ClipDetectResult& detRes)
{
    // Preprocess the image
    std::vector<TZ_FLOAT> inputTensorValues;
    if (PreprocessImage(input, inputTensorValues) != 0)
    { 
        TZLogError("Failed to preprocess input image!!!");
        return MEC_FAILED;
    }

    // Get model input/output information dynamically
    auto inputDims = m_session->GetInputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();
    auto outputDims = m_session->GetOutputTypeInfo(0).GetTensorTypeAndShapeInfo().GetShape();

    // Fix dynamic dimensions (-1)
    inputDims[0] = 1;  // Batch size
    outputDims[0] = 1;  // Batch size

    // Create input tensor
    Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value inputTensor = Ort::Value::CreateTensor<TZ_FLOAT>(
        memoryInfo, inputTensorValues.data(), inputTensorValues.size(),
        inputDims.data(), inputDims.size());

    // Allocate output tensor
    TZ_INT outputTensorSize = 1;
    for (auto dim : outputDims)
    {
        outputTensorSize *= dim;
    }
    std::vector<TZ_FLOAT> outputTensorValues(outputTensorSize);
    Ort::Value outputTensor = Ort::Value::CreateTensor<TZ_FLOAT>(
        memoryInfo, outputTensorValues.data(), outputTensorValues.size(),
        outputDims.data(), outputDims.size());

    // Get input/output names
    Ort::AllocatorWithDefaultOptions allocator;
    auto inputNamePtr = m_session->GetInputNameAllocated(0, allocator);
    auto outputNamePtr = m_session->GetOutputNameAllocated(0, allocator);

    const TZ_Int8* inputName = inputNamePtr.get();
    const TZ_Int8* outputName = outputNamePtr.get();

    // Run inference
    m_session->Run(Ort::RunOptions{nullptr}, &inputName, &inputTensor, 1,
                   &outputName, &outputTensor, 1);

    // Convert outputTensorValues (TZ_FLOAT vector) to detRes.output_data (vector<OutputData>)
    detRes.output_data.clear();  // Clear any existing data
    for (const auto& val : outputTensorValues)
    {
        OutputData output;
        output.val = val;  // Populate OutputData with the value
        detRes.output_data.push_back(output);  // Add to the result vector
    }

    return MEC_OK;
}

ClipApp::ClipApp()
    : m_inited(false),
    m_env(nullptr),
    m_session(nullptr) {}

ClipApp::~ClipApp()
{
    Dispose();
}

NAMESPACE_CLIP_END

NAMESPACE_MAS_END