// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include "testPch.h"
#include "test.h"
#include "RawApiHelpers.h"

namespace ml = Microsoft::AI::MachineLearning;

void RunOnDevice(ml::learning_model& model, ml::learning_model_device& device, InputStrategy strategy) {
  const wchar_t input_name[] = L"data_0";
  const wchar_t output_name[] = L"softmaxout_1";

  std::unique_ptr<ml::learning_model_session> session = nullptr;
  WINML_EXPECT_NO_THROW(session = std::make_unique<ml::learning_model_session>(model, device));

  std::unique_ptr<ml::learning_model_binding> binding = nullptr;
  WINML_EXPECT_NO_THROW(binding = std::make_unique<ml::learning_model_binding>(*session.get()));

  auto input_shape = std::vector<ml::tensor_shape_type>{1, 3, 224, 224};
  auto input_data = std::vector<float>(1 * 3 * 224 * 224);
  auto output_shape = std::vector<ml::tensor_shape_type>{1, 1000, 1, 1};

  std::iota(begin(input_data), end(input_data), 0.f);

  if (strategy == InputStrategy::CopyInputs) {
    WINML_EXPECT_HRESULT_SUCCEEDED(binding->bind<float>(
      input_name, _countof(input_name) - 1, input_shape.data(), input_shape.size(), input_data.data(), input_data.size()
    ));
  } else if (strategy == InputStrategy::BindAsReference) {
    WINML_EXPECT_HRESULT_SUCCEEDED(binding->bind_as_reference<float>(
      input_name, _countof(input_name) - 1, input_shape.data(), input_shape.size(), input_data.data(), input_data.size()
    ));
  } else if (strategy == InputStrategy::BindWithMultipleReferences) {
    size_t channel_size = 224 * 224;
    auto channel_buffers_sizes = std::vector<size_t>{channel_size, channel_size, channel_size};

    auto channel_buffers_pointers = std::vector<float*>{
      &input_data.at(0),
      &input_data.at(0) + channel_buffers_sizes[0],
      &input_data.at(0) + channel_buffers_sizes[0] + +channel_buffers_sizes[1]};

    WINML_EXPECT_HRESULT_SUCCEEDED(binding->bind_as_references<float>(
      input_name,
      _countof(input_name) - 1,
      channel_buffers_pointers.data(),
      channel_buffers_sizes.data(),
      channel_buffers_sizes.size()
    ));
  }

  WINML_EXPECT_HRESULT_SUCCEEDED(
    binding->bind<float>(output_name, _countof(output_name) - 1, output_shape.data(), output_shape.size())
  );

  ml::learning_model_results results = session->evaluate(*binding.get());

  float* p_buffer = nullptr;
  size_t buffer_size = 0;
  WINML_EXPECT_HRESULT_SUCCEEDED(
    0 == results.get_output(output_name, _countof(output_name) - 1, reinterpret_cast<void**>(&p_buffer), &buffer_size)
  );
}