#ifndef ARMORDETECTOR_H
#define ARMORDETECTOR_H

#include <opencv2/opencv.hpp>
#include <vector>
#include <string>
#include "config.h"

struct LightBar {
    cv::Point2f center;
    cv::Size2f size;
    float angle;          // in degrees
    float angle_rad;      // in radians
    float area;
    std::vector<cv::Point2f> center_line;  // [end1, end2]
    float center_line_length;
    std::vector<cv::Point2f> box;
};

struct ArmorPlate {
    std::string color;
    cv::Point2f center;
    double confidence;
    std::pair<LightBar, LightBar> pair;
    std::vector<cv::Point2f> corners_2d;  // Can be computed later for 3D pose
};

class ArmorDetector {
public:
    ArmorDetector();
    
    // Main detection function
    void detect(const cv::Mat& mask, const std::string& target_color, 
                std::vector<LightBar>& valid_light_bars, 
                std::vector<ArmorPlate>& armor_plates);

private:
    float angle_threshold_rad;
    
    // Helper methods
    std::vector<LightBar> extract_initial_light_bars(const cv::Mat& mask);
    std::vector<std::vector<cv::Point2f>> merge_nearby_light_bars(const std::vector<LightBar>& initial_light_bars);
    std::vector<LightBar> filter_valid_light_bars(const std::vector<std::vector<cv::Point2f>>& processed_boxes);
    std::vector<ArmorPlate> pair_light_bars_to_armor(const std::vector<LightBar>& light_bars, const std::string& target_color);
    double calculate_iou(const cv::Rect& b1, const cv::Rect& b2);
};

#endif // ARMORDETECTOR_H