#ifndef ANSLPROV_H
#define ANSLPROV_H
#pragma once
#include "ANSLPR.h"
#include <list>
#include <map>
#include <string>
#include <utility>
#include <vector>

#include "openvino/openvino.hpp"
#include "openvino/runtime/intel_gpu/properties.hpp"
#include "utils/args_helper.hpp"
#include "utils/input_wrappers.hpp"
#include "utils/common.hpp"
#include "utils/ocv_common.hpp"
#include "utils/threads_common.hpp"
#include "utils/grid_mat.hpp"
#include "monitors/presenter.h"

namespace ANSCENTER
{
    void tryPush(const std::weak_ptr<Worker>& worker, std::shared_ptr<Task>&& task);

    struct BboxAndDescr {
        enum class ObjectType {
            NONE,
            VEHICLE,
            PLATE,
        } objectType;
        cv::Rect rect;
        std::string descr;
    };

    class Detector {
    public:
        struct Result {
            std::size_t label;
            float confidence;
            cv::Rect location;
        };
        static constexpr int maxProposalCount = 200;
        static constexpr int objectSize = 7;  // Output should have 7 as a last dimension"
        Detector() = default;
        Detector(ov::Core& core, 
                 const std::string& deviceName, 
                 const std::string& xmlPath, 
                 const std::vector<float>& detectionTresholds,
                 const bool autoResize);
        ov::InferRequest createInferRequest();  
        void setImage(ov::InferRequest& inferRequest, const cv::Mat& img);
        std::list<Result> getResults(ov::InferRequest& inferRequest, cv::Size upscale, std::vector<std::string>& rawResults);
        
    private:
        bool m_autoResize;
        std::vector<float> m_detectionTresholds;
        std::string m_detectorInputName;
        std::string m_detectorOutputName;
        ov::CompiledModel m_compiled_model;
    };

    class VehicleAttributesClassifier {
    public:
        VehicleAttributesClassifier() = default;
        VehicleAttributesClassifier(ov::Core& core, 
                                   const std::string& deviceName,
                                   const std::string& xmlPath, 
                                   const bool autoResize);
        ov::InferRequest createInferRequest();
        void setImage(ov::InferRequest& inferRequest, const cv::Mat& img, const cv::Rect vehicleRect);
        std::pair<std::string, std::string> getResults(ov::InferRequest& inferRequest);
    private:
        bool m_autoResize;
        std::string m_attributesInputName;
        std::string m_outputNameForColor;
        std::string m_outputNameForType;
        ov::CompiledModel m_compiled_model;
    };

    class Lpr {
    public:
        Lpr() = default;
        Lpr(ov::Core& core, 
            const std::string& deviceName, 
            const std::string& xmlPath, 
            const bool autoResize);
        ov::InferRequest createInferRequest();
        void setImage(ov::InferRequest& inferRequest, const cv::Mat& img, const cv::Rect plateRect);
        std::string getResults(ov::InferRequest& inferRequest);
        
    private:
        bool m_autoResize;
        int m_maxSequenceSizePerPlate = 0;
        std::string m_LprInputName;
        std::string m_LprInputSeqName;
        std::string m_LprOutputName;
        ov::Layout m_modelLayout;
        ov::CompiledModel m_compiled_model;
    };

    //Utilities
    struct InferRequestsContainer {
        InferRequestsContainer() = default;
        InferRequestsContainer(const InferRequestsContainer&) = delete;
        InferRequestsContainer& operator=(const InferRequestsContainer&) = delete;

        void assign(const std::vector<ov::InferRequest>& inferRequests) {
            actualInferRequests = inferRequests;
            this->inferRequests.container.clear();
            for (auto& ir : this->actualInferRequests) {
                this->inferRequests.container.push_back(ir);
            }
        }
        ConcurrentContainer<std::vector<std::reference_wrapper<ov::InferRequest>>> inferRequests;
        std::vector<ov::InferRequest> actualInferRequests;
    };


