#include "ArmorDetector.h"
#include <iostream>
#include <cmath>
#include <algorithm>

ArmorDetector::ArmorDetector() {
    angle_threshold_rad = HORIZONTAL_ANGLE_THRESHOLD_RAD;
}

void ArmorDetector::detect(const cv::Mat& mask, const std::string& target_color, 
                          std::vector<LightBar>& valid_light_bars, 
                          std::vector<ArmorPlate>& armor_plates) {
    std::vector<LightBar> initial_light_bars = extract_initial_light_bars(mask);
    if (initial_light_bars.empty()) {
        valid_light_bars.clear();
        armor_plates.clear();
        return;
    }

    std::vector<std::vector<cv::Point2f>> processed_boxes = merge_nearby_light_bars(initial_light_bars);
    if (processed_boxes.empty()) {
        valid_light_bars.clear();
        armor_plates.clear();
        return;
    }

    valid_light_bars = filter_valid_light_bars(processed_boxes);
    if (valid_light_bars.size() < 2) {
        armor_plates.clear();
        return;
    }

    armor_plates = pair_light_bars_to_armor(valid_light_bars, target_color);
    if (armor_plates.empty()) {
        return;
    }
}

std::vector<LightBar> ArmorDetector::extract_initial_light_bars(const cv::Mat& mask) {
    std::vector<std::vector<cv::Point>> contours;
    std::vector<cv::Vec4i> hierarchy;
    cv::findContours(mask, contours, hierarchy, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);

    std::vector<LightBar> initial_light_bars;

    for (size_t i = 0; i < contours.size(); i++) {
        if (contours[i].size() < 5) continue; // Need at least 5 points to fit ellipse
        
        cv::RotatedRect rect = cv::minAreaRect(contours[i]);
        cv::Point2f vertices[4];
        rect.points(vertices);
        
        std::vector<cv::Point2f> box;
        for (int j = 0; j < 4; j++) {
            box.push_back(vertices[j]);
        }
        
        float contour_area = cv::contourArea(contours[i]);
        float rect_area = rect.size.width * rect.size.height;

        // Filter: sufficient area + high fill ratio
        if (rect_area > 30 && rect_area > 0 && (contour_area / rect_area) > 0.4) {
            LightBar light_bar;
            light_bar.center = rect.center;
            light_bar.size = rect.size;
            light_bar.angle = rect.angle;
            light_bar.area = rect_area;
            light_bar.box = box;
            
            initial_light_bars.push_back(light_bar);
        }
    }

    return initial_light_bars;
}

std::vector<std::vector<cv::Point2f>> ArmorDetector::merge_nearby_light_bars(const std::vector<LightBar>& initial_light_bars) {
    std::vector<std::vector<cv::Point2f>> processed_boxes;
    std::vector<bool> visited(initial_light_bars.size(), false);

    for (size_t i = 0; i < initial_light_bars.size(); i++) {
        if (visited[i]) continue;

        std::vector<int> nearby_indices = {static_cast<int>(i)};
        cv::Point2f center1 = initial_light_bars[i].center;

        for (size_t j = i + 1; j < initial_light_bars.size(); j++) {
            if (visited[j]) continue;
            
            cv::Point2f center2 = initial_light_bars[j].center;
            float distance = std::sqrt(std::pow(center2.x - center1.x, 2) + std::pow(center2.y - center1.y, 2));

            // Distance + IOU filtering
            if (distance < NEARBY_LIGHT_BAR_THRESHOLD) {
                cv::Rect b1 = cv::boundingRect(initial_light_bars[i].box);
                cv::Rect b2 = cv::boundingRect(initial_light_bars[j].box);
                
                double iou = calculate_iou(b1, b2);
                if (iou > LIGHT_BAR_IOU_THRESHOLD) {
                    nearby_indices.push_back(static_cast<int>(j));
                    visited[j] = true;
                }
            }
        }

        // Merge nearby light bar vertices to generate new rectangle
        std::vector<cv::Point2f> all_points;
        for (int k : nearby_indices) {
            for (const auto& point : initial_light_bars[k].box) {
                all_points.push_back(point);
            }
        }
        
        if (!all_points.empty()) {
            cv::RotatedRect merged_rect = cv::minAreaRect(all_points);
            cv::Point2f vertices[4];
            merged_rect.points(vertices);
            
            std::vector<cv::Point2f> merged_box;
            for (int k = 0; k < 4; k++) {
                merged_box.push_back(vertices[k]);
            }
            
            processed_boxes.push_back(merged_box);
            visited[i] = true;
        }
    }

    return processed_boxes;
}

