#include "ANSCustomViolence.h"
#include "HDBSCAN.h"

ANSCustomViolence::ANSCustomViolence() {
}
bool ANSCustomViolence::Destroy() {
    std::lock_guard<std::recursive_mutex> lock(_mutex);
    try {
        _movinetDetector.reset();
        _personDetector.reset();
        _fightDetector.reset();
        _prevCrowdTracks.clear();
        _crowdTracksHistory.clear();
        _crowdCropBoxes.clear();
        _isInitialized = false;
        return true;
    }
    catch (...) {
        return false;
    }
}
ANSCustomViolence::~ANSCustomViolence() {
    Destroy();
}
bool ANSCustomViolence::Initialize(const std::string& modelDirectory, float detectionScoreThreshold, std::string& labelMap) {
    std::lock_guard<std::recursive_mutex> lock(_mutex);
    try {
        std::string _personLabelMap;
        std::string _fightLabelMap;
        std::string _movienetLabelMap;
        _modelDirectory = modelDirectory;
        _detectionScoreThreshold = detectionScoreThreshold;

        _isInitialized = false;
        _prevCrowdTracks.clear();
        _nextCrowdTrackId = 1;
        _crowdCropBoxes.clear();
        _crowdFrameCount = 0;
        _crowdAlertActive = false;

        // Create model instances using factory pattern (ABI-safe)
        _personDetector = ANSLIBPtr(ANSCENTER::ANSLIB::Create(), &ANSCENTER::ANSLIB::Destroy);
        _fightDetector = ANSLIBPtr(ANSCENTER::ANSLIB::Create(), &ANSCENTER::ANSLIB::Destroy);
        _movinetDetector = ANSLIBPtr(ANSCENTER::ANSLIB::Create(), &ANSCENTER::ANSLIB::Destroy);

        _personModelType = 4;       // TENSORRT
        _personDetectionType = 1;   // DETECTION
        _fightModelType = 4;        // TENSORRT
        _fightDetectionType = 1;    // DETECTION
        _movinetModelType = 27;     // MOVIENET
        _movinetDetectionType = 0;  // CLASSIFICATION

        // Check the hardware type
        engineType = _personDetector->GetEngineType();
        if (engineType == 1) {
            // NVIDIA GPU: Use TensorRT
            _personModelType = 4;       // TENSORRT
            _personDetectionType = 1;   // DETECTION
            _fightModelType = 4;        // TENSORRT
            _fightDetectionType = 1;    // DETECTION
            _movinetModelType = 27;     // MOVIENET
            _movinetDetectionType = 0;  // CLASSIFICATION
			std::cout << "NVIDIA GPU detected. Using TensorRT" << std::endl;
        }
        else {
            // CPU/Other: Use OpenVINO
            _personModelType = 5;       // OPENVINO
            _personDetectionType = 1;   // DETECTION
            _fightModelType = 5;        // OPENVINO
            _fightDetectionType = 1;    // DETECTION
            _movinetModelType = 27;     // MOVIENET
            _movinetDetectionType = 0;  // CLASSIFICATION
			std::cout << "CPU detected. Using OpenVINO" << std::endl;
        }

        if (this->_detectionScoreThreshold < 0.25f) this->_detectionScoreThreshold = 0.25f;

        // classId: 0=Normal, 1=Violence, 2=person, 3=crowd, 4=fighting, 5=normal_fight
        labelMap = "Normal,Violence,person,crowd,fighting,normal_fight";

#ifdef VD_DEBUG
        this->_loadEngineOnCreate = true;
#endif
        int loadEngineOnCreation = _loadEngineOnCreate ? 1 : 0;
        int autoEngineDetection = 1;
        std::string licenseKey = "";

        // Load person detector model
        float personScoreThreshold = 0.5f;
        float personConfThreshold = 0.5f;
        float personNMSThreshold = 0.5f;
        int personResult = _personDetector->LoadModelFromFolder(
            licenseKey.c_str(),
            "person", "person.names",
            personScoreThreshold, personConfThreshold, personNMSThreshold,
            autoEngineDetection,
            _personModelType, _personDetectionType,
            loadEngineOnCreation,
            modelDirectory.c_str(),
            _personLabelMap);
        if (personResult != 1) return false;

        // Load fight detector model
        float fightScoreThreshold = 0.25f;
        float fightConfThreshold = 0.5f;
        float fightNMSThreshold = 0.5f;
        int fightResult = _fightDetector->LoadModelFromFolder(
            licenseKey.c_str(),
            "fightOD", "fightOD.names",
            fightScoreThreshold, fightConfThreshold, fightNMSThreshold,
            autoEngineDetection,
            _fightModelType, _fightDetectionType,
            loadEngineOnCreation,
            modelDirectory.c_str(),
            _fightLabelMap);
        if (fightResult != 1) return false;

        // Load MoviNet temporal model
        float movinetScoreThreshold = this->_detectionScoreThreshold;
        float movinetConfThreshold = 0.5f;
        float movinetNMSThreshold = 0.5f;
        int movinetResult = _movinetDetector->LoadModelFromFolder(
            licenseKey.c_str(),
            "MoviNet151", "MoviNet151.names",
            movinetScoreThreshold, movinetConfThreshold, movinetNMSThreshold,
            autoEngineDetection,
            _movinetModelType, _movinetDetectionType,
            loadEngineOnCreation,
            modelDirectory.c_str(),
            _movienetLabelMap);
        if (movinetResult != 1) return false;

        _isInitialized = true;
        return true;
    }
    catch (const std::exception& ex) {
        return false;
    }
}
bool ANSCustomViolence::OptimizeModel(bool fp16)
{
    if (!_personDetector || !_fightDetector) return false;
    int personDetectionResult = _personDetector->Optimize(fp16);
	int fightDetectionResult = _fightDetector->Optimize(fp16);
	if ( (personDetectionResult != 1) || 
        (fightDetectionResult != 1)) return false;
    else return true;
}
bool ANSCustomViolence::ConfigureParameters(CustomParams& param) {
    return true;
}
std::vector<CustomObject> ANSCustomViolence::RunInference(const cv::Mat& input)
{
    return RunInference(input, "CustomCam");
}

