#include "FaceNet.h"
namespace ANSCENTER {
    std::string ANSFaceNet::GetOpenVINODevice() {
        ov::Core core;
        std::vector<std::string> available_devices = core.get_available_devices();
        // Prioritize devices: NPU > GPU > CPU
        std::vector<std::string> priority_devices = {"GPU", "CPU" };
        for (const auto& device : priority_devices) {
            if (std::find(available_devices.begin(), available_devices.end(), device) != available_devices.end()) {
                return device; // Return the first available device based on priority
            }
        }

        // Default to CPU if no prioritized devices are found
        return "CPU";
    }

    bool ANSFaceNet::Initialize(std::string licenseKey,ModelConfig modelConfig,
        const std::string& modelZipFilePath,const std::string& modelZipPassword,std::string& labelMap) 
    {
        bool result = ANSFRBase::Initialize(licenseKey, modelConfig, modelZipFilePath, modelZipPassword, labelMap);
        if (!result) return false;
        try {
            _modelConfig = modelConfig;
            _modelConfig.modelType = ModelType::FACERECOGNIZE;
            _modelConfig.detectionType = DetectionType::FACERECOGNIZER;

            // We need to get the modelfolder from here
            std::string faceidModel = CreateFilePath(_modelFolder, "ansfacenet.xml");
            if (std::filesystem::exists(faceidModel)) {
                _modelFilePath = faceidModel;
                this->_logger.LogDebug("ANSFaceNet::Initialize.  Loading ansfacenet weight", _modelFilePath, __FILE__, __LINE__);
            }
            else {   
                this->_logger.LogError("ANSFaceNet::Initialize.  Model ansfacenet.xml file is not exist", _modelFilePath, __FILE__, __LINE__);
                return false;
            }
            m_knownPersonThresh = _modelConfig.unknownPersonThreshold;
            if (m_knownPersonThresh == 0)m_knownPersonThresh = 0.35;

            std::string deviceName = GetOpenVINODevice();
            ov::Core core;
            CnnConfig reid_config(_modelFilePath, "Face Re-Identification");
            reid_config.m_deviceName = deviceName;
            //core.set_property(deviceName, ov::hint::performance_mode(ov::hint::PerformanceMode::THROUGHPUT));
            reid_config.m_max_batch_size = 1;
            reid_config.m_core = core;
            this->faceRecognizer = std::make_unique<VectorCNN>(reid_config);
            if (faceRecognizer == nullptr) {
                this->_logger.LogFatal("ANSFaceNet::Initialize", "Failed to initialize face recognizer model", __FILE__, __LINE__);
                return false;
            }
            Init();
            _isInitialized = true;  
            return true;
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSFaceNet::Initialize", e.what(), __FILE__, __LINE__);
            return false;
        }
    }

