#include <iostream>
#include "ANSCustomCodeWeaponDetection.h"
#include "ANSCustomCodeHelmetDetection.h"
#include "ANSCustomFireNSmoke.h"
#include "ANSCustomViolence.h"
int WeaponDetection() {
    std::string modelDirectory = "C:\\Programs\\DemoAssets\\ModelsForANSVIS\\ANS_WeaponDetection(GPU)_1.0";// "C:\\Programs\\DemoAssets\\ModelsForANSVIS\\CustomModels\\ANS_WeaponDetection(GPU)_1.0";
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\Weapon\\AK47Noweapon.mp4"; // passed
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\Weapon\\BMG AK47 inside 2.mp4"; // passed
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\Weapon\\BMGWeapon.mp4"; // passed but need to improve the model

    std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\Weapon\\E112 Weapon.mp4"; // failed (too small weapon)
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\Weapon\\Glock & knife.mp4"; // failed (too small weapon)

     //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\Weapon\\Knifes & No weapon.mp4"; // failed (too small weapon)
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\Weapon\\Knifes.mp4"; // pass
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\Weapon\\No weapon.mp4"; // pass
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\Weapon\\AK47 Glock.mp4"; // pass

    std::vector<std::string> classes;
    ANSCustomWD detector;
    std::string labelMap;

    // Initialize the detector and parse the label map
    detector.Initialize(modelDirectory, 0.6,labelMap);
    std::stringstream ss(labelMap);
    while (ss.good()) {
        std::string substr;
        getline(ss, substr, ',');
        classes.push_back(substr);
    }

    // Open video capture
    cv::VideoCapture capture(videoFilePath);
    if (!capture.isOpened()) {
        std::cerr << "Error: Unable to open video file!" << std::endl;
        return -1;
    }

    // Get the original video dimensions
    int originalWidth = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_WIDTH));
    int originalHeight = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_HEIGHT));
    double aspectRatio = static_cast<double>(originalWidth) / originalHeight;

    // Create a resizable window
    const std::string windowName = "ANS Weapon Detection";
    cv::namedWindow(windowName, cv::WINDOW_NORMAL);

    while (true) {
        cv::Mat frame;
        if (!capture.read(frame)) {
            // Reset to the beginning of the video
            capture.set(cv::CAP_PROP_POS_FRAMES, 0);
            continue;
        }

        // Run inference on the frame
        std::vector<CustomObject> customObjects;
        auto start = std::chrono::system_clock::now();
        customObjects = detector.RunInference(frame);
        auto end = std::chrono::system_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        printf("Time = %lld ms\n", static_cast<long long int>(elapsed.count()));

        // Draw detection results on the frame
        if (!customObjects.empty()) {
            for (const auto& obj : customObjects) {
                cv::rectangle(frame, cv::Rect(obj.box.x, obj.box.y, obj.box.width, obj.box.height), cv::Scalar(123, 255, 123), 2);
                cv::putText(frame,
                    cv::format("%s:%d:%.2f", classes[obj.classId].c_str(), obj.classId, obj.confidence),
                    cv::Point(obj.box.x, obj.box.y - 5),
                    cv::FONT_HERSHEY_SIMPLEX,
                    0.6,
                    cv::Scalar(0, 0, 255),
                    1,
                    cv::LINE_AA);
            }
        }

        // Handle resizing and maintain aspect ratio
        int currentWidth = static_cast<int>(cv::getWindowProperty(windowName, cv::WND_PROP_AUTOSIZE));
        int currentHeight = static_cast<int>(cv::getWindowProperty(windowName, cv::WND_PROP_AUTOSIZE));

        if (currentWidth > 0 && currentHeight > 0) {
            int newWidth = currentWidth;
            int newHeight = static_cast<int>(newWidth / aspectRatio);
            if (newHeight > currentHeight) {
                newHeight = currentHeight;
                newWidth = static_cast<int>(newHeight * aspectRatio);
            }
            cv::resize(frame, frame, cv::Size(newWidth, newHeight));
        }

        // Display the frame
        cv::imshow(windowName, frame);

        // Exit the loop if 'esc' key is pressed
        if (cv::waitKey(30) == 27) {
            break;
        }
    }

    // Clean up
    capture.release();
    cv::destroyAllWindows();
    return 0;
}
int FireNSmokeDetection() {
    std::string modelDirectory = "E:\\Programs\\DemoAssets\\ModelsForANSVIS\\WorkingCustomModels\\ANS_FireSmoke_v2.0";
	std::string videoFilePath = "E:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\ANSFireFull.mp4";// passed
    //std::string videoFilePath = "E:\\Programs\\DemoAssets\\Videos\\ANSVIS_Issues\\FGFire.mp4";// passed
    //std::string videoFilePath = "E:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\HFire1.mp4"; // passed
    //std::string videoFilePath = "E:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\HFire3.mp4"; // passed
    //std::string videoFilePath = "C:\\Programs\\FireTest\\road.mp4"; // passed
    //std::string videoFilePath = "C:\\Programs\\FireTest\\video_20.mp4"; // passed
    //std::string videoFilePath = "C:\\Programs\\FireTest\\BCA1.mp4"; // passed
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\E112 v4-1.mp4"; // passed
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\E112Full.mp4"; // passed
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\SimFire.mp4"; // passed
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\BCA3.mp4"; // pass (smoke issue)
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\E112 fire.mp4"; // pass
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\e112 v3-2.mp4"; // pass
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\e112 v3-3.mp4"; // pass
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\e112 v3-4.mp4"; // pass
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\E112 v4-1.mp4"; // pass
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\Fire25m.mp4"; // passed (fire is too small and low resolution cam), might need to have  polity to further divide the image into smaller sections if it is small and no detections found.
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\Failed\\fire 25cm 480p.mp4"; // passed (fire is too small and low resolution cam), might need to have  polity to further divide the image into smaller sections if it is small and no detections found.


    //std::string videoFilePath = "C:\\Programs\\FireTest\\BCA2.mp4"; // failed (fire is too small and low resolution cam)
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\fire4.mp4"; // failed (smoke issue, too low resolution)
    //std::string videoFilePath = "C:\\Programs\\FireTest\\HFire2.mp4"; // failed as camerea is moving

   // std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\E112 v4-2.mp4"; // failed white smoke low resolution cam

    std::vector<std::string> classes;
    ANSCustomFS detector;
    std::string labelMap;

    // Optimize model  
	bool optimze = detector.OptimizeModel(true);


    // Initialize the detector and parse the label map
    detector.Initialize(modelDirectory, 0.5, labelMap);
    std::stringstream ss(labelMap);
    while (ss.good()) {
        std::string substr;
        getline(ss, substr, ',');
        classes.push_back(substr);
    }

    // Open video capture
    cv::VideoCapture capture(videoFilePath);
    if (!capture.isOpened()) {
        std::cerr << "Error: Unable to open video file!" << std::endl;
        return -1;
    }

    // Get the original video dimensions
    int originalWidth = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_WIDTH));
    int originalHeight = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_HEIGHT));
    double aspectRatio = static_cast<double>(originalWidth) / originalHeight;

    // Create a resizable window
    const std::string windowName = "ANS Fire Detection";
    cv::namedWindow(windowName, cv::WINDOW_NORMAL);

    while (true) {
        cv::Mat frame;
        if (!capture.read(frame)) {
            // Reset to the beginning of the video
            //capture.set(cv::CAP_PROP_POS_FRAMES, 0);
            //continue;
            break;
        }

        // Run inference on the frame
        std::vector<CustomObject> customObjects;
        auto start = std::chrono::system_clock::now();
        customObjects = detector.RunInference(frame);
        auto end = std::chrono::system_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        printf("Time = %lld ms\n", static_cast<long long int>(elapsed.count()));

        // Draw detection results on the frame
        if (!customObjects.empty()) {
            for (const auto& obj : customObjects) {
                cv::rectangle(frame, cv::Rect(obj.box.x, obj.box.y, obj.box.width, obj.box.height), cv::Scalar(123, 255, 123), 2);
                cv::putText(frame,
                    cv::format("%s:%d:%.2f", classes[obj.classId].c_str(), obj.classId, obj.confidence),
                    cv::Point(obj.box.x, obj.box.y - 5),
                    cv::FONT_HERSHEY_SIMPLEX,
                    0.6,
                    cv::Scalar(0, 0, 255),
                    1,
                    cv::LINE_AA);
            }
        }

        // Handle resizing and maintain aspect ratio
        int currentWidth = static_cast<int>(cv::getWindowProperty(windowName, cv::WND_PROP_AUTOSIZE));
        int currentHeight = static_cast<int>(cv::getWindowProperty(windowName, cv::WND_PROP_AUTOSIZE));

        if (currentWidth > 0 && currentHeight > 0) {
            int newWidth = currentWidth;
            int newHeight = static_cast<int>(newWidth / aspectRatio);
            if (newHeight > currentHeight) {
                newHeight = currentHeight;
                newWidth = static_cast<int>(newHeight * aspectRatio);
            }
            cv::resize(frame, frame, cv::Size(newWidth, newHeight));
        }

        // Display the frame
        cv::imshow(windowName, frame);

        // Exit the loop if 'esc' key is pressed
        if (cv::waitKey(30) == 27) {
            break;
        }
    }

    // Clean up
    capture.release();
    cv::destroyAllWindows();
    return 0;
}
int ViolenceDetection() {
    std::string modelDirectory = "E:\\Programs\\DemoAssets\\ModelsForANSVIS\\WorkingCustomModels\\ANS_Violence(GPU)_v1.0";
    std::string videoFilePath = "E:\\Programs\\DemoAssets\\Videos\\Violence\\violence1.mp4";// passed
    
    std::vector<std::string> classes;
    ANSCustomViolence detector;
    std::string labelMap;

    // Initialize the detector and parse the label map
    detector.Initialize(modelDirectory, 0.7, labelMap);
    std::stringstream ss(labelMap);
    while (ss.good()) {
        std::string substr;
        getline(ss, substr, ',');
        classes.push_back(substr);
    }

    // Open video capture
    cv::VideoCapture capture(videoFilePath);
    if (!capture.isOpened()) {
        std::cerr << "Error: Unable to open video file!" << std::endl;
        return -1;
    }

    // Get the original video dimensions
    int originalWidth = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_WIDTH));
    int originalHeight = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_HEIGHT));
    double aspectRatio = static_cast<double>(originalWidth) / originalHeight;

    // Create a resizable window
    const std::string windowName = "ANS Violence Detection";
    cv::namedWindow(windowName, cv::WINDOW_NORMAL);

    while (true) {
        cv::Mat frame;
        if (!capture.read(frame)) {
            // Reset to the beginning of the video
            //capture.set(cv::CAP_PROP_POS_FRAMES, 0);
            //continue;
            break;
        }

        // Run inference on the frame
        std::vector<CustomObject> customObjects;
        auto start = std::chrono::system_clock::now();
        customObjects = detector.RunInference(frame);
        auto end = std::chrono::system_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        printf("Time = %lld ms\n", static_cast<long long int>(elapsed.count()));

        // Draw detection results on the frame
        if (!customObjects.empty()) {
            for (const auto& obj : customObjects) {
                cv::rectangle(frame, cv::Rect(obj.box.x, obj.box.y, obj.box.width, obj.box.height), cv::Scalar(123, 255, 123), 2);
                cv::putText(frame,
                    cv::format("%s:%d:%.2f", classes[obj.classId].c_str(), obj.classId, obj.confidence),
                    cv::Point(obj.box.x, obj.box.y - 5),
                    cv::FONT_HERSHEY_SIMPLEX,
                    0.6,
                    cv::Scalar(0, 0, 255),
                    1,
                    cv::LINE_AA);
            }
        }

        // Handle resizing and maintain aspect ratio
        int currentWidth = static_cast<int>(cv::getWindowProperty(windowName, cv::WND_PROP_AUTOSIZE));
        int currentHeight = static_cast<int>(cv::getWindowProperty(windowName, cv::WND_PROP_AUTOSIZE));

        if (currentWidth > 0 && currentHeight > 0) {
            int newWidth = currentWidth;
            int newHeight = static_cast<int>(newWidth / aspectRatio);
            if (newHeight > currentHeight) {
                newHeight = currentHeight;
                newWidth = static_cast<int>(newHeight * aspectRatio);
            }
            cv::resize(frame, frame, cv::Size(newWidth, newHeight));
        }

        // Display the frame
        cv::imshow(windowName, frame);

        // Exit the loop if 'esc' key is pressed
        if (cv::waitKey(30) == 27) {
            break;
        }
    }
    // Clean up
    capture.release();
    cv::destroyAllWindows();
    return 0;
}
int ReadCorruptImage() {
    std::string imageFilePath = "C:\\ProgramData\\ANSCENTER\\ANSVIS Server\\InferenceError\\1.jpg";
    cv::Mat image = cv::imread(imageFilePath);
    if (image.empty()) {
        std::cout << "Error: Unable to read image file!" << std::endl;
        return -1;
    }
    cv::imshow("Corrupt Image", image);
    cv::waitKey(0);
    return 0;

}
int HelmetDetection() {
    std::string modelDirectory = "E:\\Programs\\DemoAssets\\ModelsForANSVIS\\WorkingCustomModels\\ANS_Helmet(GPU)_v1.0";// "C:\\Users\\nghia\\AppData\\Local\\Temp\\Models\\EngineModels\\ANS_Helmet_v2.0";// "C:\\Programs\\DemoAssets\\ModelsForANSVIS\\CustomModels\\ANS_WeaponDetection(GPU)_1.0";
    std::string videoFilePath = "E:\\Programs\\DemoAssets\\Videos\\Helmet\\HM2.mp4"; // failed (too small weapon)
    std::vector<std::string> classes;
    ANSCustomHMD detector;
    std::string labelMap;

    // Initialize the detector and parse the label map
    detector.Initialize(modelDirectory, 0.6, labelMap);
    std::stringstream ss(labelMap);
    while (ss.good()) {
        std::string substr;
        getline(ss, substr, ',');
        classes.push_back(substr);
    }

    // Open video capture
    cv::VideoCapture capture(videoFilePath);
    if (!capture.isOpened()) {
        std::cerr << "Error: Unable to open video file!" << std::endl;
        return -1;
    }

    // Get the original video dimensions
    int originalWidth = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_WIDTH));
    int originalHeight = static_cast<int>(capture.get(cv::CAP_PROP_FRAME_HEIGHT));
    double aspectRatio = static_cast<double>(originalWidth) / originalHeight;

    // Create a resizable window
    const std::string windowName = "ANS Helmet Detection";
    cv::namedWindow(windowName, cv::WINDOW_NORMAL);

    while (true) {
        cv::Mat frame;
        if (!capture.read(frame)) {
            // Reset to the beginning of the video
            capture.set(cv::CAP_PROP_POS_FRAMES, 0);
            continue;
        }

        // Run inference on the frame
        std::vector<CustomObject> customObjects;
        auto start = std::chrono::system_clock::now();
        customObjects = detector.RunInference(frame);
        auto end = std::chrono::system_clock::now();
        auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - start);
        printf("Time = %lld ms\n", static_cast<long long int>(elapsed.count()));

        // Draw detection results on the frame
        if (!customObjects.empty()) {
            for (const auto& obj : customObjects) {
                cv::rectangle(frame, cv::Rect(obj.box.x, obj.box.y, obj.box.width, obj.box.height), cv::Scalar(123, 255, 123), 2);
                cv::putText(frame,
                    cv::format("%s", obj.className.c_str()),
                    cv::Point(obj.box.x, obj.box.y - 5),
                    cv::FONT_HERSHEY_SIMPLEX,
                    0.6,
                    cv::Scalar(0, 0, 255),
                    1,
                    cv::LINE_AA);
            }
        }

        // Handle resizing and maintain aspect ratio
        int currentWidth = static_cast<int>(cv::getWindowProperty(windowName, cv::WND_PROP_AUTOSIZE));
        int currentHeight = static_cast<int>(cv::getWindowProperty(windowName, cv::WND_PROP_AUTOSIZE));

        if (currentWidth > 0 && currentHeight > 0) {
            int newWidth = currentWidth;
            int newHeight = static_cast<int>(newWidth / aspectRatio);
            if (newHeight > currentHeight) {
                newHeight = currentHeight;
                newWidth = static_cast<int>(newHeight * aspectRatio);
            }
            cv::resize(frame, frame, cv::Size(newWidth, newHeight));
        }

        // Display the frame
        cv::imshow(windowName, frame);

        // Exit the loop if 'esc' key is pressed
        if (cv::waitKey(30) == 27) {
            break;
        }
    }

    // Clean up
    capture.release();
    cv::destroyAllWindows();
    return 0;
}
int main()
{
    HelmetDetection();
    //ViolenceDetection();
    //HelmetDetection();
    //WeaponDetection();
    //ReadCorruptImage();
   // FireNSmokeDetection();
}
