#ifndef ANSPPTracker_H
#define ANSPPTracker_H
#define ANSPPTracker_API __declspec(dllexport)
#pragma once

#include <stddef.h>
#include <functional>
#include <string>
#include <vector>
#include <opencv2/core.hpp>
#include <openvino/openvino.hpp>
#include <unordered_map>
#include <deque>
#include <iostream>

#include <cstddef>
#include <map>
#include <memory>
#include <set>
#include <tuple>
#include <unordered_map>
#include <utility>
#include <vector>
#include <opencv2/imgproc.hpp>
#define PT_CHECK(cond) CV_Assert(cond);

#define PT_CHECK_BINARY(actual, expected, op) \
    CV_Assert(actual op expected);

#define PT_CHECK_EQ(actual, expected) PT_CHECK_BINARY(actual, expected, ==)
#define PT_CHECK_NE(actual, expected) PT_CHECK_BINARY(actual, expected, !=)
#define PT_CHECK_LT(actual, expected) PT_CHECK_BINARY(actual, expected, <)
#define PT_CHECK_GT(actual, expected) PT_CHECK_BINARY(actual, expected, >)
#define PT_CHECK_LE(actual, expected) PT_CHECK_BINARY(actual, expected, <=)
#define PT_CHECK_GE(actual, expected) PT_CHECK_BINARY(actual, expected, >=)
template <typename T>
std::ostream& operator<<(std::ostream& os, const std::deque<T>& v) {
    os << "[\n";
    if (!v.empty()) {
        auto itr = v.begin();
        os << *itr;
        for (++itr; itr != v.end(); ++itr)
            os << ",\n" << *itr;
    }
    os << "\n]";
    return os;
}
template <typename T>
std::ostream& operator<<(std::ostream& os, const std::vector<T>& v) {
    os << "[\n";
    if (!v.empty()) {
        auto itr = v.begin();
        os << *itr;
        for (++itr; itr != v.end(); ++itr)
            os << ",\n" << *itr;
    }
    os << "\n]";
    return os;
}

namespace ANSPPTrack {

    struct ppTrackedObject {
        cv::Rect rect;  ///< Detected object ROI (zero area if N/A).
        double confidence;  ///< Detection confidence level (-1 if N/A).
        int64_t frame_idx;  ///< Frame index where object was detected (-1 if N/A).
        int object_id;  ///< Unique object identifier (-1 if N/A).
        uint64_t timestamp;  ///< Timestamp in milliseconds.
        ppTrackedObject() : confidence(-1), frame_idx(-1), object_id(-1), timestamp(0) {}
        ppTrackedObject(const cv::Rect& rect, float confidence, int64_t frame_idx, int object_id)
            : rect(rect),
            confidence(confidence),
            frame_idx(frame_idx),
            object_id(object_id),
            timestamp(0) {}
    };
    struct ppDetectionLogEntry;
    struct ppTrackerParams;
    struct ppTrack;
    struct ppModelConfigTracker;
    class ppIImageDescriptor;
    class ppResizedImageDescriptor;
    class ppBaseModel;
    class ppVectorCNN;
    class ppDescriptor;
    class ppIDescriptorDistance;
    class ppCosDistance;
    class ppMatchTemplateDistance;
    class PedestrianTracker;

    using ppTrackedObjects = std::deque<ppTrackedObject>;
    using ppObjectTracks = std::unordered_map<int, std::deque<ppTrackedObject>>;
    using ppDetectionLog = std::vector<ppDetectionLogEntry>;

    //void SaveDetectionLogToTrajFile(const std::string& path, const DetectionLog& log);
    void PrintDetectionLog(const ppDetectionLog& log);
    void DrawPolyline(const std::vector<cv::Point>& polyline, const cv::Scalar& color, cv::Mat* image, int lwd = 5);

    struct ppDetectionLogEntry {
        std::deque<ppTrackedObject> objects;  ///< Detected objects.
        int64_t frame_idx;  ///< Processed frame index (-1 if N/A).
        double time_ms;  ///< Frame processing time in ms (-1 if N/A).
        ppDetectionLogEntry() : frame_idx(-1), time_ms(-1) {}
        ppDetectionLogEntry(const ppDetectionLogEntry& other)
            : objects(other.objects),
            frame_idx(other.frame_idx),
            time_ms(other.time_ms) {}
        ppDetectionLogEntry(ppDetectionLogEntry&& other)
            : objects(std::move(other.objects)),
            frame_idx(other.frame_idx),
            time_ms(other.time_ms) {}
        ppDetectionLogEntry& operator=(const ppDetectionLogEntry& other) = default;
        ppDetectionLogEntry& operator=(ppDetectionLogEntry&& other) {
            if (this != &other) {
                objects = std::move(other.objects);
                frame_idx = other.frame_idx;
                time_ms = other.time_ms;
            }
            return *this;
        }
    };
    struct ppTrackerParams {
        size_t min_track_duration;  ///< Min track duration in milliseconds.
        size_t forget_delay;  ///< Forget about track if the last bounding box in
        float aff_thr_fast;  ///< Affinity threshold which is used to determine if
        float aff_thr_strong;  ///< Affinity threshold which is used to determine if
        float shape_affinity_w;  ///< Shape affinity weight.

