/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include <executorch/backends/vulkan/test/utils/test_utils.h>

#include <executorch/runtime/core/exec_aten/exec_aten.h>

#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>

#include <cassert>
#include <random>

using namespace vkcompute;

//
// Operator Recording Functions
//

void record_nchw_to_buffer_op(
    api::Context* const context,
    vkapi::VulkanBuffer& src_buffer,
    api::vTensor& v_dst) {
  vkapi::PipelineBarrier pipeline_barrier{};

  context->submit_compute_job(
      get_nchw_to_tensor_shader(v_dst),
      pipeline_barrier,
      {uint32_t(v_dst.numel()), 1, 1},
      {64, 1, 1},
      {},
      VK_NULL_HANDLE,
      0,
      v_dst.buffer(
          pipeline_barrier,
          vkapi::PipelineStage::COMPUTE,
          vkapi::MemoryAccessType::WRITE),
      src_buffer,
      v_dst.sizes_ubo(),
      v_dst.strides_ubo(),
      v_dst.numel_ubo());
}

void record_buffer_to_nchw_op(
    api::Context* const context,
    api::vTensor& v_src,
    vkapi::VulkanBuffer& dst_buffer) {
  vkapi::PipelineBarrier pipeline_barrier{};
  context->submit_compute_job(
      get_tensor_to_nchw_shader(v_src),
      pipeline_barrier,
      {uint32_t(v_src.numel()), 1, 1},
      {64, 1, 1},
      {},
      VK_NULL_HANDLE,
      0,
      dst_buffer,
      v_src.buffer(pipeline_barrier, vkapi::PipelineStage::COMPUTE),
      v_src.sizes_ubo(),
      v_src.strides_ubo(),
      v_src.numel_ubo());
}

void record_nchw_to_image_op(
    api::Context* const context,
    vkapi::VulkanBuffer& src_buffer,
    api::vTensor& v_dst) {
  vkapi::PipelineBarrier pipeline_barrier{};
  vkapi::SpecVarList specialization_constants = {v_dst.hashed_layout()};

  context->submit_compute_job(
      get_nchw_to_tensor_shader(
          v_dst, context->adapter_ptr()->has_full_int8_buffers_support()),
      pipeline_barrier,
      v_dst.logical_limits(),
      adaptive_work_group_size(v_dst.logical_limits()),
      specialization_constants,
      VK_NULL_HANDLE,
      0,
      v_dst.image(
          pipeline_barrier,
          vkapi::PipelineStage::COMPUTE,
          vkapi::MemoryAccessType::WRITE),
      src_buffer,
      v_dst.sizes_ubo());
}

void record_image_to_nchw_op(
    api::Context* const context,
    api::vTensor& v_src,
    vkapi::VulkanBuffer& dst_buffer) {
  vkapi::PipelineBarrier pipeline_barrier{};
  vkapi::SpecVarList specialization_constants = {v_src.hashed_layout()};

  context->submit_compute_job(
      get_tensor_to_nchw_shader(v_src),
      pipeline_barrier,
      v_src.logical_limits(),
      adaptive_work_group_size(v_src.logical_limits()),
      specialization_constants,
      VK_NULL_HANDLE,
      0,
      dst_buffer,
      v_src.image(pipeline_barrier, vkapi::PipelineStage::COMPUTE),
      v_src.sizes_ubo());
}

void record_bitw8_image_to_nchw_nobitw8buffer_op(
    api::Context* const context,
    api::vTensor& v_src,
    api::StagingBuffer& dst_buffer) {
  vkapi::PipelineBarrier pipeline_barrier{};
  uint32_t buffer_len = utils::safe_downcast<uint32_t>(dst_buffer.numel() / 4);
  utils::uvec3 global_wg_size = {buffer_len, 1, 1};

  std::string kernel_name = "bitw8_image_to_nchw_nobitw8buffer";
  add_storage_type_suffix(kernel_name, v_src);
  add_dtype_suffix(kernel_name, v_src);

  context->submit_compute_job(
      VK_KERNEL_FROM_STR(kernel_name),
      pipeline_barrier,
      global_wg_size,
      adaptive_work_group_size(global_wg_size),
      {v_src.hashed_layout()},
      VK_NULL_HANDLE,
      0,
      dst_buffer.buffer(),
      v_src.image(pipeline_barrier, vkapi::PipelineStage::COMPUTE),
      v_src.sizes_ubo(),
      v_src.numel_ubo());
}

