// This file is part of the AliceVision project.
// Copyright (c) 2023 AliceVision contributors.
// This Source Code Form is subject to the terms of the Mozilla Public License,
// v. 2.0. If a copy of the MPL was not distributed with this file,
// You can obtain one at https://mozilla.org/MPL/2.0/.

#include "segmentation.hpp"

#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_CUDA)
    #include <cuda_runtime.h>
#endif

#include <aliceVision/image/all.hpp>
#include <aliceVision/image/imageAlgo.hpp>
#include <aliceVision/numeric/numeric.hpp>

namespace aliceVision {
namespace segmentation {

void imageToPlanes(std::vector<float>& output, const image::Image<image::RGBfColor>::Base& source)
{
    size_t planeSize = source.rows() * source.cols();

    output.resize(planeSize * 3);

    float* planeR = output.data();
    float* planeG = planeR + planeSize;
    float* planeB = planeG + planeSize;

    size_t pos = 0;
    for (int i = 0; i < source.rows(); i++)
    {
        for (int j = 0; j < source.cols(); j++)
        {
            const image::RGBfColor& rgb = source(i, j);
            planeR[pos] = rgb.r();
            planeG[pos] = rgb.g();
            planeB[pos] = rgb.b();

            pos++;
        }
    }
}

bool Segmentation::initialize()
{
    const auto& api = Ort::GetApi();

    _ortEnvironment = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "aliceVision-imageSegmentation");

    Ort::SessionOptions ortSessionOptions;

    // this is false if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_ONNX_GPU) is false
    if (_parameters.useGpu)
    {
// Disable for compilation purpose if needed
#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_ONNX_GPU)
        OrtCUDAProviderOptionsV2* cuda_options = nullptr;
        api.CreateCUDAProviderOptions(&cuda_options);
        api.SessionOptionsAppendExecutionProvider_CUDA_V2(static_cast<OrtSessionOptions*>(ortSessionOptions), cuda_options);
        api.ReleaseCUDAProviderOptions(cuda_options);

    #if defined(_WIN32) || defined(_WIN64)
        std::wstring modelWeights(_parameters.modelWeights.begin(), _parameters.modelWeights.end());
        _ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, modelWeights.c_str(), ortSessionOptions);
    #else
        _ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, _parameters.modelWeights.c_str(), ortSessionOptions);
    #endif

        _output.resize(_parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
#endif
    }
    else
    {
#if defined(_WIN32) || defined(_WIN64)
        std::wstring modelWeights(_parameters.modelWeights.begin(), _parameters.modelWeights.end());
        _ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, modelWeights.c_str(), ortSessionOptions);
#else
        _ortSession = std::make_unique<Ort::Session>(*_ortEnvironment, _parameters.modelWeights.c_str(), ortSessionOptions);
#endif
    }

    return true;
}

bool Segmentation::processImage(image::Image<IndexT>& labels, const image::Image<image::RGBfColor>& source)
{
    // Todo : handle orientation and small images smaller than model input

    // Compute the optimal resized size such that:
    // - at last one dimension fit the model
    // - both dimensions are larger or equal than the model dimensions
    int resizedHeight = 0;
    int resizedWidth = 0;
    if (source.height() < source.width())
    {
        resizedWidth =
          static_cast<int>(static_cast<double>(source.width()) * static_cast<double>(_parameters.modelHeight) / static_cast<double>(source.height()));
        if (resizedWidth < _parameters.modelWidth)
        {
            resizedWidth = _parameters.modelWidth;
            resizedHeight = static_cast<int>(static_cast<double>(resizedWidth) * static_cast<double>(_parameters.modelHeight) /
                                             static_cast<double>(_parameters.modelWidth));
        }
        else
        {
            resizedHeight = _parameters.modelHeight;
        }
    }
    else
    {
        resizedHeight =
          static_cast<int>(static_cast<double>(source.height()) * static_cast<double>(_parameters.modelWidth) / static_cast<double>(source.width()));
        if (resizedHeight < _parameters.modelHeight)
        {
            resizedHeight = _parameters.modelHeight;
            resizedWidth = static_cast<int>(static_cast<double>(resizedHeight) * static_cast<double>(_parameters.modelWidth) /
                                            static_cast<double>(_parameters.modelHeight));
        }
        else
        {
            resizedWidth = _parameters.modelWidth;
        }
    }

    // Resize image
    image::Image<image::RGBfColor> resized;
    imageAlgo::resizeImage(resizedWidth, resizedHeight, source, resized);

    // Normalize image to fit model statistics
    for (int i = 0; i < resizedHeight; i++)
    {
        for (int j = 0; j < resizedWidth; j++)
        {
            image::RGBfColor value = resized(i, j);
            resized(i, j) = (value - _parameters.center) * _parameters.scale;
        }
    }

    image::Image<IndexT> resizedLabels;
    if (!tiledProcess(resizedLabels, resized))
    {
        return false;
    }

    imageAlgo::resampleImage(source.width(), source.height(), resizedLabels, labels, false);

    return true;
}