        float motion_affinity_w;  ///< Motion affinity weight.
        float time_affinity_w;  ///< Time affinity weight.
        float min_det_conf;  ///< Min confidence of detection.

        cv::Vec2f bbox_aspect_ratios_range;  ///< Bounding box aspect ratios range.
        cv::Vec2f bbox_heights_range;  ///< Bounding box heights range.

        int predict;  ///< How many frames are used to predict bounding box in case
        float strong_affinity_thr;  ///< If 'fast' confidence is greater than this
        float reid_thr;  ///< Affinity threshold for re-identification.
        bool drop_forgotten_tracks;  ///< Drop forgotten tracks. If it's enabled it
        int max_num_objects_in_track;  ///< The number of objects in track is
        ppTrackerParams();
    };
    struct ppTrack {
        ppTrack(const std::deque<ppTrackedObject>& objs,
            const cv::Mat& last_image,
            const cv::Mat& descriptor_fast,
            const cv::Mat& descriptor_strong);
        bool empty() const {
            return objects.empty();
        }

        size_t size() const {
            return objects.size();
        }

        const ppTrackedObject& operator[](size_t i) const {
            return objects[i];
        }

        ppTrackedObject& operator[](size_t i) {
            return objects[i];
        }


        const ppTrackedObject& back() const {
            PT_CHECK(!empty());
            return objects.back();
        }

        ppTrackedObject& back() {
            PT_CHECK(!empty());
            return objects.back();
        }

        std::deque<ppTrackedObject> objects;  ///< Detected objects;
        ppTrackedObject first_object;  ///< First object in track.
        cv::Rect predicted_rect;  ///< Rectangle that represents predicted position
        cv::Mat last_image;  ///< Image of last detected object in track.
        cv::Mat descriptor_fast;  ///< Fast descriptor.
        cv::Mat descriptor_strong;  ///< Strong descriptor (reid embedding).
        size_t lost;  ///< How many frames ago track has been lost.

        size_t length;  ///< Length of a track including number of objects that were
    };

    struct ppModelConfigTracker {
        explicit ppModelConfigTracker(const std::string& path_to_model) : path_to_model(path_to_model) {}
        std::string path_to_model;
        int max_batch_size{ 1 };
    };
    
    class ppIImageDescriptor {
    public:
        virtual cv::Size size() const = 0;
        virtual void Compute(const cv::Mat& mat, cv::Mat* descr) = 0;
        virtual void Compute(const std::vector<cv::Mat>& mats, std::vector<cv::Mat>* descrs) = 0;
        virtual ~ppIImageDescriptor() {}
    };
    class ppResizedImageDescriptor : public ppIImageDescriptor {
    public:
        explicit ppResizedImageDescriptor(const cv::Size& descr_size, const cv::InterpolationFlags interpolation)
            : descr_size_(descr_size),
            interpolation_(interpolation) {
            PT_CHECK_GT(descr_size.width, 0);
            PT_CHECK_GT(descr_size.height, 0);
        }
        cv::Size size() const override {
            return descr_size_;
        }
        void Compute(const cv::Mat& mat, cv::Mat* descr) override {
            PT_CHECK(descr != nullptr);
            PT_CHECK(!mat.empty());
            cv::resize(mat, *descr, descr_size_, 0, 0, interpolation_);
        }
        void Compute(const std::vector<cv::Mat>& mats, std::vector<cv::Mat>* descrs) override {
            PT_CHECK(descrs != nullptr);
            descrs->resize(mats.size());
            for (size_t i = 0; i < mats.size(); i++) {
                Compute(mats[i], &(descrs[i]));
            }
        }

    private:
        cv::Size descr_size_;
        cv::InterpolationFlags interpolation_;
    };
    class ppBaseModel {
    public:
        using Config = ppModelConfigTracker;
        ppBaseModel(const Config& config,
            const ov::Core& core,
            const std::string& deviceName);
        virtual ~ppBaseModel() {}
        void Load();
        const std::string modelType = "Person Re-Identification";

    protected:
        void InferBatch(const std::vector<cv::Mat>& frames,
            const std::function<void(const ov::Tensor&, size_t)>& results_fetcher) const;