void record_conv2d_prepack_weights_op(
    api::Context* const context,
    vkapi::VulkanBuffer& src_buffer,
    api::vTensor& v_dst,
    const std::vector<int64_t>& original_sizes,
    const bool transposed) {
  vkapi::PipelineBarrier pipeline_barrier{};

  std::string kernel_name;
  if (transposed) {
    kernel_name = "conv_transpose2d";
  } else {
    kernel_name = "conv2d";
  }
  kernel_name += "_prepack_weights";
  add_dtype_suffix(kernel_name, v_dst);
  vkapi::ShaderInfo shader = VK_KERNEL_FROM_STR(kernel_name);

  api::ParamsBuffer original_sizes_ubo(
      context, utils::make_ivec4(original_sizes, /*reverse = */ true));

  vkapi::SpecVarList specialization_constants = {};
  context->submit_compute_job(
      shader,
      pipeline_barrier,
      v_dst.logical_limits(),
      adaptive_work_group_size(v_dst.logical_limits()),
      specialization_constants,
      VK_NULL_HANDLE,
      0,
      v_dst.image(
          pipeline_barrier,
          vkapi::PipelineStage::COMPUTE,
          vkapi::MemoryAccessType::WRITE),
      src_buffer,
      v_dst.sizes_ubo(),
      original_sizes_ubo.buffer());
}

void record_binary_op(
    api::Context* const context,
    const std::string& op_name,
    api::vTensor& v_in1,
    api::vTensor& v_in2,
    api::vTensor& v_dst) {
  std::string kernel_name = "binary_" + op_name + "_nobroadcast__test";
  add_dtype_suffix(kernel_name, v_dst);

  vkapi::PipelineBarrier pipeline_barrier{};
  vkapi::SpecVarList specialization_constants = {};
  context->submit_compute_job(
      VK_KERNEL_FROM_STR(kernel_name),
      pipeline_barrier,
      v_dst.logical_limits(),
      adaptive_work_group_size(v_dst.logical_limits()),
      specialization_constants,
      VK_NULL_HANDLE,
      0,
      v_dst.image(
          pipeline_barrier,
          vkapi::PipelineStage::COMPUTE,
          vkapi::MemoryAccessType::WRITE),
      v_in1.image(pipeline_barrier, vkapi::PipelineStage::COMPUTE),
      v_in2.image(pipeline_barrier, vkapi::PipelineStage::COMPUTE),
      v_dst.sizes_ubo());
}

void execute_and_check_add(
    api::vTensor& a,
    api::vTensor& b,
    api::vTensor& c,
    float a_val,
    float b_val) {
  // Fill input tensors
  fill_vtensor(a, a_val);
  fill_vtensor(b, b_val);

  // a + b = c
  record_binary_op(api::context(), "add", a, b, c);

  // Extract output tensor
  std::vector<float> data_out = extract_vtensor(c);

  // Check output
  for (size_t i = 0; i < data_out.size(); ++i) {
    CHECK_VALUE(data_out, i, (a_val + b_val));
  }
}

void record_index_fill_buffer(api::Context* context, api::vTensor& v_ten) {
  std::string kernel_name("idx_fill_buffer");
  switch (v_ten.dtype()) {
    case vkapi::kFloat:
      kernel_name += "_float";
      break;
    case vkapi::kHalf:
      kernel_name += "_half";
      break;
    case vkapi::kQInt8:
      kernel_name += "_int8";
      break;
    case vkapi::kQUInt8:
      kernel_name += "_uint8";
      break;
    default:
      throw std::runtime_error("Unsupported dtype");
      break;
  }

  api::ParamsBuffer params(api::context(), int32_t(v_ten.numel()));

  {
    vkapi::PipelineBarrier pipeline_barrier{};
    vkapi::SpecVarList specialization_constants = {};
    api::context()->submit_compute_job(
        VK_KERNEL_FROM_STR(kernel_name),
        pipeline_barrier,
        {uint32_t(v_ten.numel()), 1, 1},
        {64, 1, 1},
        specialization_constants,
        VK_NULL_HANDLE,
        0,
        v_ten.buffer(
            pipeline_barrier,
            vkapi::PipelineStage::COMPUTE,
            vkapi::MemoryAccessType::READ),
        params.buffer());
  }
}

