//
// Created by DefTruth on 2021/9/20.
//

#include "rvm.h"
#include "lite/ort/core/ort_utils.h"
#include "lite/utils.h"

using ortcv::RobustVideoMatting;

RobustVideoMatting::RobustVideoMatting(const std::string &_onnx_path, unsigned int _num_threads) :
    log_id(_onnx_path.data()), num_threads(_num_threads)
{
#ifdef LITE_WIN32
  std::wstring _w_onnx_path(lite::utils::to_wstring(_onnx_path));
  onnx_path = _w_onnx_path.data();
#else
  onnx_path = _onnx_path.data();
#endif
  ort_env = Ort::Env(ORT_LOGGING_LEVEL_ERROR, log_id);
  // 0. session options
  Ort::SessionOptions session_options;
  session_options.SetIntraOpNumThreads(num_threads);
  session_options.SetGraphOptimizationLevel(
      GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
  session_options.SetLogSeverityLevel(4);
  // 1. session
  // GPU Compatibility.
#ifdef USE_CUDA
  OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0); // C API stable.
#endif
  ort_session = new Ort::Session(ort_env, onnx_path, session_options);
#if LITEORT_DEBUG
  std::cout << "Load " << onnx_path << " done!" << std::endl;
#endif
}

RobustVideoMatting::~RobustVideoMatting()
{
  if (ort_session)
    delete ort_session;
  ort_session = nullptr;
}

int64_t RobustVideoMatting::value_size_of(const std::vector<int64_t> &dims)
{
  if (dims.empty()) return 0;
  int64_t value_size = 1;
  for (const auto &size: dims) value_size *= size;
  return value_size;
}

std::vector<Ort::Value> RobustVideoMatting::transform(const cv::Mat &mat)
{
  cv::Mat src = mat.clone();
  const unsigned int img_height = mat.rows;
  const unsigned int img_width = mat.cols;
  std::vector<int64_t> &src_dims = dynamic_input_node_dims.at(0); // (1,3,h,w)
  // update src height and width
  src_dims.at(2) = img_height;
  src_dims.at(3) = img_width;
  // assume that rxi's dims and value_handler was updated by last step in a while loop.
  std::vector<int64_t> &r1i_dims = dynamic_input_node_dims.at(1); // (1,?,?h,?w)
  std::vector<int64_t> &r2i_dims = dynamic_input_node_dims.at(2); // (1,?,?h,?w)
  std::vector<int64_t> &r3i_dims = dynamic_input_node_dims.at(3); // (1,?,?h,?w)
  std::vector<int64_t> &r4i_dims = dynamic_input_node_dims.at(4); // (1,?,?h,?w)
  std::vector<int64_t> &dsr_dims = dynamic_input_node_dims.at(5); // (1)
  int64_t src_value_size = this->value_size_of(src_dims); // (1*3*h*w)
  int64_t r1i_value_size = this->value_size_of(r1i_dims); // (1*?*?h*?w)
  int64_t r2i_value_size = this->value_size_of(r2i_dims); // (1*?*?h*?w)
  int64_t r3i_value_size = this->value_size_of(r3i_dims); // (1*?*?h*?w)
  int64_t r4i_value_size = this->value_size_of(r4i_dims); // (1*?*?h*?w)
  int64_t dsr_value_size = this->value_size_of(dsr_dims); // 1
  dynamic_src_value_handler.resize(src_value_size);

  // normalize & RGB
  cv::cvtColor(src, src, cv::COLOR_BGR2RGB); // (h,w,3)
  src.convertTo(src, CV_32FC3, 1.0f / 255.0f, 0.f); // 0.~1.

  // convert to tensor.
  std::vector<Ort::Value> input_tensors;
  input_tensors.emplace_back(ortcv::utils::transform::create_tensor(
      src, src_dims, memory_info_handler, dynamic_src_value_handler,
      ortcv::utils::transform::CHW
  ));
  input_tensors.emplace_back(Ort::Value::CreateTensor<float>(
      memory_info_handler, dynamic_r1i_value_handler.data(),
      r1i_value_size, r1i_dims.data(), r1i_dims.size()
  ));
  input_tensors.emplace_back(Ort::Value::CreateTensor<float>(
      memory_info_handler, dynamic_r2i_value_handler.data(),
      r2i_value_size, r2i_dims.data(), r2i_dims.size()
  ));
  input_tensors.emplace_back(Ort::Value::CreateTensor<float>(
      memory_info_handler, dynamic_r3i_value_handler.data(),
      r3i_value_size, r3i_dims.data(), r3i_dims.size()
  ));
  input_tensors.emplace_back(Ort::Value::CreateTensor<float>(
      memory_info_handler, dynamic_r4i_value_handler.data(),
      r4i_value_size, r4i_dims.data(), r4i_dims.size()
  ));
  input_tensors.emplace_back(Ort::Value::CreateTensor<float>(
      memory_info_handler, dynamic_dsr_value_handler.data(),
      dsr_value_size, dsr_dims.data(), dsr_dims.size()
  ));

  return input_tensors;
}