        Config config;
        ov::Core core;
        std::string device_name;
        ov::Layout input_layout;
        ov::CompiledModel compiled_model;
        mutable ov::InferRequest infer_request;
        mutable ov::Tensor input_tensor;
        ov::Shape input_shape;
        ov::Tensor output_tensor;
        ov::Shape output_shape;
    };
    class ppVectorCNN : public ppBaseModel {
    public:
        ppVectorCNN(const ppModelConfigTracker& config, const ov::Core& core, const std::string& deviceName);
        void Compute(const cv::Mat& image, cv::Mat* vector) const;
        void Compute(const std::vector<cv::Mat>& images, std::vector<cv::Mat>* vectors) const;

        int size() const {
            return result_size;
        }
    private:
        int result_size;  // Length of result
    };
    class ppDescriptor : public ppIImageDescriptor {
    private:
        ppVectorCNN handler;
    public:
        ppDescriptor(const ppModelConfigTracker& config,
            const ov::Core& core,
            const std::string& deviceName)
            : handler(config, core, deviceName) {}
        cv::Size size() const override {
            return cv::Size(1, handler.size());
        }
        void Compute(const cv::Mat& mat, cv::Mat* descr) override {
            handler.Compute(mat, descr);
        }
        void Compute(const std::vector<cv::Mat>& mats, std::vector<cv::Mat>* descrs) override {
            handler.Compute(mats, descrs);
        }
    };
    class ppIDescriptorDistance {
    public:
        virtual float Compute(const cv::Mat& descr1, const cv::Mat& descr2) = 0;
        virtual std::vector<float> Compute(const std::vector<cv::Mat>& descrs1, const std::vector<cv::Mat>& descrs2) = 0;
        virtual ~ppIDescriptorDistance() {}
    };
    class ppCosDistance : public ppIDescriptorDistance {
    public:
        explicit ppCosDistance(const cv::Size& descriptor_size);
        float Compute(const cv::Mat& descr1, const cv::Mat& descr2) override;
        std::vector<float> Compute(const std::vector<cv::Mat>& descrs1, const std::vector<cv::Mat>& descrs2) override;

    private:
        cv::Size descriptor_size_;
    };
    class ppMatchTemplateDistance : public ppIDescriptorDistance {
    public:
        ppMatchTemplateDistance(int type = cv::TemplateMatchModes::TM_CCORR_NORMED, float scale = -1, float offset = 1)
            : type_(type),
            scale_(scale),
            offset_(offset) {}
        float Compute(const cv::Mat& descr1, const cv::Mat& descr2) override;
        std::vector<float> Compute(const std::vector<cv::Mat>& descrs1, const std::vector<cv::Mat>& descrs2) override;
        virtual ~ppMatchTemplateDistance() {}

    private:
        int type_;  ///< Method of MatchTemplate function computation.
        float scale_;  ///< Scale parameter for the distance. Final distance is
        float offset_;  ///< Offset parameter for the distance. Final distance is
    };
    class PedestrianTracker {
    public:
        using Descriptor = std::shared_ptr<ppIImageDescriptor>;
        using Distance = std::shared_ptr<ppIDescriptorDistance>;
        explicit PedestrianTracker(const ppTrackerParams& params = ppTrackerParams());
        virtual ~PedestrianTracker() {}
        void Process(const cv::Mat& frame, const std::deque<ppTrackedObject>& detections, uint64_t timestamp);
        const ppTrackerParams& params() const;
        const Descriptor& descriptor_fast() const;
        void set_descriptor_fast(const Descriptor& val);
        const Descriptor& descriptor_strong() const;
        void set_descriptor_strong(const Descriptor& val);
        const Distance& distance_fast() const;
        void set_distance_fast(const Distance& val);
        const Distance& distance_strong() const;
        void set_distance_strong(const Distance& val);
        ppDetectionLog GetDetectionLog(const bool valid_only) const;
        std::unordered_map<size_t, std::vector<cv::Point>> GetActiveTracks() const;
        std::deque<ppTrackedObject> TrackedDetections() const;
        cv::Mat DrawActiveTracks(const cv::Mat& frame);
        bool IsTrackForgotten(size_t id) const;
        const std::unordered_map<size_t, ppTrack>& tracks() const;
        bool IsTrackValid(size_t track_id) const;
        void DropForgottenTracks();

    private:
        struct Match {
            int64_t frame_idx1;
            int64_t frame_idx2;
            cv::Rect rect1;
            cv::Rect rect2;
            cv::Rect pr_rect1;
            bool pr_label;
            bool gt_label;

            Match() {}

            Match(const ppTrackedObject& a, const cv::Rect& a_pr_rect, const ppTrackedObject& b, bool pr_label)
                : frame_idx1(a.frame_idx),
                frame_idx2(b.frame_idx),
                rect1(a.rect),
                rect2(b.rect),
                pr_rect1(a_pr_rect),
                pr_label(pr_label),
                gt_label(a.object_id == b.object_id) {
                PT_CHECK_NE(frame_idx1, frame_idx2);
            }
        };