void record_scalar_add_buffer(
    api::Context* context,
    api::vTensor& v_ten,
    float offset) {
  vkapi::PipelineBarrier pipeline_barrier{};
  vkapi::SpecVarList specialization_constants = {SV(offset)};
  std::string kernel = "scalar_add_buffer";
  add_dtype_suffix(kernel, v_ten);
  api::context()->submit_compute_job(
      VK_KERNEL_FROM_STR(kernel),
      pipeline_barrier,
      {uint32_t(v_ten.numel()), 1, 1},
      {64, 1, 1},
      specialization_constants,
      VK_NULL_HANDLE,
      0,
      v_ten.buffer(
          pipeline_barrier,
          vkapi::PipelineStage::COMPUTE,
          vkapi::MemoryAccessType::READ | vkapi::MemoryAccessType::WRITE),
      v_ten.numel_ubo());
}

void record_reference_matmul(
    api::Context* context,
    api::vTensor& out,
    api::vTensor& mat1,
    api::vTensor& mat2) {
  vkapi::PipelineBarrier pipeline_barrier{};
  api::context()->submit_compute_job(
      VK_KERNEL(reference_matmul),
      pipeline_barrier,
      {uint32_t(out.size(1)), uint32_t(out.size(0)), 1},
      {64, 1, 1},
      {},
      VK_NULL_HANDLE,
      0,
      out.buffer(
          pipeline_barrier,
          vkapi::PipelineStage::COMPUTE,
          vkapi::MemoryAccessType::WRITE),
      mat1.buffer(pipeline_barrier, vkapi::PipelineStage::COMPUTE),
      mat2.buffer(pipeline_barrier, vkapi::PipelineStage::COMPUTE),
      out.sizes_ubo(),
      out.strides_ubo(),
      mat1.sizes_ubo(),
      mat1.strides_ubo(),
      mat2.sizes_ubo(),
      mat2.strides_ubo());
}

void record_matmul_texture3d(
    api::Context* context,
    api::vTensor& out,
    api::vTensor& mat1,
    api::vTensor& mat2) {
  std::string kernel_name = "matmul_naive";
  kernel_name.reserve(kShaderNameReserve);
  add_storage_type_suffix(kernel_name, out.storage_type());
  add_dtype_suffix(kernel_name, out.dtype());

  utils::uvec3 global_wg_size = out.logical_limits();

  vkapi::PipelineBarrier pipeline_barrier{};
  api::context()->submit_compute_job(
      VK_KERNEL_FROM_STR(kernel_name),
      pipeline_barrier,
      global_wg_size,
      {8, 8, 1},
      {out.hashed_layout(), mat1.hashed_layout(), mat2.hashed_layout()},
      VK_NULL_HANDLE,
      0,
      out.image(
          pipeline_barrier,
          vkapi::PipelineStage::COMPUTE,
          vkapi::MemoryAccessType::WRITE),
      mat1.image(pipeline_barrier, vkapi::PipelineStage::COMPUTE),
      mat2.image(pipeline_barrier, vkapi::PipelineStage::COMPUTE),
      out.sizes_ubo(),
      out.logical_limits_ubo(),
      mat1.sizes_ubo(),
      mat2.sizes_ubo());
}

//
// Input & Output Utilities
//

#define FORALL_SUPPORTED_TYPES(_) \
  _(uint8_t, Byte)                \
  _(int8_t, Char)                 \
  _(int32_t, Int)                 \
  _(executorch::aten::Half, Half) \
  _(float, Float)                 \
  _(int8_t, QInt8)

void fill_vtensor(api::vTensor& vten, std::vector<float>& data) {
  api::StagingBuffer staging_buffer(api::context(), vten.dtype(), data.size());

#define CASE(ctype, name)                                     \
  case vkapi::ScalarType::name: {                             \
    std::vector<ctype> data_converted;                        \
    data_converted.resize(data.size());                       \
    for (int i = 0; i < data.size(); ++i) {                   \
      data_converted[i] = ctype(data[i]);                     \
    }                                                         \
    staging_buffer.copy_from(                                 \
        data_converted.data(), vten.staging_buffer_nbytes()); \
  } break;

  switch (vten.dtype()) {
    FORALL_SUPPORTED_TYPES(CASE)
    default:
      VK_THROW("Unsupported dtype");
  }

#undef CASE

  if (vten.storage_type() == utils::StorageType::BUFFER) {
    record_nchw_to_buffer_op(api::context(), staging_buffer.buffer(), vten);
  } else {
    record_nchw_to_image_op(api::context(), staging_buffer.buffer(), vten);
  }
}

void fill_vtensor(api::vTensor& vten, float val, bool iota) {
  std::vector<float> vten_data(vten.staging_buffer_numel());
  if (iota) {
    std::iota(vten_data.begin(), vten_data.end(), val);
  } else {
    std::fill(vten_data.begin(), vten_data.end(), val);
  }

  fill_vtensor(vten, vten_data);
}

