#include "RetinaFaceDetector.h"
#include "Utility.h"
namespace ANSCENTER {
    bool ANSRETINAFD::Initialize(std::string licenseKey, ModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, std::string& labelMap) {
        bool result = ANSFDBase::Initialize(licenseKey, modelConfig, modelZipFilePath, modelZipPassword, labelMap);
        // We do not need to check for the license 
        _licenseValid = true;
        if (!_licenseValid) return false;
        try {
            labelMap = "Face";
            _modelConfig = modelConfig;
            _modelConfig.modelType = ModelType::FACEDETECT;
            _modelConfig.detectionType = DetectionType::FACEDETECTOR;

            // We need to get the modelfolder from here
            std::string onnxfile = CreateFilePath(_modelFolder, "retinaface.onnx");
            if (std::filesystem::exists(onnxfile)) {
                _modelFilePath = onnxfile;
                this->_logger.LogDebug("ANSRETINAFD::Initialize.  Loading retina weight", _modelFilePath, __FILE__, __LINE__);
            }
            else {
                this->_logger.LogError("ANSRETINAFD::Initialize.  Model retinaface.onnx file is not exist", _modelFilePath, __FILE__, __LINE__);
                return false;
            }
            std::string params_file;
            std::string TrtCacheFile = CreateFilePath(_modelFolder, "retinaface.cache");;
            auto option = fastdeploy::RuntimeOption();

            // Check if the system has NVIDIA GPU
            EngineType engineType = ANSLicenseHelper::CheckHardwareInformation();
            if (engineType == EngineType::NVIDIA_GPU) // NVIDIA CUDA
            {
                option.UseGpu();
                option.UseTrtBackend();
                option.trt_option.SetShape("images", { 1, 3, 640, 640 });
                option.trt_option.enable_fp16 = false;
                option.trt_option.serialize_file = TrtCacheFile;
            }
            else {
                // Otherwise it use CPU
                option.UseCpu();
                option.UseOrtBackend();
            }
            auto format = fastdeploy::ModelFormat::ONNX;
            this->_faceDectector = std::make_unique<fastdeploy::vision::facedet::RetinaFace>(onnxfile, params_file, option, format);
            if (!_faceDectector->Initialized()) {
                this->_logger.LogFatal("ANSRETINAFD::Initialize", "Failed to initialize face detector model", __FILE__, __LINE__);
                return false;
            }
            _isInitialized = true;
            return true;
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSRETINAFD::Initialize", e.what(), __FILE__, __LINE__);
            return false;
        }
    }

    bool ANSRETINAFD::LoadModel(const std::string& modelZipFilePath, const std::string& modelZipPassword) {
        try {
            // We need to get the _modelFolder
            bool result = ANSFDBase::LoadModel(modelZipFilePath, modelZipPassword);
            if (!result) return false;

            // We need to get the modelfolder from here
            std::string onnxfile = CreateFilePath(_modelFolder, "retinaface.onnx");
            if (std::filesystem::exists(onnxfile)) {
                _modelFilePath = onnxfile;
                this->_logger.LogDebug("ANSRETINAFD::LoadModel.  Loading retinaface weight", _modelFilePath, __FILE__, __LINE__);
            }
            else {
                this->_logger.LogError("ANSRETINAFD::LoadModel.  Model retinaface.onnx file is not exist", _modelFilePath, __FILE__, __LINE__);
                return false;
            }
            return true;
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSRETINAFD::LoadModel", e.what(), __FILE__, __LINE__);
            return false;
        }

    }
    bool ANSRETINAFD::OptimizeModel(bool fp16, std::string& optimizedModelFolder){
        if (!FileExist(_modelFilePath)) {
            optimizedModelFolder = "";
            return false;
        }
        optimizedModelFolder = GetParentFolder(_modelFilePath);
        std::string params_file;
        std::string TrtCacheFile = CreateFilePath(_modelFolder, "retinaface.cache");;
        auto option = fastdeploy::RuntimeOption();
        EngineType engineType = ANSLicenseHelper::CheckHardwareInformation();
        if (engineType == EngineType::NVIDIA_GPU) // NVIDIA CUDA
        {
            option.UseGpu();
            option.UseTrtBackend();
            option.trt_option.SetShape("images", { 1, 3, 640, 640 });
            option.trt_option.enable_fp16 = fp16;
            option.trt_option.serialize_file = TrtCacheFile;
        }
        else {
            option.UseCpu();
            option.UseOrtBackend();
        }
        auto format = fastdeploy::ModelFormat::ONNX;
        auto model = fastdeploy::vision::facedet::RetinaFace(_modelFilePath, params_file, option, format);
        if (!model.Initialized()) {
            this->_logger.LogError("ANSRETINAFD::OptimizeModel.", "Failed to initialize", __FILE__, __LINE__);
            return false;
        }
        return true;
    }
    std::vector<Object> ANSRETINAFD::RunInference(const cv::Mat& input) {
        std::vector<Object> output;
        output.clear();
        if (!_licenseValid) {
    	if (_modelLoading.load()) return {};
            this->_logger.LogError("ANSRETINAFD::RunInference", "Invalid license", __FILE__, __LINE__);
            return output;
        }
        if (!_isInitialized) {
			this->_logger.LogError("ANSRETINAFD::RunInference", "Model is not initialized", __FILE__, __LINE__);
			return output;
		}
        try {
            bool croppedFace = false; // Check if the image is cropped face image
            cv::Mat im = input.clone();
            // We know that the image sizes <=300 px, it is likely that image is cropped for face only
            if ((input.size[0] <= 300) || (input.size[1] <= 300)) croppedFace = true;
            if (croppedFace) cv::copyMakeBorder(input, im, 200, 200, 200, 200, cv::BORDER_REPLICATE);

            fastdeploy::vision::FaceDetectionResult res;
            if (!_faceDectector->Predict(&im, &res)) {
                this->_logger.LogFatal("ANSRETINAFD::RunInference", "Failed to predict.", __FILE__, __LINE__);
                return output;
            }
            if (res.boxes.size() > 0) {
                // Peform face alignment
                std::vector<cv::Mat> detectedFaces =
                    fastdeploy::vision::utils::AlignFaceWithFivePoints(im, res);
                if (res.boxes.size() == detectedFaces.size()) {
                    for (int i = 0; i < res.boxes.size(); i++)
                    {
                        Object result;
                        float confidence = res.scores[i];
                        if (confidence >= _modelConfig.detectionScoreThreshold) {
                            int x_min = res.boxes[i][0];
                            int y_min = res.boxes[i][1];
                            int x_max = res.boxes[i][2];
                            int y_max = res.boxes[i][3];
                            result.classId = 0;
                            result.className = "Face";
                            result.confidence = confidence;
                            result.box.x = x_min;
                            result.box.y = y_min;
                            if (croppedFace) {
                                if (x_min <= 200) x_min = 200;
                                if (y_min <= 200) y_min = 200;
                                result.box.x = x_min - 200;
                                result.box.y = y_min - 200;
                            }
                            result.box.width = x_max - x_min;
                            result.box.height = y_max - y_min;
                            result.mask = detectedFaces.at(i).clone();
							result.cameraId = "RETINAFACECAM";
                            output.push_back(result);
                        }
                    }
                }
                //detectedFaces.clear(); 
            }
            im.release();
            res.Clear();
            res.Free();
            EnqueueDetection(output, "RETINAFACECAM");
            return output;
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSRETINAFD::RunInference", e.what(), __FILE__, __LINE__);
            return output;
        }
    }