    bool ANSFaceNet::LoadModel(const std::string& modelZipFilePath, const std::string& modelZipPassword) {
        try {
            // We need to get the _modelFolder
            bool result = ANSFRBase::LoadModel(modelZipFilePath, modelZipPassword);
            if (!result) return false;
           // We need to get the modelfolder from here
            std::string faceidModel = CreateFilePath(_modelFolder, "ansfacenet.xml");
            if (std::filesystem::exists(faceidModel)) {
                _modelFilePath = faceidModel;
                this->_logger.LogDebug("ANSFaceNet::LoadModel.  Loading ansfacenet weight", _modelFilePath, __FILE__, __LINE__);
            }
            else {
                this->_logger.LogError("ANSFaceNet::Initialize.  Model ansfacenet.xml file is not exist", _modelFilePath, __FILE__, __LINE__);
                return false;
            }
            return true;
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ArcFace50::LoadModel", e.what(), __FILE__, __LINE__);
            return false;
        }
    }
    bool ANSFaceNet::OptimizeModel(bool fp16, std::string& optimizedModelFolder) {
        if (!FileExist(_modelFilePath)) {
            optimizedModelFolder = "";
            return false;
        }
        return true;
    }
    std::vector<float> ANSFaceNet::Feature(const cv::Mat& image, const ANSCENTER::Object& bBox) {
        std::lock_guard<std::recursive_mutex> lock(_mutex);
        std::vector<float> embeddingResult;
        try {
            if (image.empty()) return embeddingResult;
            if ((image.cols < 10) || (image.rows < 10)) return embeddingResult;
            std::vector<cv::Mat> embeddings;
            std::vector<cv::Mat> face_rois;
            face_rois.clear();
            embeddings.clear();
            face_rois.push_back(bBox.mask);
            faceRecognizer->Compute(face_rois, &embeddings);
            embeddingResult.assign(embeddings[0].begin<float>(), embeddings[0].end<float>());
            return embeddingResult;
        }

        catch (std::exception& e) {
            std::vector<float> embeddingResult;
            embeddingResult.clear();
            this->_logger.LogFatal("ANSFaceNet::Feature", e.what(), __FILE__, __LINE__);
            return embeddingResult;
        }
    }
    std::vector<FaceResultObject> ANSFaceNet::Match(
        const cv::Mat& input,
        const std::vector<ANSCENTER::Object>& bBox,
        const std::map<std::string, std::string>& userDict) {

        std::vector<FaceResultObject> resultObjects;

        // Early validation before locking
        if (input.empty()) {
            return resultObjects;
        }

        if (input.cols < 10 || input.rows < 10) {
            return resultObjects;
        }

        std::lock_guard<std::recursive_mutex> lock(_mutex);

        if (!_isInitialized) {
            _logger.LogError("ANSFaceNet::Match", "Model is not initialized", __FILE__, __LINE__);
            return resultObjects;
        }

        try {
            // Convert grayscale to 3-channel BGR if needed
            cv::Mat processedImage;
            if (input.channels() == 1) {
                cv::cvtColor(input, processedImage, cv::COLOR_GRAY2BGR);
            }
            else {
                processedImage = input;
            }

            // Get embeddings
            std::vector<std::vector<float>> detectedEmbeddings = Forward(processedImage, bBox);

            // Search for matches
            std::vector<std::string> names;
            std::vector<float> sims;
            std::tie(names, sims) = SearchForFaces(detectedEmbeddings);

            // Check if we got results
            if (names.empty()) {
                _logger.LogError("ANSFaceNet::Match", "No face is match", __FILE__, __LINE__);
                return resultObjects;
            }

            // Pre-reserve result space
            const size_t resultCount = std::min(names.size(), bBox.size());
            resultObjects.reserve(resultCount);

            // Build result objects
            for (size_t i = 0; i < resultCount; ++i) {
                FaceResultObject resultObject;

                // Determine if face is known or unknown
                const bool isUnknown = (sims[i] > m_knownPersonThresh);

                if (isUnknown) {
                    resultObject.isUnknown = true;
                    resultObject.userId = "0";
                    resultObject.userName = "Unknown";
                    resultObject.confidence = 1.0f; // 100% unknown
                }
                else {
                    resultObject.isUnknown = false;
                    resultObject.userId = names[i];

                    // Safe map lookup
                    auto it = userDict.find(names[i]);
                    resultObject.userName = (it != userDict.end()) ? it->second : names[i];

                    resultObject.confidence = 1.0f - sims[i];
                }

                resultObject.similarity = sims[i];

                // Copy bounding box (no clamping - keeping original logic)
                resultObject.box.x = bBox[i].box.x;
                resultObject.box.y = bBox[i].box.y;
                resultObject.box.width = bBox[i].box.width;
                resultObject.box.height = bBox[i].box.height;

                // Copy additional data
                resultObject.mask = bBox[i].mask;
                resultObject.cameraId = bBox[i].cameraId;
                resultObject.trackId = bBox[i].trackId;
                resultObject.polygon = bBox[i].polygon;
                resultObject.kps = bBox[i].kps;
                resultObject.extraInformation = bBox[i].extraInfo;

                resultObjects.push_back(std::move(resultObject));
            }

            return resultObjects;
        }
        catch (const std::exception& e) {
            _logger.LogFatal("ANSFaceNet::Match", e.what(), __FILE__, __LINE__);
            return resultObjects;
        }
    }
    cv::Mat ANSFaceNet::GetCropFace(const cv::Mat& input, const ANSCENTER::Object& bBox) {
        try {
            cv::Mat processedImage;
            if (input.channels() == 1) {
                cv::cvtColor(input, processedImage, cv::COLOR_GRAY2BGR);
            }
            else {
                processedImage = input;
            }
            std::vector<ANSCENTER::Object> outputBbox;
            outputBbox.push_back(bBox);
            std::vector<struct CroppedFace> crFaces;
            crFaces.clear();
            ANSFRHelper::GetCroppedFaces(processedImage, outputBbox, FACE_WIDTH, FACE_HEIGHT, crFaces);
            return crFaces[0].faceMat;
        }
        catch (std::exception& e) {
            cv::Mat result;
            this->_logger.LogFatal("ANSFaceNet::GetCropFace", e.what(), __FILE__, __LINE__);
            return result;
        }
    }
    bool ANSFaceNet::LoadEngine(const std::string xmlModelPath, bool engineOptimisation) {
        try {
            if (!FileExist(xmlModelPath)) {
                this->_logger.LogError("ANSFaceNet::LoadEngine", "Cannot find the raw XML model file.", __FILE__, __LINE__);
                return false;
            }        
            return true;
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSFaceNet::LoadEngine", e.what(), __FILE__, __LINE__);
            return false;
        }
    }
    // Private methods, can be replacable
    void ANSFaceNet::AddEmbedding(const std::string& className, float embedding[]) {
        //std::lock_guard<std::recursive_mutex> lock(_mutex);
        try {
            // Check if faiss_index is initialized
            if (!faiss_index) {
                this->_logger.LogFatal("ANSFaceNet::AddEmbedding", "Search_index is not initialized.", __FILE__, __LINE__);
                return;
            }
            // Convert embedding array to vector and check size
            std::vector<float> vec(embedding, embedding + FACE_EMBEDDING_SIZE);
            if (vec.size() != FACE_EMBEDDING_SIZE) {
                this->_logger.LogFatal("ANSFaceNet::AddEmbedding", "Embedding size does not match expected output dimension of 512.", __FILE__, __LINE__);
                return;
            }
            // Add class name and get index
            classNames.push_back(className);

            // Add embedding to faiss_index
            faiss_index->add(1, vec.data());
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSFaceNet::AddEmbedding", e.what(), __FILE__, __LINE__);
        }

    }
    void ANSFaceNet::AddEmbedding(const std::string& className, const std::vector<float>& embedding) {
        //std::lock_guard<std::recursive_mutex> lock(_mutex);
        try {
            // Ensure faiss_index is initialized
            if (!faiss_index) {
                this->_logger.LogFatal("ANSFaceNet::AddEmbedding", "Search_index is not initialized.", __FILE__, __LINE__);
                return;
            }
            // Check embedding size
            if (embedding.size() != FACE_EMBEDDING_SIZE) {
                this->_logger.LogFatal("ANSFaceNet::AddEmbedding", "Embedding size does not match expected output dimension of 512.", __FILE__, __LINE__);
                return;
            }
            // Add class name
            classNames.push_back(className);
            // Add embedding to faiss_index
            faiss_index->add(1, embedding.data());
        }
        catch (const std::exception& e) {
            this->_logger.LogFatal("ANSFaceNet::AddEmbedding", e.what(), __FILE__, __LINE__);
        }
    }
    std::vector<std::vector<float>> ANSFaceNet::Forward(const cv::Mat& input, std::vector<ANSCENTER::Object> outputBbox) {
        //std::lock_guard<std::recursive_mutex> lock(_mutex);
        std::vector<std::vector<float>> detectedEmbeddings;
        try {
            std::vector<cv::Mat> embeddings;
            std::vector<cv::Mat> face_rois;
            detectedEmbeddings.clear();
            if (outputBbox.size() > 0) {
                for (int i = 0; i < outputBbox.size(); i++) {
                    face_rois.clear();
                    embeddings.clear();
                    std::vector<float> embeddingRs;
                    face_rois.push_back(outputBbox[i].mask);
                    faceRecognizer->Compute(face_rois, &embeddings); 
                    embeddingRs.assign(embeddings[0].begin<float>(), embeddings[0].end<float>());
                    detectedEmbeddings.push_back(embeddingRs);
                }
            }
            return detectedEmbeddings;
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSFaceNet::Forward", e.what(), __FILE__, __LINE__);
            return detectedEmbeddings;
        }
    }
    std::tuple<std::vector<std::string>, std::vector<float>> ANSFaceNet::SearchForFaces(const std::vector<std::vector<float>>& detectedEmbeddings)
    {
        //std::lock_guard<std::recursive_mutex> lock(_mutex);
        std::vector<std::string> detectedUsers;
        std::vector<float> simValues;
        try {
            // Check if there are class names available
            if (!classNames.empty()) {
                const int k = 3;  // Number of nearest neighbors to retrieve
                // Process each detected embedding
                for (const auto& embedding : detectedEmbeddings) {
                    // Prepare vectors to hold search results
                    std::vector<faiss::idx_t> indices(k);
                    std::vector<float> distances(k);

                    // Perform the search for the k nearest neighbors
                    faiss_index->search(1, embedding.data(), k, distances.data(), indices.data());

                    // Find the index with the maximum distance
                    auto min_it = std::min_element(distances.begin(), distances.end());
                    int best_index = std::distance(distances.begin(), min_it);

                    // Map the index to the corresponding class name and calculate similarity
                    std::vector<float> matchEmbedding(faiss_index->d);
                    faiss_index->reconstruct(indices[best_index], matchEmbedding.data());
                    float similarity = 1 - CosineSimilarity(embedding, matchEmbedding, false);
                    detectedUsers.push_back(classNames.at(indices[best_index]));
                    simValues.push_back(std::abs(similarity));

                }
            }
            else {
                // If no class names are available, mark all users as "unknown"
                detectedUsers.assign(detectedEmbeddings.size(), "0");
                simValues.assign(detectedEmbeddings.size(), 1.0f);
            }

            return std::make_tuple(detectedUsers, simValues);

        }
        catch (const std::exception& e) {
            // Log the error and return default values for the failed search
            detectedUsers.assign(1, "0");
            simValues.assign(1, 1.0f);
            this->_logger.LogFatal("ANSFaceNet::SearchForFaces", e.what(), __FILE__, __LINE__);
            return std::make_tuple(detectedUsers, simValues);
        }
    }
    void ANSFaceNet::Init() {
        try {
            classNames.clear();
            if (faiss_index) {
                faiss_index->reset();
            }
            else {
                faiss_index = std::make_unique<faiss::IndexFlatL2>(FACE_EMBEDDING_SIZE);
            }
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSFaceNet::Init", e.what(), __FILE__, __LINE__);
        }
    }
    ANSFaceNet::~ANSFaceNet() {
        try {
			Destroy();
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSFaceNet::~ANSFaceNet()", e.what(), __FILE__, __LINE__);
        }
    }
    bool ANSFaceNet::Destroy() {
        try {   
            // Clear all resources safely
            classNames.clear();

            // Reset and release faiss_index
            if (faiss_index) {
                faiss_index->reset();
                faiss_index.reset(); // Release shared pointer ownership
            }
			faceRecognizer.reset();
            return true;
        }
        catch (std::exception& e) {
            this->_logger.LogFatal("ANSFaceNet::Destroy", e.what(), __FILE__, __LINE__);
			return false;
        }
    }
}