        const ppObjectTracks all_tracks(bool valid_only) const;
        static float ShapeAffinity(float w, const cv::Rect& trk, const cv::Rect& det);
        static float MotionAffinity(float w, const cv::Rect& trk, const cv::Rect& det);

        static float TimeAffinity(float w, const float& trk, const float& det);

        cv::Rect PredictRect(size_t id, size_t k, size_t s) const;

        cv::Rect PredictRectSmoothed(size_t id, size_t k, size_t s) const;

        cv::Rect PredictRectSimple(size_t id, size_t k, size_t s) const;

        void SolveAssignmentProblem(const std::set<size_t>& track_ids,
            const std::deque<ppTrackedObject>& detections,
            const std::vector<cv::Mat>& descriptors,
            float thr,
            std::set<size_t>* unmatched_tracks,
            std::set<size_t>* unmatched_detections,
            std::set<std::tuple<size_t, size_t, float>>* matches);

        void ComputeFastDesciptors(const cv::Mat& frame,
            const std::deque<ppTrackedObject>& detections,
            std::vector<cv::Mat>* descriptors);

        void ComputeDissimilarityMatrix(const std::set<size_t>& active_track_ids,
            const std::deque<ppTrackedObject>& detections,
            const std::vector<cv::Mat>& fast_descriptors,
            cv::Mat* dissimilarity_matrix);

        std::vector<float> ComputeDistances(const cv::Mat& frame,
            const std::deque<ppTrackedObject>& detections,
            const std::vector<std::pair<size_t, size_t>>& track_and_det_ids,
            std::map<size_t, cv::Mat>* det_id_to_descriptor);

        std::map<size_t, std::pair<bool, cv::Mat>> StrongMatching(
            const cv::Mat& frame,
            const std::deque<ppTrackedObject>& detections,
            const std::vector<std::pair<size_t, size_t>>& track_and_det_ids);

        std::vector<std::pair<size_t, size_t>> GetTrackToDetectionIds(
            const std::set<std::tuple<size_t, size_t, float>>& matches);

        float AffinityFast(const cv::Mat& descriptor1,
            const ppTrackedObject& obj1,
            const cv::Mat& descriptor2,
            const ppTrackedObject& obj2);

        float Affinity(const ppTrackedObject& obj1, const ppTrackedObject& obj2);

        void AddNewTrack(const cv::Mat& frame,
            const ppTrackedObject& detection,
            const cv::Mat& fast_descriptor,
            const cv::Mat& descriptor_strong = cv::Mat());

        void AddNewTracks(const cv::Mat& frame,
            const std::deque<ppTrackedObject>& detections,
            const std::vector<cv::Mat>& descriptors_fast);

        void AddNewTracks(const cv::Mat& frame,
            const std::deque<ppTrackedObject>& detections,
            const std::vector<cv::Mat>& descriptors_fast,
            const std::set<size_t>& ids);

        void AppendToTrack(const cv::Mat& frame,
            size_t track_id,
            const ppTrackedObject& detection,
            const cv::Mat& descriptor_fast,
            const cv::Mat& descriptor_strong);

        bool EraseTrackIfBBoxIsOutOfFrame(size_t track_id);

        bool EraseTrackIfItWasLostTooManyFramesAgo(size_t track_id);

        bool UpdateLostTrackAndEraseIfItsNeeded(size_t track_id);

        void UpdateLostTracks(const std::set<size_t>& track_ids);

        const std::set<size_t>& active_track_ids() const;

        std::deque<ppTrackedObject> FilterDetections(const std::deque<ppTrackedObject>& detections) const;
        bool IsTrackForgotten(const ppTrack& track) const;

        // Parameters of the pipeline.
        ppTrackerParams params_;

        // Indexes of active tracks.
        std::set<size_t> active_track_ids_;

        // Descriptor fast (base classifier).
        Descriptor descriptor_fast_;

        // Distance fast (base classifier).
        Distance distance_fast_;

        // Descriptor strong (reid classifier).
        Descriptor descriptor_strong_;

        // Distance strong (reid classifier).
        Distance distance_strong_;

        // All tracks.
        std::unordered_map<size_t, ppTrack> tracks_;

        // Previous frame image.
        cv::Size prev_frame_size_;

        struct pair_hash {
            std::size_t operator()(const std::pair<size_t, size_t>& p) const {
                PT_CHECK(p.first < 1e6 && p.second < 1e6);
                return static_cast<size_t>(p.first * 1e6 + p.second);
            }
        };

        // Distance between current active tracks.
        std::unordered_map<std::pair<size_t, size_t>, float, pair_hash> tracks_dists_;

        // Number of all current tracks.
        size_t tracks_counter_;
        cv::Size frame_size_;
        std::vector<cv::Scalar> colors_;
        uint64_t prev_timestamp_;
    };
}
#endif