    std::vector<Object> ANSRETINAFD::RunInference(const cv::Mat& input, const std::string& camera_id) {
        std::vector<Object> output;
        output.clear();
        if (!_licenseValid) {
    	if (_modelLoading.load()) return {};
            this->_logger.LogError("ANSRETINAFD::RunInference", "Invalid license", __FILE__, __LINE__);
            return output;
        }
        if (!_isInitialized) {
            this->_logger.LogError("ANSRETINAFD::RunInference", "Model is not initialized", __FILE__, __LINE__);
            return output;
        }
        try {
            bool croppedFace = false; // Check if the image is cropped face image
            cv::Mat im = input.clone();
            // We know that the image sizes <=300 px, it is likely that image is cropped for face only
            if ((input.size[0] <= 300) || (input.size[1] <= 300)) croppedFace = true;
            if (croppedFace) cv::copyMakeBorder(input, im, 200, 200, 200, 200, cv::BORDER_REPLICATE);

            fastdeploy::vision::FaceDetectionResult res;
            if (!_faceDectector->Predict(&im, &res)) {
                this->_logger.LogFatal("ANSRETINAFD::RunInference", "Failed to predict.", __FILE__, __LINE__);
                return output;
            }
            if (res.boxes.size() > 0) {
                // Peform face alignment
                std::vector<cv::Mat> detectedFaces =
                    fastdeploy::vision::utils::AlignFaceWithFivePoints(im, res);
                if (res.boxes.size() == detectedFaces.size()) {
                    for (int i = 0; i < res.boxes.size(); i++)
                    {
                        Object result;
                        float confidence = res.scores[i];
                        if (confidence >= _modelConfig.detectionScoreThreshold) {
                            int x_min = res.boxes[i][0];
                            int y_min = res.boxes[i][1];
                            int x_max = res.boxes[i][2];
                            int y_max = res.boxes[i][3];
                            result.classId = 0;
                            result.className = "Face";
                            result.confidence = confidence;
                            result.box.x = x_min;
                            result.box.y = y_min;
                            if (croppedFace) {
                                if (x_min <= 200) x_min = 200;
                                if (y_min <= 200) y_min = 200;
                                result.box.x = x_min - 200;
                                result.box.y = y_min - 200;
                            }
                            result.box.width = x_max - x_min;
                            result.box.height = y_max - y_min;
                            result.mask = detectedFaces.at(i).clone();
							result.cameraId = camera_id;
                            output.push_back(result);
                        }
                    }
                }
                //detectedFaces.clear(); 
            }
            im.release();
            res.Clear();
            res.Free();
			EnqueueDetection(output,camera_id);
            return output;
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSRETINAFD::RunInference", e.what(), __FILE__, __LINE__);
            return output;
        }
    }


    ANSRETINAFD::~ANSRETINAFD() {
        try {
    //        if (_faceDectector) {
    //            delete _faceDectector;
				//_faceDectector = nullptr;
    //        }
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSRETINAFD::Destroy", e.what(), __FILE__, __LINE__);
        }
    }

    bool ANSRETINAFD::Destroy() {
        try {
            //if (_faceDectector) {
            //    delete _faceDectector;
            //    _faceDectector = nullptr;
            //}
            return true;
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSRETINAFD::Destroy", e.what(), __FILE__, __LINE__);
            return false;
        }
    }
}