    // stores all global data for tasks
    struct Context {
        Context(const std::vector<std::shared_ptr<InputChannel>>& inputChannels,
                const Detector& detector,
                const VehicleAttributesClassifier& vehicleAttributesClassifier, 
                const Lpr& lpr,
                uint64_t lastFrameId,
                uint64_t nireq,
                bool isVideo,
                std::size_t nclassifiersireq, 
                std::size_t nrecognizersireq) :
                readersContext{ inputChannels, std::vector<int64_t>(inputChannels.size(), -1), std::vector<std::mutex>(inputChannels.size()) },
                inferTasksContext{ detector },
                detectionsProcessorsContext{ vehicleAttributesClassifier, lpr },
                videoFramesContext{ std::vector<uint64_t>(inputChannels.size(), lastFrameId), std::vector<std::mutex>(inputChannels.size()) },
                nireq{ nireq },
                isVideo{ isVideo },
                freeDetectionInfersCount{ 0 },
                frameCounter{ 0 }
        {
            std::vector<ov::InferRequest> detectorInferRequests;
            std::vector<ov::InferRequest> attributesInferRequests;
            std::vector<ov::InferRequest> lprInferRequests;
            detectorInferRequests.reserve(nireq);
            attributesInferRequests.reserve(nclassifiersireq);
            lprInferRequests.reserve(nrecognizersireq);
            std::generate_n(std::back_inserter(detectorInferRequests), nireq, [&] {return inferTasksContext.detector.createInferRequest(); });
            std::generate_n(std::back_inserter(attributesInferRequests), nclassifiersireq, [&] {return detectionsProcessorsContext.vehicleAttributesClassifier.createInferRequest(); });
            std::generate_n(std::back_inserter(lprInferRequests), nrecognizersireq, [&] {return detectionsProcessorsContext.lpr.createInferRequest(); });
            detectorsInfers.assign(detectorInferRequests);
            attributesInfers.assign(attributesInferRequests);
            platesInfers.assign(lprInferRequests);
        }
        struct {
            std::vector<std::shared_ptr<InputChannel>> inputChannels;
            std::vector<int64_t> lastCapturedFrameIds;
            std::vector<std::mutex> lastCapturedFrameIdsMutexes;
            std::weak_ptr<Worker> readersWorker;
        } readersContext;

        struct {
            Detector detector;
            std::weak_ptr<Worker> inferTasksWorker;
        } inferTasksContext;

        struct {
            VehicleAttributesClassifier vehicleAttributesClassifier;
            Lpr lpr;
            std::weak_ptr<Worker> detectionsProcessorsWorker;
        } detectionsProcessorsContext;
        struct {
            std::vector<uint64_t> lastframeIds;
            std::vector<std::mutex> lastFrameIdsMutexes;
        } videoFramesContext;

        std::weak_ptr<Worker> resAggregatorsWorker;
        std::mutex classifiersAggregatorPrintMutex;
        uint64_t nireq;
        bool isVideo;
        std::atomic<std::vector<ov::InferRequest>::size_type> freeDetectionInfersCount;
        std::atomic<uint32_t> frameCounter;
        InferRequestsContainer detectorsInfers, attributesInfers, platesInfers;
        PerformanceMetrics metrics;
        std::list<BboxAndDescr> boxesAndDescrs;
    };


    
    
    // End of Context
    class ReborningVideoFrame : public VideoFrame {
    public:
        ReborningVideoFrame(Context& context, 
                            const unsigned sourceID, 
                            const int64_t frameId, 
                            const cv::Mat& frame = cv::Mat()) :
                            VideoFrame{ sourceID, frameId, frame },
                            context(context) {}  
        virtual ~ReborningVideoFrame();
        Context& context;
    };

    // draws results on the frame
    class ResAggregator : public Task {
    public:
        ResAggregator(const VideoFrame::Ptr& sharedVideoFrame, 
                    std::list<BboxAndDescr>&& boxesAndDescrs) :
                    Task{ sharedVideoFrame, 4.0 }, 
                    boxesAndDescrs{ std::move(boxesAndDescrs) } {}

        bool isReady() override {
            return true;
        }
        void process() override;
    private:
        std::list<BboxAndDescr> boxesAndDescrs;
    };

    // waits for all classifiers and recognisers accumulating results
    class ClassifiersAggregator {
    public:
        std::vector<std::string> rawDetections;
        ConcurrentContainer<std::list<std::string>> rawAttributes;
        ConcurrentContainer<std::list<std::string>> rawDecodedPlates;