void RobustVideoMatting::detect(const cv::Mat &mat, types::MattingContent &content,
                                float downsample_ratio, bool video_mode)
{
  if (mat.empty()) return;
  // 0. set dsr at runtime.
  dynamic_dsr_value_handler.at(0) = downsample_ratio;

  // 1. make input tensors, src, rxi, dsr
  std::vector<Ort::Value> input_tensors = this->transform(mat);
  // 2. inference, fgr, pha, rxo.
  auto output_tensors = ort_session->Run(
      Ort::RunOptions{nullptr}, input_node_names.data(),
      input_tensors.data(), num_inputs, output_node_names.data(),
      num_outputs
  );
  // 3. generate matting
  this->generate_matting(output_tensors, content);
  // 4. update context (needed for video detection.)
  if (video_mode)
  {
    context_is_update = false; // init state.
    this->update_context(output_tensors);
  }

}


void RobustVideoMatting::detect_video(const std::string &video_path,
                                      const std::string &output_path,
                                      std::vector<types::MattingContent> &contents,
                                      bool save_contents, float downsample_ratio,
                                      unsigned int writer_fps)
{
  // 0. init video capture
  cv::VideoCapture video_capture(video_path);
  const unsigned int width = video_capture.get(cv::CAP_PROP_FRAME_WIDTH);
  const unsigned int height = video_capture.get(cv::CAP_PROP_FRAME_HEIGHT);
  const unsigned int frame_count = video_capture.get(cv::CAP_PROP_FRAME_COUNT);
  if (!video_capture.isOpened())
  {
    std::cout << "Can not open video: " << video_path << "\n";
    return;
  }
  // 1. init video writer
  cv::VideoWriter video_writer(output_path, cv::VideoWriter::fourcc('m', 'p', '4', 'v'),
                               writer_fps, cv::Size(width, height));
  if (!video_writer.isOpened())
  {
    std::cout << "Can not open writer: " << output_path << "\n";
    return;
  }

  // 2. matting loop
  cv::Mat mat;
  unsigned int i = 0;
  while (video_capture.read(mat))
  {
    i += 1;
    types::MattingContent content;
    this->detect(mat, content, downsample_ratio, true); // video_mode true
    // 3. save contents and writing out.
    if (content.flag)
    {
      if (save_contents) contents.push_back(content);
      if (!content.merge_mat.empty()) video_writer.write(content.merge_mat);
    }
    // 4. check context states.
    if (!context_is_update) break;
#ifdef LITEORT_DEBUG
    std::cout << i << "/" << frame_count << " done!" << "\n";
#endif
  }

  // 5. release
  video_capture.release();
  video_writer.release();

}

