#include <iostream>
#include <fstream>
#include <opencv2/opencv.hpp>
#include "face_detection.h"
#include "face_mask.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

std::string kmodel_path = "face_detection_320.kmodel"; // Path to the compiled kmodel file
float obj_thres = 0.6;                  // Confidence threshold for detections
float nms_thres = 0.3;                  // Non-Maximum Suppression threshold

std::string mask_kmodel_path = "face_mask.kmodel";
float mask_thres = 0.5;
int debug_mode = 0;                     // Debug mode flag (0 = off)

// Face detection model pointer
FaceDetection *face_detection = nullptr;
// Face emotion model pointer
FaceMask *face_mask = nullptr;


// Storage for detection results
std::vector<FaceDetectionInfo> face_results;
// Storage for mask results
FaceMaskInfo mask_result;


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


void setup()
{
    pl = new PipeLine(sensor_id, display_type, ai_frame_width, ai_frame_height, display_width, display_height);
    face_detection = new FaceDetection(kmodel_path, obj_thres, nms_thres,{ai_frame_height,ai_frame_width,3}, debug_mode);
    face_mask = new FaceMask(mask_kmodel_path, mask_thres,{ai_frame_height,ai_frame_width,3}, debug_mode);
}

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
    face_results.clear();

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

    // Run face detection inference on the input tensor
    face_detection->run(input_tensor,face_results);

    for(auto &face:face_results){
        face_mask->run(input_tensor,face.sparse_kps.points,mask_result);
        face_mask->draw_results(pl->osd_frame, face.bbox, mask_result);
    }

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

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