        explicit ClassifiersAggregator(const VideoFrame::Ptr& sharedVideoFrame) :
            sharedVideoFrame{ sharedVideoFrame } {}
        ~ClassifiersAggregator() {
            std::mutex& printMutex = static_cast<ReborningVideoFrame*>(sharedVideoFrame.get())->context.classifiersAggregatorPrintMutex;
            printMutex.lock();
            if (!rawDetections.empty()) {
                slog::debug << "Frame #: " << sharedVideoFrame->frameId << slog::endl;
                slog::debug << rawDetections;
                // destructor assures that none uses the container
                for (const std::string& rawAttribute : rawAttributes.container) {
                    slog::debug << rawAttribute << slog::endl;
                }
                for (const std::string& rawDecodedPlate : rawDecodedPlates.container) {
                    slog::debug << rawDecodedPlate << slog::endl;
                }
            }
            printMutex.unlock();
            tryPush(static_cast<ReborningVideoFrame*>(sharedVideoFrame.get())->context.resAggregatorsWorker,
                std::make_shared<ResAggregator>(sharedVideoFrame, std::move(boxesAndDescrs)));
        }

        void push(BboxAndDescr&& bboxAndDescr) {
            boxesAndDescrs.lockedPushBack(std::move(bboxAndDescr));
        }
        const VideoFrame::Ptr sharedVideoFrame;

    private:
        ConcurrentContainer<std::list<BboxAndDescr>> boxesAndDescrs;
    };

    // extracts detections from blob InferRequests and runs classifiers and recognisers
    class DetectionsProcessor : public Task {
    public:
        DetectionsProcessor(VideoFrame::Ptr sharedVideoFrame, ov::InferRequest* inferRequest) :
                            Task{ sharedVideoFrame, 1.0 }, 
                            inferRequest{ inferRequest }, 
                            requireGettingNumberOfDetections{ true } 
        {

        }

        DetectionsProcessor(VideoFrame::Ptr sharedVideoFrame, 
                            std::shared_ptr<ClassifiersAggregator>&& classifiersAggregator, 
                            std::list<cv::Rect>&& vehicleRects,
                            std::list<cv::Rect>&& plateRects) :
                            Task{ sharedVideoFrame, 1.0 }, 
                            classifiersAggregator{ std::move(classifiersAggregator) }, 
                            inferRequest{ nullptr },
                            vehicleRects{ std::move(vehicleRects) }, 
                            plateRects{ std::move(plateRects) },
                            requireGettingNumberOfDetections{ false } 
        {

        }

        bool isReady() override;
        void process() override;

    private:
        std::shared_ptr<ClassifiersAggregator> classifiersAggregator; // when no one stores this object we will draw
        ov::InferRequest* inferRequest;
        std::list<cv::Rect> vehicleRects;
        std::list<cv::Rect> plateRects;
        std::vector<std::reference_wrapper<ov::InferRequest>> reservedAttributesRequests;
        std::vector<std::reference_wrapper<ov::InferRequest>> reservedLprRequests;
        bool requireGettingNumberOfDetections;

    };
    // runs detection
    class InferTask : public Task {
    public:
        explicit InferTask(VideoFrame::Ptr sharedVideoFrame) :
            Task{ sharedVideoFrame, 5.0 } {}
        bool isReady() override;
        void process() override;
    };

    class Reader : public Task {
    public:
        explicit Reader(VideoFrame::Ptr sharedVideoFrame) :
            Task{ sharedVideoFrame, 2.0 } {}
        bool isReady() override;
        void process() override;
    };


	class ANSLPR_API ANSALPR_OV :public ANSALPR {
	private:
        std::string _vehicleLPModel; //model that detects both vehicle and license plate
        std::string _vehicleAtModel; //model that detects vehicle attributes
        std::string _lprModel;       //model that recognise license plate
        Detector* _detector = nullptr;
        VehicleAttributesClassifier* _vehicleAttributesClassifier = nullptr;
        Lpr* _lpr = nullptr;
        [[nodiscard]] std::string VectorDetectionToJsonString(const std::vector<ALPRObject>& dets);
	public:
		ANSALPR_OV();
		~ANSALPR_OV();
		[[nodiscard]] bool Initialize(const std::string& licenseKey, const std::string& modelZipFilePath, const std::string& modelZipPassword) override;
		[[nodiscard]] bool Inference(const cv::Mat& input, std::string& lprResult);
		[[nodiscard]] bool Inference(const cv::Mat& input, const std::vector<cv::Rect>& Bbox, std::string& lprResult);
        [[nodiscard]] bool Destroy() override;
	};
}
#endif