#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>
#include "yolov8.h"
#include "utils.h"
#include "pipeline.h"

// Sensor and display configuration
int sensor_id = 2;                   // ID of the camera sensor
std::string display_type = "ST7701"; // Display type
size_t ai_frame_width = 640;         // Width of AI input frame
size_t ai_frame_height = 360;        // Height of AI input frame
size_t display_width = 800;          // Display width
size_t display_height = 480;         // Display height

// YOLO model and task configuration
std::string model_type = "yolov8";       // Model type
std::string task_type = "detect";        // Task type: detection
std::string task_mode = "video";         // Task mode: video input
std::string kmodel_path = "yolov8n_224.kmodel"; // Path to the compiled kmodel file
float conf_thres = 0.35;                 // Confidence threshold for detections
float nms_thres = 0.65;                  // Non-Maximum Suppression threshold
float mask_thres = 0.5;                  // Mask threshold (for segmentation)
std::vector<std::string> labels = {      // List of class labels
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
    "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
    "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
    "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
    "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
    "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
    "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
    "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"
};
int debug_mode = 0;               // Debug mode flag (0 = off)

// YOLOv8 model pointer
Yolov8 *yolov8 = nullptr;

// Storage for detection results
std::vector<YOLOBbox> yolo_results;

// Pipeline pointer for camera capture and display
PipeLine* pl;

// Setup function: initialize pipeline and YOLO model
void setup() {
    pl = new PipeLine(sensor_id, display_type, ai_frame_width, ai_frame_height, display_width, display_height);
    yolov8 = new Yolov8(task_type, task_mode, kmodel_path, {ai_frame_width, ai_frame_height}, labels, conf_thres, nms_thres, mask_thres, debug_mode);
}

// Main loop: capture, run detection, draw results, and display
void loop() {
    PROFILE_SCOPE("total:", 1);       // Profiling scope for performance measurement

    // Get the input tensor from the pipeline
    runtime_tensor input_tensor = pl->get_tensor();

    // Clear previous detection results
    yolo_results.clear();

    // Clear the OSD frame (overlay) with transparent black
    pl->osd_frame.setTo(cv::Scalar(0, 0, 0, 0));

    // Run YOLOv8 inference on the input tensor
    yolov8->run(input_tensor, yolo_results);

    // Draw detection results onto the OSD frame
    yolov8->draw_results(pl->osd_frame, yolo_results);

    // Release the input tensor back to the pipeline
    pl->release_tensor();

    // Display the final image with detections
    pl->show_image();
}
