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

void setup() {
  // put your setup code here, to run once:
  /************* Shift (Bit Shift) Sample *************/
  // Create a synthetic 16-bit image with constant value
  size_t img_w=320;
  size_t img_h=320;
  cv::Mat shift_img(img_h, img_w, CV_16UC3, cv::Scalar(120, 120, 120));
  cv::imwrite("shift_ori.jpg", shift_img);

  AI2D ai2d_shift=AI2D();
  ai2d_shift.set_ai2d_dtype(ai2d_format::RAW16, ai2d_format::RAW16,
                            typecode_t::dt_uint16, typecode_t::dt_uint16);
  // Shift right by 1 bit
  ai2d_shift.set_shift(1);
  dims_t shift_input_shape{1,img_h, img_w,3};
  dims_t shift_output_shape{1,img_h, img_w,3};
  ai2d_shift.build(shift_input_shape,shift_output_shape);

  runtime_tensor shift_input_tensor = host_runtime_tensor::create(
      typecode_t::dt_uint16,
      shift_input_shape,
      { (gsl::byte *)shift_img.data, 320*320*3*sizeof(uint16_t)},
      true,
      hrt::pool_shared
  ).expect("cannot create input tensor");
  hrt::sync(shift_input_tensor, sync_op_t::sync_write_back, true).unwrap();
  
  runtime_tensor ai2d_shift_tensor = host_runtime_tensor::create(typecode_t::dt_uint16, shift_output_shape, hrt::pool_shared).expect("cannot create input tensor");
  ai2d_shift.run(shift_input_tensor,ai2d_shift_tensor);

  // Save shift result
  auto shift_input_buf = ai2d_shift_tensor.impl()->to_host().unwrap()->buffer().as_host().unwrap().map(map_access_::map_write).unwrap().buffer();
  char* shift_output_data=reinterpret_cast<char *>(shift_input_buf.data());
  cv::Mat shift_out_img(img_h, img_w, CV_16UC3, shift_output_data);
  cv::imwrite("shift_output.jpg", shift_out_img);
  /*******************************************/
}

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

}
