#include <iostream>
#include <string>
#include <vector>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
#include <filesystem>
#include <fstream>

#include "util.h"

// Helper function to check if a file exists
bool fileExists(const std::string &path)
{
    return std::filesystem::exists(path);
}

// Function to read class names from a file
std::vector<std::string> _readClassNameList(const std::string &path)
{
    std::vector<std::string> classes;

    // Check if file exists
    if (!fileExists(path))
    {
        throw std::runtime_error("Class names file not found: " + path);
    }

    // Try to open and read file
    std::ifstream file(path);
    if (!file.is_open())
    {
        throw std::runtime_error("Unable to open class names file: " + path);
    }

    std::string line;
    while (std::getline(file, line))
    {
        if (!line.empty())
        {
            classes.push_back(line);
        }
    }

    if (classes.empty())
    {
        throw std::runtime_error("No classes found in file: " + path);
    }

    return classes;
}

int main(int argc, char *argv[])
{
    try
    {
        // Open Video Capture
        cv::VideoCapture capture = cv::VideoCapture(0);
        if (!capture.isOpened())
        {
            std::cerr << "Failed to open camera device" << std::endl;
            return -1;
        }

        // Read Class Name List and Color Table
        const std::string list = "coco.names";
        const std::vector<std::string> classes = _readClassNameList(list);
        const std::vector<cv::Scalar> colors = getClassColors(classes.size());

        // Debug: Print the size of the colors vector
        std::cout << "Number of colors: " << colors.size() << std::endl;

        // Read Darknet
        const std::string model = "yolov4.weights";
        const std::string config = "yolov4.cfg";
        cv::dnn::Net net = cv::dnn::readNet(model, config);
        if (net.empty())
        {
            std::cerr << "Failed to load network" << std::endl;
            return -1;
        }

        // Set Preferable Backend
        net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);

        // Set Preferable Target
        net.setPreferableTarget(cv::dnn::DNN_TARGET_OPENCL);

        while (true)
        {
            // Read Frame
            cv::Mat frame;
            capture >> frame;
            if (frame.empty())
            {
                cv::waitKey(0);
                break;
            }
            if (frame.channels() == 4)
            {
                cv::cvtColor(frame, frame, cv::COLOR_BGRA2BGR);
            }

            // Create Blob from Input Image
            cv::Mat blob = cv::dnn::blobFromImage(frame, 1 / 255.f, cv::Size(416, 416), cv::Scalar(), true, false);

            // Set Input Blob
            net.setInput(blob);

            // Run Forward Network
            std::vector<cv::Mat> detections;
            net.forward(detections, getOutputsNames(net));

            // Draw Region
            std::vector<int32_t> class_ids;
            std::vector<float> confidences;
            std::vector<cv::Rect> rectangles;
            for (cv::Mat &detection : detections)
            {
                if (detection.empty())
                {
                    std::cerr << "Detection matrix is empty!" << std::endl;
                    continue;
                }
                for (int32_t i = 0; i < detection.rows; i++)
                {
                    cv::Mat region = detection.row(i);

                    // Retrieve Max Confidence and Class Index
                    cv::Mat scores = region.colRange(5, detection.cols);
                    cv::Point class_id;
                    double confidence;
                    cv::minMaxLoc(scores, 0, &confidence, 0, &class_id);

                    // Check Confidence
                    constexpr float threshold = 0.2;
                    if (threshold > confidence)
                    {
                        continue;
                    }

                    // Retrieve Object Position
                    const int32_t x_center = static_cast<int32_t>(region.at<float>(0) * frame.cols);
                    const int32_t y_center = static_cast<int32_t>(region.at<float>(1) * frame.rows);
                    const int32_t width = static_cast<int32_t>(region.at<float>(2) * frame.cols);
                    const int32_t height = static_cast<int32_t>(region.at<float>(3) * frame.rows);
                    const cv::Rect rectangle = cv::Rect(x_center - (width / 2), y_center - (height / 2), width, height);

                    // Add Class ID, Confidence, Rectangle
                    class_ids.push_back(class_id.x);
                    confidences.push_back(confidence);
                    rectangles.push_back(rectangle);
                }
            }

            // Remove Overlap Rectangles using Non-Maximum Suppression
            constexpr float confidence_threshold = 0.5; // Confidence
            constexpr float nms_threshold = 0.5;        // IoU (Intersection over Union)
            std::vector<int32_t> indices;
            cv::dnn::NMSBoxes(rectangles, confidences, confidence_threshold, nms_threshold, indices);

            // Draw Rectangle
            for (const int32_t &index : indices)
            {
                // Bounds checking
                if (class_ids[index] >= colors.size())
                {
                    std::cerr << "Color index out of bounds: " << class_ids[index] << " (max: " << colors.size() - 1 << ")" << std::endl;
                    continue;
                }

                const cv::Rect rectangle = rectangles[index];
                const cv::Scalar color = colors[class_ids[index]];

                // Debug: Print the index and color
                std::cout << "Drawing rectangle with color index: " << class_ids[index] << std::endl;

                constexpr int32_t thickness = 3;
                cv::rectangle(frame, rectangle, color, thickness);

                std::string label = classes[class_ids[index]] + ": " + std::to_string(static_cast<int>(confidences[index] * 100)) + "%";
                
                int baseLine;
                cv::Size labelSize = cv::getTextSize(label, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
                int top = std::max(rectangle.y, labelSize.height);
                cv::rectangle(frame, cv::Point(rectangle.x, top - labelSize.height),
                              cv::Point(rectangle.x + labelSize.width, top + baseLine), color, cv::FILLED);
                cv::putText(frame, label, cv::Point(rectangle.x, top), cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(255, 255, 255), 1);
            }

            // Show Image
            cv::imshow("Object Detection", frame);
            const int32_t key = cv::waitKey(1);
            if (key == 'q')
            {
                break;
            }
        }

        cv::destroyAllWindows();
        return 0;
    }
    catch (const std::exception &e)
    {
        std::cerr << "Error: " << e.what() << std::endl;
        return -1;
    }
}

// cloned and fixed from https://github.com/UnaNancyOwen/OpenCVDNNSample/tree/master