std::vector<float> create_random_float_buffer(
    const size_t numel,
    const float min,
    const float max) {
  std::vector<float> data(numel);
  std::default_random_engine rng;
  std::uniform_real_distribution<float> dist(min, max);

  for (size_t i = 0; i < data.size(); ++i) {
    data[i] = dist(rng);
  }
  return data;
}

std::vector<uint8_t> create_random_uint8_buffer(
    const size_t numel,
    const uint8_t min,
    const uint8_t max) {
  std::vector<uint8_t> data(numel);
  std::default_random_engine rng;
  std::uniform_real_distribution<float> dist(min, max);

  for (size_t i = 0; i < data.size(); ++i) {
    data[i] = (uint8_t)dist(rng);
  }
  return data;
}

void fill_vtensor(
    ComputeGraph& graph,
    const IOValueRef idx,
    float val,
    bool iota) {
  vTensorPtr t = graph.get_tensor(idx.value);
  std::vector<float> data(t->numel());
  if (t->storage_type() != utils::kBuffer) {
    data.resize(t->staging_buffer_numel());
  }
  if (iota) {
    std::iota(data.begin(), data.end(), val);
  } else {
    std::fill(data.begin(), data.end(), val);
  }

  graph.copy_into_staging(idx.staging, data.data(), data.size());
}

void extract_vtensor(api::vTensor& vten, std::vector<float>& data) {
  api::StagingBuffer staging_buffer(
      api::context(), vten.dtype(), vten.staging_buffer_numel());

  if (vten.storage_type() == utils::StorageType::BUFFER) {
    record_buffer_to_nchw_op(api::context(), vten, staging_buffer.buffer());
  } else {
    record_image_to_nchw_op(api::context(), vten, staging_buffer.buffer());
  }

  vkapi::VulkanFence fence = api::context()->fences().get_fence();
  api::context()->submit_cmd_to_gpu(fence.get_submit_handle());
  fence.wait();

#define CASE(ctype, name)                                     \
  case vkapi::ScalarType::name: {                             \
    std::vector<ctype> data_converted(data.size());           \
    staging_buffer.copy_to(                                   \
        data_converted.data(), vten.staging_buffer_nbytes()); \
    for (int i = 0; i < data.size(); ++i) {                   \
      data[i] = float(data_converted[i]);                     \
    }                                                         \
  } break;

  switch (vten.dtype()) {
    FORALL_SUPPORTED_TYPES(CASE)
    default:
      VK_THROW("Unsupported dtype");
  }

#undef CASE
}

//
// Context Management
//

void submit_to_gpu() {
  vkapi::VulkanFence fence = api::context()->fences().get_fence();
  api::context()->submit_cmd_to_gpu(fence.get_submit_handle());
  fence.wait();
}

vkapi::Allocation allocate_memory_for(const api::vTensor& vten) {
  VmaAllocationCreateInfo alloc_create_info =
      api::context()->adapter_ptr()->vma().gpuonly_resource_create_info();
  return api::context()->adapter_ptr()->vma().create_allocation(
      vten.get_memory_requirements(), alloc_create_info);
}

VmaTotalStatistics get_vma_stats() {
  return api::context()->adapter_ptr()->vma().get_memory_statistics();
}

size_t get_vma_allocation_count() {
  return get_vma_stats().total.statistics.allocationCount;
}

//
// Graph Test Utilities
//

void execute_graph_and_check_output(
    ComputeGraph& graph,
    std::vector<float> input_vals,
    std::vector<float> expected_outputs) {
  assert(input_vals.size() == graph.inputs().size());
  assert(expected_outputs.size() == graph.outputs().size());

  for (size_t i = 0; i < graph.inputs().size(); ++i) {
    fill_vtensor(graph, graph.inputs().at(i), input_vals.at(i));
  }

  graph.execute();

  for (size_t i = 0; i < graph.outputs().size(); ++i) {
    IOValueRef out_ioval = graph.outputs().at(i);
    vTensorPtr t_out = graph.get_tensor(out_ioval.value);

    std::vector<float> output_data(t_out->staging_buffer_numel());
    graph.copy_from_staging(
        out_ioval.staging, output_data.data(), output_data.size());

    for (size_t j = 0; j < t_out->numel(); ++j) {
      CHECK_VALUE(output_data, j, expected_outputs.at(i));
    }
  }
}

bool check_close(float a, float b, float atol, float rtol) {
  float max = std::max(std::abs(a), std::abs(b));
  float diff = std::abs(a - b);
  return diff <= (atol + rtol * max);
}