#include <iostream>
#include <sstream>
#include <vector>
#include <string>
#include <chrono>
#include <opencv2/opencv.hpp>
#include "ANSCustomFireNSmoke.h"

int FireNSmokeDetection() {
    std::string modelDirectory = "C:\\Projects\\ANSVIS\\Models\\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)
    //std::string videoFilePath = "C:\\Programs\\DemoAssets\\Videos\\FireNSmoke\\Failed\\fire 25cm 480p.mp4"; // passed (fire is too small and low resolution cam)

    //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 camera 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 optimize = 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)) {
            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 main() {
    return FireNSmokeDetection();
}