std::vector<LightBar> ArmorDetector::filter_valid_light_bars(const std::vector<std::vector<cv::Point2f>>& processed_boxes) {
    std::vector<LightBar> valid_light_bars;

    for (const auto& box : processed_boxes) {
        if (box.size() < 4) continue;
        
        cv::RotatedRect rect = cv::minAreaRect(box);
        cv::Point2f center = rect.center;
        float width = rect.size.width;
        float height = rect.size.height;
        
        float length = width > height ? width : height;
        float bar_width = width > height ? height : width;
        
        float main_angle = rect.angle;
        if (width > height) {
            // 保持原始角度不变
        } else {
            main_angle += 90;
        }
        
        // Normalize angle to [-90, 90]
        main_angle = std::fmod(main_angle, 180);
        if (main_angle > 90) main_angle -= 180;
        if (main_angle < -90) main_angle += 180;
        
        float main_angle_rad = main_angle * CV_PI / 180.0f;

        cv::Point2f end1(
            center.x + (length / 2) * std::cos(main_angle_rad),
            center.y + (length / 2) * std::sin(main_angle_rad)
        );
        cv::Point2f end2(
            center.x - (length / 2) * std::cos(main_angle_rad),
            center.y - (length / 2) * std::sin(main_angle_rad)
        );

        float rect_area = length * bar_width;
        float contour_area = cv::contourArea(box);
        float area_ratio = rect_area > 0 ? contour_area / rect_area : 0;
        
        // Filter: sufficient area + high fill ratio + non-horizontal
        if (rect_area > 40 && area_ratio > 0.4 && std::abs(main_angle_rad) > angle_threshold_rad) {
            LightBar light_bar;
            light_bar.center = center;
            light_bar.size = cv::Size2f(length, bar_width);
            light_bar.angle = main_angle;
            light_bar.angle_rad = main_angle_rad;
            light_bar.area = rect_area;
            light_bar.center_line = {end1, end2};
            light_bar.center_line_length = length;
            light_bar.box = box;
            
            valid_light_bars.push_back(light_bar);
        }
    }

    return valid_light_bars;
}

std::vector<ArmorPlate> ArmorDetector::pair_light_bars_to_armor(const std::vector<LightBar>& light_bars, const std::string& target_color) {
    std::vector<ArmorPlate> armor_plates;
    std::vector<bool> used(light_bars.size(), false);

    for (size_t i = 0; i < light_bars.size(); i++) {
        if (used[i]) continue;
        const LightBar& lb1 = light_bars[i];

        for (size_t j = i + 1; j < light_bars.size(); j++) {
            if (used[j]) continue;
            const LightBar& lb2 = light_bars[j];

            // Angle difference filtering
            float angle_diff = std::abs(lb1.angle_rad - lb2.angle_rad);
            float angle_diff_2 = 2.0f * CV_PI - angle_diff;
            angle_diff = std::min(angle_diff, angle_diff_2);
            if (angle_diff > ARMOR_ANGLE_DIFF_THRESHOLD) {
                continue;
            }

            // Distance ratio filtering
            float dx = lb2.center.x - lb1.center.x;
            float dy = lb2.center.y - lb1.center.y;
            float distance = std::sqrt(dx * dx + dy * dy);
            float avg_length = (lb1.center_line_length + lb2.center_line_length) / 2;
            float distance_ratio = avg_length > 0 ? distance / avg_length : 0;
            if (distance_ratio <= ARMOR_DISTANCE_RATIO_MIN || distance_ratio >= ARMOR_DISTANCE_RATIO_MAX) {
                continue;
            }

            // Length difference filtering
            float length_diff_ratio = avg_length > 0 ? 
                std::abs(lb1.center_line_length - lb2.center_line_length) / avg_length : 0;
            if (length_diff_ratio > ARMOR_LENGTH_DIFF_RATIO) {
                continue;
            }

            cv::Point2f armor_center(
                (lb1.center.x + lb2.center.x) / 2,
                (lb1.center.y + lb2.center.y) / 2
            );
            
            double confidence = (lb1.area + lb2.area) / (distance + 1);

            ArmorPlate armor_plate;
            armor_plate.color = target_color;
            armor_plate.center = armor_center;
            armor_plate.confidence = confidence;
            armor_plate.pair = std::make_pair(lb1, lb2);
            armor_plate.corners_2d = std::vector<cv::Point2f>(); // Will be computed later if needed
            
            armor_plates.push_back(armor_plate);

            used[i] = true;
            used[j] = true;
            break;
        }
    }

    // Sort by confidence in descending order
    std::sort(armor_plates.begin(), armor_plates.end(), 
              [](const ArmorPlate& a, const ArmorPlate& b) {
                  return a.confidence > b.confidence;
              });
    
    return armor_plates;
}

double ArmorDetector::calculate_iou(const cv::Rect& b1, const cv::Rect& b2) {
    cv::Point2f inter_tl(std::max(b1.x, b2.x), std::max(b1.y, b2.y));
    cv::Point2f inter_br(std::min(b1.x + b1.width, b2.x + b2.width), 
                         std::min(b1.y + b1.height, b2.y + b2.height));
    
    if (inter_br.x <= inter_tl.x || inter_br.y <= inter_tl.y) {
        return 0.0;
    }

    float inter_area = (inter_br.x - inter_tl.x) * (inter_br.y - inter_tl.y);
    float b1_area = b1.width * b1.height;
    float b2_area = b2.width * b2.height;
    float union_area = b1_area + b2_area - inter_area;
    
    return union_area > 0 ? inter_area / union_area : 0.0;
}