//   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 <glog/logging.h>
#include <yaml-cpp/yaml.h>

#include <iostream>
#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 Detection {

// Object for storing all preprocessed data
    class ImageBlob {
    public:
        // image width and height
        std::vector<float> im_shape_;
        // Buffer for image data after preprocessing
        std::vector<float> im_data_;
        // in net data shape(after pad)
        std::vector<float> in_net_shape_;
        // Evaluation image width and height
        // std::vector<float>  eval_im_size_f_;
        // Scale factor for image size to origin image size
        std::vector<float> scale_factor_;
    };

// Abstraction of preprocessing opration class
    class PreprocessOp {
    public:
        virtual void Init(const YAML::Node &item) = 0;

        virtual void Run(cv::Mat *im, ImageBlob *data) = 0;
    };

    class InitInfo : public PreprocessOp {
    public:
        virtual void Init(const YAML::Node &item) {}

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

    class NormalizeImage : public PreprocessOp {
    public:
        virtual void Init(const YAML::Node &item) {
            mean_ = item["mean"].as < std::vector < float >> ();
            scale_ = item["std"].as < std::vector < float >> ();
            is_scale_ = item["is_scale"].as<bool>();
        }

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

    private:
        // CHW or HWC
        std::vector<float> mean_;
        std::vector<float> scale_;
        bool is_scale_;
    };

    class Permute : public PreprocessOp {
    public:
        virtual void Init(const YAML::Node &item) {}

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

    class Resize : public PreprocessOp {
    public:
        virtual void Init(const YAML::Node &item) {
            interp_ = item["interp"].as<int>();
            // max_size_ = item["target_size"].as<int>();
            keep_ratio_ = item["keep_ratio"].as<bool>();
            target_size_ = item["target_size"].as < std::vector < int >> ();
        }

        // Compute best resize scale for x-dimension, y-dimension
        std::pair<double, double> GenerateScale(const cv::Mat &im);

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

    private:
        int interp_ = 2;
        bool keep_ratio_;
        std::vector<int> target_size_;
        std::vector<int> in_net_shape_;
    };

// Models with FPN need input shape % stride == 0
    class PadStride : public PreprocessOp {
    public:
        virtual void Init(const YAML::Node &item) {
            stride_ = item["stride"].as<int>();
        }

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

    private:
        int stride_;
    };

    class Preprocessor {
    public:
        void Init(const YAML::Node &config_node) {
            // initialize image info at first
            ops_["InitInfo"] = std::make_shared<InitInfo>();
            for (int i = 0; i < config_node.size(); ++i) {
                if (config_node[i]["DetResize"].IsDefined()) {
                    ops_["Resize"] = std::make_shared<Resize>();
                    ops_["Resize"]->Init(config_node[i]["DetResize"]);
                }

                if (config_node[i]["DetNormalizeImage"].IsDefined()) {
                    ops_["NormalizeImage"] = std::make_shared<NormalizeImage>();
                    ops_["NormalizeImage"]->Init(config_node[i]["DetNormalizeImage"]);
                }

                if (config_node[i]["DetPermute"].IsDefined()) {
                    ops_["Permute"] = std::make_shared<Permute>();
                    ops_["Permute"]->Init(config_node[i]["DetPermute"]);
                }

                if (config_node[i]["DetPadStrid"].IsDefined()) {
                    ops_["PadStride"] = std::make_shared<PadStride>();
                    ops_["PadStride"]->Init(config_node[i]["DetPadStrid"]);
                }
            }
        }

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

    public:
        static const std::vector <std::string> RUN_ORDER;

    private:
        std::unordered_map <std::string, std::shared_ptr<PreprocessOp>> ops_;
    };

} // namespace Detection