//   Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include <yaml-cpp/yaml.h>

#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

namespace PaddleX {

// Object for storing all preprocessed data
class ImageBlob {
 public:
  // Original image height and width
  std::vector<int> ori_im_size_ = std::vector<int>(2);
  // Newest image height and width after process
  std::vector<int> new_im_size_ = std::vector<int>(2);
  // Image height and width before resize
  std::vector<std::vector<int>> im_size_before_resize_;
  // Reshape order
  std::vector<std::string> reshape_order_;
  // Resize scale
  float scale = 1.0;
  // Buffer for image data after preprocessing
  std::vector<float> im_data_;

  void clear() {
    ori_im_size_.clear();
    new_im_size_.clear();
    im_size_before_resize_.clear();
    reshape_order_.clear();
    im_data_.clear();
  }
};

// Abstraction of preprocessing opration class
class Transform {
 public:
  virtual void Init(const YAML::Node& item) = 0;
  virtual bool Run(cv::Mat* im, ImageBlob* data) = 0;
  virtual void SetPaddingSize(int max_h, int max_w) {}
};

class Normalize : public Transform {
 public:
  virtual void Init(const YAML::Node& item) {
    mean_ = item["mean"].as<std::vector<float>>();
    std_ = item["std"].as<std::vector<float>>();
  }

  virtual bool Run(cv::Mat* im, ImageBlob* data);

 private:
  std::vector<float> mean_;
  std::vector<float> std_;
};

class ResizeByShort : public Transform {
 public:
  virtual void Init(const YAML::Node& item) {
    short_size_ = item["short_size"].as<int>();
    if (item["max_size"].IsDefined()) {
      max_size_ = item["max_size"].as<int>();
    } else {
      max_size_ = -1;
    }
  }
  virtual bool Run(cv::Mat* im, ImageBlob* data);

 private:
  float GenerateScale(const cv::Mat& im);
  int short_size_;
  int max_size_;
};

class ResizeByLong : public Transform {
 public:
  virtual void Init(const YAML::Node& item) {
    long_size_ = item["long_size"].as<int>();
  }
  virtual bool Run(cv::Mat* im, ImageBlob* data);

 private:
  int long_size_;
};

class Resize : public Transform {
 public:
  virtual void Init(const YAML::Node& item) {
    if (item["target_size"].IsScalar()) {
      height_ = item["target_size"].as<int>();
      width_ = item["target_size"].as<int>();
      interp_ = item["interp"].as<std::string>();
    } else if (item["target_size"].IsSequence()) {
      std::vector<int> target_size = item["target_size"].as<std::vector<int>>();
      width_ = target_size[0];
      height_ = target_size[1];
    }
    if (height_ <= 0 || width_ <= 0) {
      std::cerr << "[Resize] target_size should greater than 0" << std::endl;
      exit(-1);
    }
  }
  virtual bool Run(cv::Mat* im, ImageBlob* data);

 private:
  int height_;
  int width_;
  std::string interp_;
};

class CenterCrop : public Transform {
 public:
  virtual void Init(const YAML::Node& item) {
    if (item["crop_size"].IsScalar()) {
      height_ = item["crop_size"].as<int>();
      width_ = item["crop_size"].as<int>();
    } else if (item["crop_size"].IsSequence()) {
      std::vector<int> crop_size = item["crop_size"].as<std::vector<int>>();
      width_ = crop_size[0];
      height_ = crop_size[1];
    }
  }
  virtual bool Run(cv::Mat* im, ImageBlob* data);

 private:
  int height_;
  int width_;
};

class Padding : public Transform {
 public:
  virtual void Init(const YAML::Node& item) {
    if (item["coarsest_stride"].IsDefined()) {
      coarsest_stride_ = item["coarsest_stride"].as<int>();
      if (coarsest_stride_ < 1) {
        std::cerr << "[Padding] coarest_stride should greater than 0"
                  << std::endl;
        exit(-1);
      }
    }
    if (item["target_size"].IsDefined()) {
      if (item["target_size"].IsScalar()) {
        width_ = item["target_size"].as<int>();
        height_ = item["target_size"].as<int>();
      } else if (item["target_size"].IsSequence()) {
        width_ = item["target_size"].as<std::vector<int>>()[0];
        height_ = item["target_size"].as<std::vector<int>>()[1];
      }
    }
  }
  virtual bool Run(cv::Mat* im, ImageBlob* data);
  virtual void SetPaddingSize(int max_h, int max_w);
 private:
  int coarsest_stride_ = -1;
  int width_ = 0;
  int height_ = 0;
  int max_height_ = 0;
  int max_width_ = 0;
};

class Transforms {
 public:
  void Init(const YAML::Node& node, bool to_rgb = true);
  std::shared_ptr<Transform> CreateTransform(const std::string& name);
  bool Run(cv::Mat* im, ImageBlob* data);
  void SetPaddingSize(int max_h, int max_w);
 private:
  std::vector<std::shared_ptr<Transform>> transforms_;
  bool to_rgb_ = true;
  int max_h_ = 0;
  int max_w_ = 0;
};

}  // namespace PaddleX