#include <iostream>
#include <fstream>
#include <ai2d.h>
#include <opencv2/opencv.hpp>

// Test image path and size
std::string image_path="test.jpg";

void setup() {
  // put your setup code here, to run once:
  cv::Mat input_image=cv::imread(image_path);
  size_t img_w=input_image.cols;
  size_t img_h=input_image.rows;
  int pad_top=30;
  int pad_bottom=30;
  int pad_left=50;
  int pad_right=50;
  size_t out_w=img_w+pad_left+pad_right;
  size_t out_h=img_h+pad_top+pad_bottom;

  // Create input tensor from image data
  dims_t input_shape{1,img_h,img_w,3};
  runtime_tensor input_tensor = host_runtime_tensor::create(
      typecode_t::dt_uint8,
      input_shape,
      { (gsl::byte *)input_image.data, input_shape[0]*input_shape[1]*input_shape[2]*input_shape[3]},
      true,
      hrt::pool_shared
  ).expect("cannot create input tensor");
  hrt::sync(input_tensor, sync_op_t::sync_write_back, true).unwrap();

  /************* Padding Sample *************/
  AI2D ai2d_pad=AI2D();
  ai2d_pad.set_ai2d_dtype(ai2d_format::RGB_packed, ai2d_format::RGB_packed,
                          typecode_t::dt_uint8, typecode_t::dt_uint8);
  // Pad params: top=30, bottom=30, left=50, right=50, constant color=128,128,128
  std::vector<int> pad_param = {0, 0, 0, 0,pad_top,pad_bottom,pad_left,pad_right};
  ai2d_pad.set_pad(pad_param, ai2d_pad_mode::constant, {128,128,128});
  dims_t pad_input_shape{1,img_h,img_w,3};
  dims_t pad_output_shape{1,out_h,out_w,3};
  ai2d_pad.build(pad_input_shape,pad_output_shape);
  runtime_tensor ai2d_pad_tensor = host_runtime_tensor::create(typecode_t::dt_uint8, pad_output_shape, hrt::pool_shared).expect("cannot create input tensor");
  ai2d_pad.run(input_tensor,ai2d_pad_tensor);

  // Save pad result
  auto pad_input_buf = ai2d_pad_tensor.impl()->to_host().unwrap()->buffer().as_host().unwrap().map(map_access_::map_write).unwrap().buffer();
  char* pad_output_data=reinterpret_cast<char *>(pad_input_buf.data());
  cv::Mat pad_out_img(out_h,out_w, CV_8UC3, pad_output_data);
  cv::imwrite("pad_output.jpg", pad_out_img);
  /****************************************/
}

void loop() {
  // put your main code here, to run repeatedly:

}