bool Segmentation::tiledProcess(image::Image<IndexT>& labels, const image::Image<image::RGBfColor>& source)
{
    // Compute the theoretical tiles count
    int cwidth = divideRoundUp(source.width(), _parameters.modelWidth);
    int cheight = divideRoundUp(source.height(), _parameters.modelHeight);

    image::Image<ScoredLabel> scoredLabels(source.width(), source.height(), true, {0, 0.0f});

    // Loop over tiles
    for (int i = 0; i < cheight; i++)
    {
        // Compute starting point with overlap on previous
        int y = std::max(0, int(i * _parameters.modelHeight - _parameters.overlapRatio * _parameters.modelHeight));
        int ly = y + _parameters.modelHeight;

        // If we are on the end border, shift on the other side
        int shifty = source.height() - ly;
        if (shifty < 0)
        {
            y = std::max(0, y + shifty);
        }

        for (int j = 0; j < cwidth; j++)
        {
            // Compute starting point with overlap on previous
            int x = std::max(0, int(j * _parameters.modelWidth - _parameters.overlapRatio * _parameters.modelWidth));
            int lx = x + _parameters.modelWidth;

            // If we are on the end border, shift on the other side
            int shiftx = source.width() - lx;
            if (shiftx < 0)
            {
                x = std::max(0, x + shiftx);
            }

            // x and y contains the position of the tile in the input image
            auto& block = source.block(y, x, _parameters.modelHeight, _parameters.modelWidth);

            // Compute tile
            image::Image<ScoredLabel> tileLabels(_parameters.modelWidth, _parameters.modelHeight, true, {0, 0.0f});

            if (_parameters.useGpu)
            {
                processTileGPU(tileLabels, block);
            }
            else
            {
                processTile(tileLabels, block);
            }

            // Update the global labeling
            mergeLabels(scoredLabels, tileLabels, x, y);
        }
    }

    labels = scoredLabels.cast<IndexT>();

    return true;
}

bool Segmentation::mergeLabels(image::Image<ScoredLabel>& labels, image::Image<ScoredLabel>& tileLabels, int tileX, int tileY)
{
    for (int i = 0; i < tileLabels.height(); i++)
    {
        int y = i + tileY;
        for (int j = 0; j < tileLabels.width(); j++)
        {
            int x = j + tileX;

            if (tileLabels(i, j).score > labels(y, x).score)
            {
                labels(y, x) = tileLabels(i, j);
            }
        }
    }

    return true;
}

bool Segmentation::labelsFromOutputTensor(image::Image<ScoredLabel>& labels, Ort::Value& modelOutput)
{
    for (int outputY = 0; outputY < _parameters.modelHeight; outputY++)
    {
        for (int outputX = 0; outputX < _parameters.modelWidth; outputX++)
        {
            int maxClasse = 0;
            int maxVal = 0;

            for (int classe = 0; classe < _parameters.classes.size(); classe++)
            {
                const std::vector<int64_t> coords = {0,classe,outputY,outputX};
                const float val = modelOutput.At<float>(coords);
                if (val > maxVal)
                {
                    maxVal = val;
                    maxClasse = classe;
                }
            }

            labels(outputY, outputX) = {static_cast<IndexT>(maxClasse), static_cast<float>(maxVal)};
        }
    }

    return true;
}

bool Segmentation::processTile(image::Image<ScoredLabel>& labels, const image::Image<image::RGBfColor>::Base& source)
{
    ALICEVISION_LOG_TRACE("Process tile using cpu");
    Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);

    std::vector<const char*> inputNames{"input"};
    std::vector<const char*> outputNames{"output"};
    std::vector<int64_t> inputDimensions = {1, 3, _parameters.modelHeight, _parameters.modelWidth};

    std::vector<float> transformedInput;
    imageToPlanes(transformedInput, source);

    Ort::Value inputTensors =
      Ort::Value::CreateTensor<float>(memInfo, transformedInput.data(), transformedInput.size(), inputDimensions.data(), inputDimensions.size());

    std::vector<Ort::Value> outTensor;

    try
    {
        outTensor = _ortSession->Run(Ort::RunOptions{nullptr}, inputNames.data(), &inputTensors, 1, outputNames.data(), 1);
    }
    catch (const Ort::Exception& exception)
    {
        ALICEVISION_LOG_ERROR("ERROR running model inference: " << exception.what());
        return false;
    }

    if (!labelsFromOutputTensor(labels, outTensor[0]))
    {
        return false;
    }

    std::vector<float> output(_parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);
    auto *outTData = outTensor.front().GetTensorMutableData<float>();
    output.assign(outTData, outTData + _parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);

    return true;
}

bool Segmentation::processTileGPU(image::Image<ScoredLabel>& labels, const image::Image<image::RGBfColor>::Base& source)
{
    ALICEVISION_LOG_TRACE("Process tile using gpu");
#if ALICEVISION_IS_DEFINED(ALICEVISION_HAVE_CUDA)
    Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault);

    std::vector<const char*> inputNames{"input"};
    std::vector<const char*> outputNames{"output"};
    std::vector<int64_t> inputDimensions = {1, 3, _parameters.modelHeight, _parameters.modelWidth};

    std::vector<float> transformedInput;
    imageToPlanes(transformedInput, source);

    std::vector<Ort::Value> inputTensors;
    inputTensors.emplace_back(Ort::Value::CreateTensor<float>(memInfo,
                                                              transformedInput.data(),
                                                              transformedInput.size(),
                                                              inputDimensions.data(),
                                                              inputDimensions.size()));

    std::vector<Ort::Value> outTensor;

    try
    {
        outTensor = _ortSession->Run(Ort::RunOptions{nullptr}, inputNames.data(), inputTensors.data(), 1, outputNames.data(), 1);
    }
    catch (const Ort::Exception& exception)
    {
        ALICEVISION_LOG_ERROR("ERROR running model inference: " << exception.what());
        return false;
    }

    if (!labelsFromOutputTensor(labels, outTensor[0]))
    {
        return false;
    }

    auto *outTData = outTensor.front().GetTensorMutableData<float>();
    _output.assign(outTData, outTData + _parameters.classes.size() * _parameters.modelHeight * _parameters.modelWidth);

#endif

    return true;
}

}  // namespace segmentation
}  // namespace aliceVision