#include <iostream>
#include <string>
#include <vector>
#include <chrono>
#include <thread>
#include <memory>

#include <opencv2/opencv.hpp>

#include "config.h"
#include "MindVisionCamera.h"  // 使用MindVision相机
#include "ImagePreprocessor.h"
#include "ArmorDetector.h"
#include "KalmanFilter.h"
#include "Visualizer.h"
#include "BallisticPredictor.h"
#include "TTLCommunicator.h"

// Function to output control data to TTL device (with enable control)
void output_control_data(const cv::Point2f* ballistic_point,
                        const std::string& target_color,
                        int frame_counter,
                        TTLCommunicator* ttl_communicator,
                        const cv::Point2f& img_center,
                        bool use_ttl,
                        bool armor_detected) {
    // Only send data when TTL is enabled, meets frame interval, AND armor is detected
    if (use_ttl && frame_counter % 5 == 0 && armor_detected) {
        // Only send if we have a valid ballistic point to send
        if (ballistic_point != nullptr) {
            // Calculate offset (based on actual image center)
            int ballistic_offset_x = static_cast<int>(ballistic_point->x - img_center.x);
            int ballistic_offset_y = static_cast<int>(ballistic_point->y - img_center.y);

            // Color simplification mapping
            std::string simplified_color = target_color;
            if (target_color == "red") simplified_color = "r";
            else if (target_color == "blue") simplified_color = "b";

            // Construct send string
            std::string send_str = "s " + simplified_color + " " + std::to_string(ballistic_offset_x) + " " + std::to_string(ballistic_offset_y);

            // Send data
            if (ttl_communicator != nullptr) {
                ttl_communicator->send_data(send_str);
            }
        }
        // Note: When ballistic_point is nullptr, no TTL data is sent (fulfilling the requirement)
    }
}

void set_camera_resolution(MindVisionCamera, int width, int height) {
    // The resolution is set during camera initialization in MindVision
    // We already set it in MindVisionCamera constructor, so this is for logging only
    std::cout << "Camera resolution already set to: " << width << "x" << height << " during initialization" << std::endl;
}

