#include <onnxruntime/cpu_provider_factory.h>
#include <onnxruntime/onnxruntime_cxx_api.h>
#include <png.h>

#include <algorithm>
#include <array>
#include <cassert>
#include <cmath>
#include <iostream>
#include <span>
#include <string>
#include <vector>

// https://onnxruntime.ai/docs/api/c/namespace_ort.html
// https://github.com/microsoft/onnxruntime-inference-examples/

class ImageData;

class OnnxMnist
{
public:
    OnnxMnist() noexcept = default;

    auto run(ImageData& image) noexcept -> std::size_t;

    static constexpr const int kWidth{28};
    static constexpr const int kHeight{28};

private:
    // From session we can get the input and output count
    Ort::Session session_{Ort::Env{}, "model.onnx", Ort::SessionOptions{nullptr}};
};

class ImageData
{
public:
    explicit ImageData(std::string_view input_file);

    auto getData() noexcept -> std::span<float> { return data_; }
    auto getWidth() const noexcept { return width_; }
    auto getHeight() const noexcept { return height_; }

private:
    std::size_t width_{0};
    std::size_t height_{0};

    std::vector<float> data_;
};

template<typename T>
void softmax(std::span<T> input)
{
    auto row_max = *std::ranges::max_element(input);
    std::vector<float> result(input.size());
    float sum{0.0};
    for (std::size_t i = 0; i != input.size(); ++i) {
        sum += result[i] = std::exp(input[i] - row_max);
    }
    for (std::size_t i = 0; i != input.size(); ++i) {
        input[i] = result[i] / sum;
    }
}

auto main(int argc, char* argv[]) -> int
{
    if (argc < 2) {
        std::cout << "usage: <input_file>" << std::endl;
        return -1;
    }

    std::string_view filename{argv[1]};  // NOLINT
    ImageData image{filename};

    if (image.getWidth() != OnnxMnist::kWidth || image.getHeight() != OnnxMnist::kHeight) {
        std::cout << "Resize to image to 28x28" << std::endl;
        return -1;
    }

    OnnxMnist model;
    auto result = model.run(image);
    std::cout << "Result : " << result << std::endl;
    return 0;
}

auto OnnxMnist::run(ImageData& image) noexcept -> std::size_t
{
    const char* kInputNames[]{"Input3"};             // NOLINT
    const char* kOutputNames[]{"Plus214_Output_0"};  // NOLINT

    constexpr std::array<int64_t, 4> kInputShape{1, 1, kWidth, kHeight};
    constexpr std::array<int64_t, 2> kOutputShape{1, 10};

    const auto memory_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU);

    auto image_data = image.getData();
    Ort::Value const input_tensor = Ort::Value::CreateTensor<float>(
        memory_info, image_data.data(), image_data.size(), kInputShape.data(), kInputShape.size());

    std::array<float, kOutputShape[1]> results_{};
    auto output_tensor = Ort::Value::CreateTensor<float>(
        memory_info, results_.data(), results_.size(), kOutputShape.data(), kOutputShape.size());

    const Ort::RunOptions run_options;
    session_.Run(run_options, kInputNames, &input_tensor, 1, kOutputNames, &output_tensor, 1);

    softmax<float>(results_);
    for (std::size_t i = 0; i < results_.size(); i++) {
        std::cout << i << " : " << results_[i] << std::endl;
    }

    auto result = std::distance(results_.begin(), std::max_element(results_.begin(), results_.end()));

    return result;
}

ImageData::ImageData(std::string_view input_file)
{
    // The control structure used by libpng
    // Initialize the 'png_image' structure.
    png_image image;
    memset(&image, 0, (sizeof image));
    image.version = PNG_IMAGE_VERSION;
    if (png_image_begin_read_from_file(&image, input_file.data()) == 0) {
        return;
    }
    image.format = PNG_FORMAT_BGR;

    const std::size_t input_data_length = PNG_IMAGE_SIZE(image);  // NOLINT
    std::vector<png_byte> buffer(input_data_length, 0);
    if (png_image_finish_read(&image, nullptr /*background*/, buffer.data(), 0 /*row_stride*/, nullptr /*colormap*/) ==
        0) {
        return;
    }

    width_ = image.width;
    height_ = image.height;

    const auto stride = height_ * width_;
    data_.resize(height_ * width_);

    for (std::size_t i = 0; i < stride; ++i) {
        data_[i] = buffer[i * 3] == 0 ? 1.0F : 0.0F;
    }
}