#include "Visualizer.h"
#include <sstream>

Visualizer::Visualizer() {
    // Constructor implementation
}

Visualizer::~Visualizer() {
    // Destructor implementation
}

cv::Mat& Visualizer::draw_light_bars(cv::Mat& frame, const std::vector<LightBar>& light_bars, const std::string& target_color) {
    cv::Scalar color = (target_color == "red") ? cv::Scalar(0, 0, 255) : cv::Scalar(255, 0, 0); // Red or Blue in BGR
    
    for (const auto& light_bar : light_bars) {
        // Draw rotated rectangle for each light bar
        for (size_t i = 0; i < light_bar.box.size(); i++) {
            cv::Point2f pt1 = light_bar.box[i];
            cv::Point2f pt2 = light_bar.box[(i + 1) % light_bar.box.size()];
            cv::line(frame, pt1, pt2, color, 2);
        }
    }
    
    return frame;
}

cv::Mat& Visualizer::draw_armor_plate(cv::Mat& frame, const ArmorPlate& armor_plate) {
    cv::circle(frame, armor_plate.center, 10, cv::Scalar(0, 255, 0), 2); // Green circle around armor center
    
    // Draw confidence text
    std::stringstream ss;
    ss << "Conf: " << std::fixed << std::setprecision(2) << armor_plate.confidence;
    cv::putText(frame, ss.str(), cv::Point(armor_plate.center.x + 15, armor_plate.center.y - 15),
                cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 255, 0), 1);
    
    return frame;
}

cv::Mat& Visualizer::draw_offset_text(cv::Mat& frame, const cv::Point2f* display_center, const std::string& /*target_color*/, bool is_predicted) {
    if (display_center == nullptr) {
        return frame;
    }
    
    cv::Point2f img_center(frame.cols / 2.0f, frame.rows / 2.0f);
    int offset_x = static_cast<int>(display_center->x - img_center.x);
    int offset_y = static_cast<int>(display_center->y - img_center.y);
    
    std::stringstream ss;
    ss << "Offset: (" << offset_x << ", " << offset_y << ")";
    
    // Use different color based on whether it's predicted or actual
    cv::Scalar color = is_predicted ? cv::Scalar(0, 255, 255) : cv::Scalar(0, 255, 0); // Yellow for predicted, Green for actual
    cv::putText(frame, ss.str(), cv::Point(20, 30), cv::FONT_HERSHEY_SIMPLEX, 0.7, color, 2);
    
    return frame;
}

cv::Mat& Visualizer::draw_crosshair(cv::Mat& frame, const cv::Point2f& center, const cv::Scalar& color, int size) {
    // Draw horizontal line
    cv::line(frame, 
             cv::Point(static_cast<int>(center.x - size), static_cast<int>(center.y)), 
             cv::Point(static_cast<int>(center.x + size), static_cast<int>(center.y)), 
             color, 2);
    
    // Draw vertical line
    cv::line(frame, 
             cv::Point(static_cast<int>(center.x), static_cast<int>(center.y - size)), 
             cv::Point(static_cast<int>(center.x), static_cast<int>(center.y + size)), 
             color, 2);
    
    return frame;
}

cv::Mat& Visualizer::draw_ballistic_point(cv::Mat& frame, const cv::Point2f* ballistic_point) {
    if (ballistic_point != nullptr) {
        // Draw a red crosshair at the ballistic prediction point
        cv::circle(frame, *ballistic_point, 8, cv::Scalar(0, 0, 255), 2); // Red circle
        cv::line(frame, 
                 cv::Point(static_cast<int>(ballistic_point->x - 10), static_cast<int>(ballistic_point->y)), 
                 cv::Point(static_cast<int>(ballistic_point->x + 10), static_cast<int>(ballistic_point->y)), 
                 cv::Scalar(0, 0, 255), 2); // Red horizontal line
        cv::line(frame, 
                 cv::Point(static_cast<int>(ballistic_point->x), static_cast<int>(ballistic_point->y - 10)), 
                 cv::Point(static_cast<int>(ballistic_point->x), static_cast<int>(ballistic_point->y + 10)), 
                 cv::Scalar(0, 0, 255), 2); // Red vertical line
        
        // Label as ballistic point
        cv::putText(frame, "BP", cv::Point(static_cast<int>(ballistic_point->x + 15), static_cast<int>(ballistic_point->y - 15)),
                    cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 255), 1);
    }
    
    return frame;
}