void RobustVideoMatting::generate_matting(std::vector<Ort::Value> &output_tensors,
                                          types::MattingContent &content)
{
  Ort::Value &fgr = output_tensors.at(0); // fgr (1,3,h,w) 0.~1.
  Ort::Value &pha = output_tensors.at(1); // pha (1,1,h,w) 0.~1.
  auto fgr_dims = fgr.GetTypeInfo().GetTensorTypeAndShapeInfo().GetShape();
  auto pha_dims = pha.GetTypeInfo().GetTensorTypeAndShapeInfo().GetShape();
  const unsigned int height = fgr_dims.at(2); // output height
  const unsigned int width = fgr_dims.at(3); // output width
  const unsigned int channel_step = height * width;
  // fast assign & channel transpose(CHW->HWC).
  float *fgr_ptr = fgr.GetTensorMutableData<float>();
  float *pha_ptr = pha.GetTensorMutableData<float>();
  cv::Mat rmat(height, width, CV_32FC1, fgr_ptr);
  cv::Mat gmat(height, width, CV_32FC1, fgr_ptr + channel_step);
  cv::Mat bmat(height, width, CV_32FC1, fgr_ptr + 2 * channel_step);
  cv::Mat pmat(height, width, CV_32FC1, pha_ptr);
  rmat *= 255.;
  bmat *= 255.;
  gmat *= 255.;
  cv::Mat rest = 1. - pmat;
  cv::Mat mbmat = bmat.mul(pmat) + rest * 153.;
  cv::Mat mgmat = gmat.mul(pmat) + rest * 255.;
  cv::Mat mrmat = rmat.mul(pmat) + rest * 120.;
  std::vector<cv::Mat> fgr_channel_mats, merge_channel_mats;
  fgr_channel_mats.push_back(bmat);
  fgr_channel_mats.push_back(gmat);
  fgr_channel_mats.push_back(rmat);
  merge_channel_mats.push_back(mbmat);
  merge_channel_mats.push_back(mgmat);
  merge_channel_mats.push_back(mrmat);

  content.pha_mat = pmat;
  cv::merge(fgr_channel_mats, content.fgr_mat);
  cv::merge(merge_channel_mats, content.merge_mat);
  content.fgr_mat.convertTo(content.fgr_mat, CV_8UC3);
  content.merge_mat.convertTo(content.merge_mat, CV_8UC3);

  content.flag = true;
}


void RobustVideoMatting::update_context(std::vector<Ort::Value> &output_tensors)
{
  // 0. update context for video matting.
  Ort::Value &r1o = output_tensors.at(2); // fgr (1,?,?h,?w)
  Ort::Value &r2o = output_tensors.at(3); // pha (1,?,?h,?w)
  Ort::Value &r3o = output_tensors.at(4); // pha (1,?,?h,?w)
  Ort::Value &r4o = output_tensors.at(5); // pha (1,?,?h,?w)
  auto r1o_dims = r1o.GetTypeInfo().GetTensorTypeAndShapeInfo().GetShape();
  auto r2o_dims = r2o.GetTypeInfo().GetTensorTypeAndShapeInfo().GetShape();
  auto r3o_dims = r3o.GetTypeInfo().GetTensorTypeAndShapeInfo().GetShape();
  auto r4o_dims = r4o.GetTypeInfo().GetTensorTypeAndShapeInfo().GetShape();
  // 1. update rxi's shape according to last rxo
  dynamic_input_node_dims.at(1) = r1o_dims;
  dynamic_input_node_dims.at(2) = r2o_dims;
  dynamic_input_node_dims.at(3) = r3o_dims;
  dynamic_input_node_dims.at(4) = r4o_dims;
  // 2. update rxi's value according to last rxo
  int64_t new_r1i_value_size = this->value_size_of(r1o_dims); // (1*?*?h*?w)
  int64_t new_r2i_value_size = this->value_size_of(r2o_dims); // (1*?*?h*?w)
  int64_t new_r3i_value_size = this->value_size_of(r3o_dims); // (1*?*?h*?w)
  int64_t new_r4i_value_size = this->value_size_of(r4o_dims); // (1*?*?h*?w)
  dynamic_r1i_value_handler.resize(new_r1i_value_size);
  dynamic_r2i_value_handler.resize(new_r2i_value_size);
  dynamic_r3i_value_handler.resize(new_r3i_value_size);
  dynamic_r4i_value_handler.resize(new_r4i_value_size);
  float *new_r1i_value_ptr = r1o.GetTensorMutableData<float>();
  float *new_r2i_value_ptr = r2o.GetTensorMutableData<float>();
  float *new_r3i_value_ptr = r3o.GetTensorMutableData<float>();
  float *new_r4i_value_ptr = r4o.GetTensorMutableData<float>();
  std::memcpy(dynamic_r1i_value_handler.data(), new_r1i_value_ptr, new_r1i_value_size * sizeof(float));
  std::memcpy(dynamic_r2i_value_handler.data(), new_r2i_value_ptr, new_r2i_value_size * sizeof(float));
  std::memcpy(dynamic_r3i_value_handler.data(), new_r3i_value_ptr, new_r3i_value_size * sizeof(float));
  std::memcpy(dynamic_r4i_value_handler.data(), new_r4i_value_ptr, new_r4i_value_size * sizeof(float));
  context_is_update = true;
}