int main(int /*argc*/, char* /*argv*/[]) {
    std::string target_color = "red";
    int cam_id = 0;
    cv::Size default_resolution(640, 480); // This will be set by MindVisionCamera constructor
    bool use_ttl = true; // Set to false to disable TTL communication

    // Define optional resolution list (adjust based on camera support)
    std::vector<cv::Size> resolutions = {
        cv::Size(320, 240),   // Low resolution, high frame rate
        cv::Size(640, 480),   // Standard resolution
        cv::Size(1280, 720),  // HD resolution
        cv::Size(1920, 1080)  // Full HD resolution
    };

    // Find the index of default resolution
    int res_index = 1; // Default to index 1 (640x480)
    for (size_t i = 0; i < resolutions.size(); i++) {
        if (resolutions[i] == default_resolution) {
            res_index = static_cast<int>(i);
            break;
        }
    }

    // Initialize TTL communication (only when enabled)
    TTLCommunicator* ttl = nullptr;
    if (use_ttl) {
        ttl = new TTLCommunicator(TTL_BAUDRATE);
        if (!ttl->connect()) {
            std::cout << "Warning: Cannot establish TTL connection, will continue running but not send data" << std::endl;
        }
    } else {
        std::cout << "TTL communication disabled" << std::endl;
    }

    // Initialize visual processing modules with MindVision camera
    MindVisionCamera camera(cam_id, target_color);
    if (!camera.is_opened) {
        std::cerr << "Cannot open MindVision camera!" << std::endl;
        return -1;
    }

    // Initialize ballistic predictor (adjustable bullet speed, e.g., 16m/s)
    BallisticPredictor ballistic_predictor(16.0f);
    // Record target speed (obtained from Kalman filter)
    cv::Point2f target_speed(0.0f, 0.0f);

    // Set initial resolution
    set_camera_resolution(camera, resolutions[res_index].width, resolutions[res_index].height);
    std::cout << "Initial camera resolution: " << resolutions[res_index].width << "x" << resolutions[res_index].height << std::endl;

    ImagePreprocessor preprocessor;
    ArmorDetector detector;
    KalmanFilter kalman_tracker;
    Visualizer visualizer;

    int frame_counter = 0; // Counter to control output frequency
    int max_consecutive_predicts = 20; // Maximum consecutive prediction times
    int consecutive_predicts = 0; // Current consecutive prediction count

    cv::Mat frame;
    try {
        while (true) {
            if (!camera.read_frame(frame)) {
                std::cout << "Cannot read from MindVision camera, exiting!" << std::endl;
                break;
            }

            // Get actual image size and calculate center
            cv::Point2f img_center(frame.cols / 2.0f, frame.rows / 2.0f);

            // Preprocessing
            cv::Mat mask, color_only_frame;
            preprocessor.process_frame(frame, target_color, mask, color_only_frame);

            // Armor plate detection
            std::vector<LightBar> valid_light_bars;
            std::vector<ArmorPlate> armor_plates;
            detector.detect(mask, target_color, valid_light_bars, armor_plates);

            // Kalman filter tracking (modified part: get target speed)
            cv::Point2f* current_center = nullptr;
            if (!armor_plates.empty()) {
                current_center = &(armor_plates[0].center);
            }

            cv::Point2f* predicted_center = nullptr;

            if (current_center != nullptr) {
                // Has detection result: update Kalman filter, reset consecutive prediction count
                kalman_tracker.update(*current_center);
                predicted_center = new cv::Point2f(kalman_tracker.predict());

                // Get velocity from the Kalman filter
                target_speed = cv::Point2f(0.0f, 0.0f); // Kalman filter in OpenCV doesn't directly expose velocity
                consecutive_predicts = 0;
            } else {
                // No detection result: only use Kalman prediction, limit consecutive predictions
                consecutive_predicts++;
                if (consecutive_predicts < max_consecutive_predicts) {
                    cv::Point2f temp_pred = kalman_tracker.predict();
                    if (temp_pred.x != 0 || temp_pred.y != 0) { // Check if prediction is valid
                        predicted_center = new cv::Point2f(temp_pred);
                    }
                } else {
                    predicted_center = nullptr;
                    target_speed = cv::Point2f(0.0f, 0.0f);
                }
            }

            // Determine display center
            cv::Point2f* display_center = current_center;
            if (display_center == nullptr && predicted_center != nullptr) {
                display_center = predicted_center;
            }
            bool is_predicted = (display_center != nullptr) && (current_center == nullptr);

            // Calculate ballistic prediction point
            cv::Point2f* ballistic_point = ballistic_predictor.predict_ballistic_point(
                predicted_center, img_center, target_speed
            );

            auto current_time = std::chrono::high_resolution_clock::now();

            // Visualization
            visualizer.draw_light_bars(frame, valid_light_bars, target_color);
            if (!armor_plates.empty()) {
                visualizer.draw_armor_plate(frame, armor_plates[0]);
            }
            visualizer.draw_offset_text(frame, display_center, target_color, is_predicted);
            visualizer.draw_ballistic_point(frame, ballistic_point);

            // Output control data to TTL (passing use_ttl to control whether to send)
            // Only send TTL data if actual armor plates were detected (not just predicted)
            bool armor_detected = !armor_plates.empty();
            // Only send TTL data when we have an actual detection, not just prediction
            bool should_send_ttl = armor_detected && (display_center != nullptr);
            output_control_data(display_center, target_color, frame_counter, ttl, img_center, use_ttl, should_send_ttl);

            // Display windows
            cv::imshow("Armor Detection", frame);
            cv::imshow(target_color + " Mask", mask);
            cv::imshow(target_color + " Only", color_only_frame);

            // Exit on 'q' key press
            if (cv::waitKey(1) == 'q') {
                break;
            }

            frame_counter++;

            // Control max frame rate (100 FPS)
            auto end_time = std::chrono::high_resolution_clock::now();
            auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - current_time).count();
            if (elapsed < 10) { // 100 FPS = 10ms per frame
                std::this_thread::sleep_for(std::chrono::milliseconds(30 - elapsed));
            }

            // Clean up dynamically allocated memory
            if (predicted_center != nullptr) {
                delete predicted_center;
                predicted_center = nullptr;
            }
        }
    }
    catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
    }

    // Cleanup
    camera.release();
    cv::destroyAllWindows();

    // Only close TTL connection when enabled and initialized
    if (use_ttl && ttl != nullptr) {
        ttl->close();
        delete ttl;
        ttl = nullptr;
    }

    return 0;
}