#include "ANSLIB.h"

const int SEQUENCE_FRAMES = 16;
const int VALID_VIOLENCE_FRAMES = 8;
const float DEFAULT_FIGHT_DETECTION_THRESHOLD = 0.5f;

//#define VD_DEBUG
class CUSTOM_API ANSCustomViolence : public IANSCustomClass
{
private:
    using ANSLIBPtr = std::unique_ptr<ANSCENTER::ANSLIB, decltype(&ANSCENTER::ANSLIB::Destroy)>;

    std::recursive_mutex    _mutex;
    int                     engineType{ 0 };
    ANSLIBPtr _personDetector{ nullptr, &ANSCENTER::ANSLIB::Destroy };
    ANSLIBPtr _fightDetector{ nullptr, &ANSCENTER::ANSLIB::Destroy };
    ANSLIBPtr _movinetDetector{ nullptr, &ANSCENTER::ANSLIB::Destroy };

    // ----- Model config -----
    std::string _fightLabelMap;
    std::string _fightClassName;
    std::string _fightModelName;
    std::string _personModelName;
    std::string _personClassName;
    std::string _personLabelMap;

    int _fightModelType;
    int _fightDetectionType;
    int _personModelType;
    int _personDetectionType;
    int _movinetModelType;
    int _movinetDetectionType;

    // ----- Crowd tracking -----
    struct CrowdTrackInfo {
        cv::Rect bbox;
        int trackId;
        int frameIdx;
        std::vector<int> personIds;
    };
    std::vector<CrowdTrackInfo> _prevCrowdTracks;
    std::vector<CrowdTrackInfo> _crowdTracksHistory;
    int _nextCrowdTrackId = 1;
    std::map<int, cv::Rect> _crowdCropBoxes;
    int _crowdMaxAge = 3;

    // ----- Alert state -----
    int  _crowdFrameCount = 0;
    bool _crowdAlertActive = false;
    int  _noFightFrameCount = 0;
    static constexpr int ALERT_DECAY_FRAMES = 30;

    // ----- Per-camera frame counter -----
    std::map<std::string, int> _frameCounterPerCamera;

    // ----- MoviNet bbox stabilization per crowd track -----
    struct TrackCropState {
        cv::Rect lockedBbox;
        bool     isLocked = false;
        int      framesSinceLock = 0;
    };
    std::map<int, TrackCropState> _trackCropStates;

    // ----- Initialization flag -----
    bool _isInitialized{ false };

    // ----- Helpers -----
    static std::string MakeMovinetCameraId(const std::string& camera_id, int trackId) {
        return camera_id + "_crowd_" + std::to_string(trackId);
    }

    static float CalculateOverlap(const cv::Rect& r1, const cv::Rect& r2) {
        cv::Rect inter = r1 & r2;
        if (inter.empty()) return 0.0f;
        float interArea = static_cast<float>(inter.width * inter.height);
        float minArea = std::min(
            static_cast<float>(r1.width * r1.height),
            static_cast<float>(r2.width * r2.height));
        return (minArea > 0) ? (interArea / minArea) : 0.0f;
    }

    static cv::Rect ExpandBbox(const cv::Rect& bbox, float scale, const cv::Size& frameSize) {
        int cx = bbox.x + bbox.width / 2;
        int cy = bbox.y + bbox.height / 2;
        int nw = static_cast<int>(bbox.width * scale);
        int nh = static_cast<int>(bbox.height * scale);
        cv::Rect expanded(cx - nw / 2, cy - nh / 2, nw, nh);
        return expanded & cv::Rect(0, 0, frameSize.width, frameSize.height);
    }
public:
    bool Initialize(const std::string& modelDirectory, float detectionScoreThreshold,
        std::string& labelMap) override;
    bool OptimizeModel(bool fp16) override;
    bool ConfigureParameters(CustomParams& param) override;
    std::vector<CustomObject> RunInference(const cv::Mat& input) override;
    std::vector<CustomObject> RunInference(const cv::Mat& input,
        const std::string& camera_id) override;
    bool Destroy() override;
    ANSCustomViolence();
    ~ANSCustomViolence();
};