std::vector<CustomObject> ANSCustomViolence::RunInference(const cv::Mat& input, const std::string& camera_id)
{
    std::lock_guard<std::recursive_mutex> lock(_mutex);
    std::vector<CustomObject> results;

    try {
        if (!_isInitialized || !_personDetector || !_fightDetector) {
            return results;
        }

        // Per-camera frame counter (fixes global counter bug for multi-camera)
        int& frameCounter = _frameCounterPerCamera[camera_id];
        frameCounter++;
        const int currentFrameIdx = frameCounter;

        // ----- Constants -----
        const float personConfThr = static_cast<float>(_detectionScoreThreshold);
        const float clusterEps = 600.0f;
        const int   clusterMinPts = 2;
        const float fightConfThr = DEFAULT_FIGHT_DETECTION_THRESHOLD;
        const float bboxExpandScale = 1.5f;
        const float bboxDriftThresh = 0.5f;
        const int   bboxDriftCheckAt = 4;

        bool currentFrameHasCrowd = false;

        // ==============================================================
        // Helper: Run fight detector on a bbox region
        // ==============================================================
        auto runFightDetector = [&](const cv::Rect& safeBbox) -> CustomObject {
            cv::Mat cropped = input(safeBbox).clone();
            std::vector<ANSCENTER::Object> fightResults;
            _fightDetector->RunInference(cropped, camera_id.c_str(), fightResults);

            bool hasFight = false;
            float fightConf = 0.0f;
            for (const auto& obj : fightResults) {
                float conf = static_cast<float>(obj.confidence);
                if (obj.classId == 1 && conf >= fightConfThr) {
                    hasFight = true;
                    fightConf = conf;
                    break;
                }
                else if (obj.classId == 0 && conf > fightConf) {
                    fightConf = conf;
                }
            }

            CustomObject fightObj;
            fightObj.classId = hasFight ? 4 : 5;
            fightObj.className = hasFight ? "fighting" : "normal_fight";
            fightObj.confidence = fightConf;
            fightObj.box = safeBbox;
            fightObj.cameraId = camera_id;
            return fightObj;
            };

        // ==============================================================
        // Helper: Feed MoviNet with bbox-stabilized crop
        // ==============================================================
        auto runMovinet = [&](const cv::Rect& currentBbox, int trackId) {
            if (!_movinetDetector) return;

            auto& cropState = _trackCropStates[trackId];

            // Lock bbox at the start of each 16-frame window
            if (!cropState.isLocked) {
                cropState.lockedBbox = currentBbox;
                cropState.isLocked = true;
                cropState.framesSinceLock = 0;
            }

            cropState.framesSinceLock++;

            // Check for significant crowd drift — re-lock if crowd moved too far
            if (cropState.framesSinceLock > bboxDriftCheckAt) {
                float overlap = CalculateOverlap(cropState.lockedBbox, currentBbox);
                if (overlap < bboxDriftThresh) {
                    cropState.lockedBbox = currentBbox;
                    cropState.framesSinceLock = 1;
                }
            }

            // Crop with the locked bbox
            cv::Rect lockedRoi = cropState.lockedBbox
                & cv::Rect(0, 0, input.cols, input.rows);
            if (lockedRoi.width <= 0 || lockedRoi.height <= 0) return;

            cv::Mat movinetCrop = input(lockedRoi).clone();
            std::string movinetCamId = MakeMovinetCameraId(camera_id, trackId);

            std::vector<ANSCENTER::Object> movinetResults;
            _movinetDetector->RunInference(movinetCrop, movinetCamId.c_str(), movinetResults);

            for (const auto& mr : movinetResults) {
                CustomObject movinetObj;
                movinetObj.classId = (mr.classId == 1) ? 1 : 0;
                movinetObj.className = mr.className;
                movinetObj.confidence = mr.confidence;
                movinetObj.box = lockedRoi;
                movinetObj.cameraId = camera_id;
                results.push_back(movinetObj);

                // Inference completed — unlock for next window
                cropState.isLocked = false;
                cropState.framesSinceLock = 0;
            }
            };

        // ==============================================================
        // Helper: Combined fight + movinet for a region
        // ==============================================================
        auto processCrowdRegion = [&](const cv::Rect& expandedBbox, int trackId) {
            cv::Rect safeBbox = expandedBbox & cv::Rect(0, 0, input.cols, input.rows);
            if (safeBbox.width <= 0 || safeBbox.height <= 0) return;

            // Fight detection (always uses live bbox)
            CustomObject fightObj = runFightDetector(safeBbox);
            results.push_back(fightObj);

            // MoviNet (uses locked bbox for spatial consistency)
            runMovinet(safeBbox, trackId);
            };

        // ==============================================================
        // BLOCK 1: Person detection
        // ==============================================================
        std::vector<ANSCENTER::Object> personResults;
        _personDetector->RunInference(input, camera_id.c_str(), personResults);

        std::vector<cv::Rect> boxes;
        std::vector<int> personIds;

        if (!personResults.empty()) {
            boxes.reserve(personResults.size());
            personIds.reserve(personResults.size());

            for (const auto& p : personResults) {
                float conf = static_cast<float>(p.confidence);
                if (conf < personConfThr) continue;

                boxes.push_back(p.box);
                personIds.push_back(p.trackId);

                CustomObject personObj;
                personObj.classId = 2;
                personObj.className = "person";
                personObj.confidence = conf;
                personObj.box = p.box;
                personObj.cameraId = camera_id;
                personObj.trackId = p.trackId;
                results.push_back(personObj);
            }
        }

        // ==============================================================
        // BLOCK 2: Clustering + Crowd tracking
        // ==============================================================
        std::vector<cv::Rect> crowdBoxes;
        std::vector<int> groupTrackIds;

        if (!boxes.empty()) {
            auto clusters = HDBSCAN::cluster(boxes, clusterEps, clusterMinPts);
            std::vector<std::vector<int>> crowdPersonIds;

            for (const auto& c : clusters) {
                if (static_cast<int>(c.size()) < clusterMinPts) continue;

                cv::Rect clusterBox;
                bool first = true;
                std::vector<int> ids;

                for (int idx : c) {
                    if (idx < 0 || idx >= static_cast<int>(boxes.size())) continue;
                    if (first) { clusterBox = boxes[idx]; first = false; }
                    else { clusterBox |= boxes[idx]; }
                    ids.push_back(personIds[idx]);
                }

                if (!first) {
                    crowdBoxes.push_back(clusterBox);
                    crowdPersonIds.push_back(std::move(ids));
                }
            }

            // --- Assign crowd track IDs by person-ID overlap ---
            if (!crowdBoxes.empty()) {
                currentFrameHasCrowd = true;
                _crowdFrameCount++;
                if (_crowdFrameCount >= 2) _crowdAlertActive = true;
                _noFightFrameCount = 0;

                groupTrackIds.reserve(crowdBoxes.size());

                for (size_t gIdx = 0; gIdx < crowdBoxes.size(); ++gIdx) {
                    std::set<int> pidSet(crowdPersonIds[gIdx].begin(),
                        crowdPersonIds[gIdx].end());
                    int assignedId = -1;
                    int bestOverlap = 0;

                    for (const auto& prev : _prevCrowdTracks) {
                        int overlap = 0;
                        for (int pid : pidSet) {
                            if (std::find(prev.personIds.begin(),
                                prev.personIds.end(), pid)
                                != prev.personIds.end())
                                ++overlap;
                        }
                        if (overlap > bestOverlap && overlap >= 1) {
                            bestOverlap = overlap;
                            assignedId = prev.trackId;
                        }
                    }

                    if (assignedId == -1) assignedId = _nextCrowdTrackId++;
                    groupTrackIds.push_back(assignedId);
                }

                // Update tracking state
                _prevCrowdTracks.clear();
                for (size_t gIdx = 0; gIdx < crowdBoxes.size(); ++gIdx) {
                    CrowdTrackInfo info;
                    info.bbox = crowdBoxes[gIdx];
                    info.trackId = groupTrackIds[gIdx];
                    info.frameIdx = currentFrameIdx;
                    info.personIds = crowdPersonIds[gIdx];
                    _prevCrowdTracks.push_back(info);
                    _crowdTracksHistory.push_back(info);
                }

                // Prune old history
                const int cutoff = currentFrameIdx - _crowdMaxAge;
                _crowdTracksHistory.erase(
                    std::remove_if(_crowdTracksHistory.begin(),
                        _crowdTracksHistory.end(),
                        [cutoff](const CrowdTrackInfo& c) {
                            return c.frameIdx < cutoff;
                        }),
                    _crowdTracksHistory.end());

                // Output crowd objects
                for (size_t gIdx = 0; gIdx < crowdBoxes.size(); ++gIdx) {
                    CustomObject crowdObj;
                    crowdObj.classId = 3;
                    crowdObj.className = "crowd";
                    crowdObj.confidence = 1.0f;
                    crowdObj.box = crowdBoxes[gIdx];
                    crowdObj.cameraId = camera_id;
                    crowdObj.trackId = groupTrackIds[gIdx];
                    results.push_back(crowdObj);
                }
            }
        }

        // ==============================================================
        // BLOCK 3: Fight detection + MoviNet for each active crowd
        // ==============================================================
        for (size_t gIdx = 0; gIdx < crowdBoxes.size(); ++gIdx) {
            cv::Rect expanded = ExpandBbox(crowdBoxes[gIdx], bboxExpandScale, input.size());
            int trackId = groupTrackIds[gIdx];
            _crowdCropBoxes[trackId] = expanded;

            processCrowdRegion(expanded, trackId);
        }

        // ==============================================================
        // BLOCK 4: Alert mode — no crowd, monitor last known region
        // ==============================================================
        if (_crowdAlertActive && !currentFrameHasCrowd && !_crowdCropBoxes.empty()) {
            int nearestTrackId = -1;
            if (!_prevCrowdTracks.empty()) {
                nearestTrackId = _prevCrowdTracks[0].trackId;
            }
            else {
                for (const auto& [tid, _] : _crowdCropBoxes) {
                    if (tid > nearestTrackId) nearestTrackId = tid;
                }
            }

            if (nearestTrackId != -1 && _crowdCropBoxes.count(nearestTrackId)) {
                processCrowdRegion(_crowdCropBoxes[nearestTrackId], nearestTrackId);
            }

            // Decay alert mode
            _noFightFrameCount++;
            if (_noFightFrameCount >= ALERT_DECAY_FRAMES) {
                _crowdAlertActive = false;
                _crowdCropBoxes.clear();
                _noFightFrameCount = 0;
            }
        }
        else if (!currentFrameHasCrowd) {
            _crowdFrameCount = 0;
        }

        // ==============================================================
        // BLOCK 5: Clean up local stale tracking state
        // ==============================================================
        if (currentFrameIdx % 50 == 0) {
            std::set<int> activeIds;
            for (const auto& ct : _prevCrowdTracks) activeIds.insert(ct.trackId);
            for (const auto& [tid, _] : _crowdCropBoxes) activeIds.insert(tid);

            for (auto it = _trackCropStates.begin(); it != _trackCropStates.end(); ) {
                if (activeIds.find(it->first) == activeIds.end()) {
                    it = _trackCropStates.erase(it);
                }
                else {
                    ++it;
                }
            }
            for (auto it = _crowdCropBoxes.begin(); it != _crowdCropBoxes.end(); ) {
                if (activeIds.find(it->first) == activeIds.end()) {
                    it = _crowdCropBoxes.erase(it);
                }
                else {
                    ++it;
                }
            }
        }
        return results;
    }
    catch (const std::exception& e) {
        std::cerr << "RunInference error: " << e.what() << std::endl;
        return results;
    }
}
