Heat Seeker Robot

#include <M5Unified.h> 
#include <M5_Thermal2.h>
#include <Wire.h> 
#include "Hat_8Servos.h" 

// --- DEVICE CONFIG ---
#define THERMAL_SDA_PIN 32
#define THERMAL_SCL_PIN 33
#define THERMAL_I2C_SPEED 400000UL
#define THERMAL_I2C_ADDR 0x32
#define SERVOS_SDA_PIN 0
#define SERVOS_SCL_PIN 26
#define SERVOS_I2C_SPEED 100000UL 
#define SERVOS_DEVICE_ADDR 0x36

M5_Thermal2 thermal2;
Hat_8Servos drive; 

// --- SERVO/MOTOR CONFIG ---
uint8_t currentServoStates[8] = {90, 90, 90, 90, 90, 90, 90, 90}; 

bool isForwardMode = false;
bool moveForward = false;
bool targetAquired = false;
bool moveForwardAlertSwitch = false;
int moveForwardAlertCounter = 0;
int moveForwardAlertCounterLimit = 10;

// --- THERMAL & DISPLAY CONFIG ---
#define FIXED_MIN_TEMP 20.0 
#define FIXED_MAX_TEMP 40.0 
#define DISPLAY_SCALE 5
#define THRESHOLD_TEMP 33.0 
#define THERMAL_CENTER_X 16 
#define BASE_SPEED 95      
#define MAX_SPEED_ADJUST 40 
#define TFT_MAGENTA M5.Display.color565(255, 0, 255)

// --- PID CONSTANTS (Requires Tuning!) ---
float Kp = 2.5;  
float Ki = 0.001; 
float Kd = 0.5;   

// --- PID STATE VARIABLES ---
long lastTime = 0;
float error = 0;
float lastError = 0;
float integral = 0;

// =========================================================================
// --- HELPER FUNCTIONS ---
// =========================================================================

void setThermalBus() {
    Wire.end(); 
    Wire.begin(THERMAL_SDA_PIN, THERMAL_SCL_PIN, THERMAL_I2C_SPEED);
    delay(25);
}

void setServoBus() {
    Wire.end();
    Wire.begin(SERVOS_SDA_PIN, SERVOS_SCL_PIN, SERVOS_I2C_SPEED);
    delay(25);
}

void handleLEDAndBuzzer(){
    if(isForwardMode)
    {   
        if(targetAquired){
            if(moveForward){
                thermal2.setLed(64, 0, 0);
                if(!moveForwardAlertSwitch){
                    thermal2.buzzerOn();               
                    moveForwardAlertSwitch = true;
                }    
            }
            else{
                thermal2.setLed(4, 0, 0);
            }
        }
        else{
            thermal2.setLed(0, 4, 0);
        }
    }
    else{
        if(targetAquired){
            thermal2.setLed(4, 0, 0);
        }
        else{
            thermal2.setLed(0, 0, 4);
        }
    }

    if(moveForwardAlertSwitch){
        moveForwardAlertCounter += 1;
    
        if (!(moveForwardAlertCounter > 0 && moveForwardAlertCounter <= moveForwardAlertCounterLimit)){
            thermal2.buzzerOff();
            moveForwardAlertCounter = 0;
            moveForwardAlertSwitch = false;
        }
    }
}

uint16_t getColor(float temp, float min_val, float max_val) {
    temp = constrain(temp, min_val, max_val);
    float normalizedTemp = (max_val - min_val) > 0 ? (temp - min_val) / (max_val - min_val) : 0;
    byte r, g, b;

    if (normalizedTemp < 0.25) {
        r = 0; g = (byte)(255 * (normalizedTemp / 0.25)); b = 255;
    } else if (normalizedTemp < 0.5) {
        r = 0; g = 255; b = (byte)(255 * (1 - (normalizedTemp - 0.25) / 0.25));
    } else if (normalizedTemp < 0.75) {
        r = (byte)(255 * ((normalizedTemp - 0.5) / 0.25)); g = 255; b = 0;
    } else {
        r = 255; g = (byte)(255 * (1 - (normalizedTemp - 0.75) / 0.25)); b = 0;
    }
    return M5.Display.color565(r, g, b);
}

float computePID(int current_centroid_x) {
    long now = millis();
    float timeChange = (float)(now - lastTime);
    error = (float)current_centroid_x - (float)THERMAL_CENTER_X;
    float proportional = Kp * error;
    integral += (error * timeChange);
    integral = constrain(integral, -2000, 2000); 
    float integral_term = Ki * integral;
    float derivative = (error - lastError) / timeChange;
    float derivative_term = Kd * derivative;
    float output = proportional + integral_term + derivative_term;
    lastError = error;
    lastTime = now;
    return constrain(output, -MAX_SPEED_ADJUST, MAX_SPEED_ADJUST);
}


// =========================================================================
// --- SETUP ---
// =========================================================================

void setup() {
    Serial.begin(115200);
    delay(500); 
    M5.begin();

    // Display Setup
    if (M5.Display.width() < M5.Display.height()) {
        M5.Display.setRotation(M5.Display.getRotation() ^ 1);
    }
    M5.Ex_I2C.begin(); // Not strictly needed, but included for completeness

    // 1. Thermal Init
    Serial.println("-> Initializing Thermal Unit...");
    setThermalBus(); 
    if (!thermal2.begin(&Wire, THERMAL_I2C_ADDR)) {
         Serial.println("❌ Thermal2 init failed!");
    }
    thermal2.setRefreshRate(6);
    thermal2.setNoiseFilterLevel(8);
    thermal2.alarmOff(); // Keep control logic simple

    thermal2.ledOn();
    thermal2.buzzerOff();

    thermal2.setBuzzer(300, 128);
    thermal2.setLed(0, 0, 4);

    // 2. Servo Init
    Serial.println("-> Initializing Servo Hat...");
    drive.begin(&Wire, SERVOS_SDA_PIN, SERVOS_SCL_PIN, SERVOS_DEVICE_ADDR);
    setServoBus(); 
    drive.enableServoPower(1); 
    drive.setServoAnglesFromArray(currentServoStates); 

    Serial.println("*** PID Hotspot Follower Ready ***");
    lastTime = millis();
}

// =========================================================================
// --- MAIN LOOP ---
// =========================================================================

void loop() {
    M5.update();
    M5.Display.startWrite(); // Begin fast display updates

    int centroid_x_px = -1; // Horizontal centroid (PID Input)
    int centroid_y_px = -1; // Vertical centroid (For display only)
    static int frame_count = 0;

    // 1. --- Switch Bus and Read Thermal Data (Input) ---
    setThermalBus(); 
    
    if (thermal2.btnWasClicked()) {
        isForwardMode = !isForwardMode;
    }

    handleLEDAndBuzzer();

    if (thermal2.update()) {
        frame_count++;
        auto temp_data = thermal2.getTemperatureData();
        targetAquired = false;
        
        // **Centroid Calculation Variables**
        int hot_pixels_count = 0;
        long sum_x_pixels = 0;
        long sum_y_pixels = 0; // Needed for display

        // **Thermal Visualization and Centroid Calculation Loop**
        for (int idx = 0; idx < 384; ++idx) {
            int y = idx >> 4;
            int x = ((idx & 15) << 1) + ((y & 1) != temp_data.getSubPage());
            float pixel_temp = temp_data.getPixelTemperature(idx);
            
            // Centroid Logic: Check Threshold and Accumulate
            if (pixel_temp > THRESHOLD_TEMP) {
                sum_x_pixels += x;
                sum_y_pixels += y;
                hot_pixels_count++;
            }

            // Draw the thermal image pixel
            uint16_t color = getColor(pixel_temp, FIXED_MIN_TEMP, FIXED_MAX_TEMP);
            M5.Display.fillRect(70 + x * DISPLAY_SCALE, 8 + y * DISPLAY_SCALE, DISPLAY_SCALE, DISPLAY_SCALE, color);
        }

        // Final Centroid Calculation
        if (hot_pixels_count > 0) {
            centroid_x_px = sum_x_pixels / hot_pixels_count;
            centroid_y_px = sum_y_pixels / hot_pixels_count; 
            targetAquired = true;
        }

        // --- DISPLAY Centroid and Stats (Your original code) ---      
        M5.Display.setCursor(0, 100);
        if (hot_pixels_count > 0) {
            // Draw tracking dot
            int draw_x = 70 + centroid_x_px * DISPLAY_SCALE + (DISPLAY_SCALE / 2);
            int draw_y = 8 + centroid_y_px * DISPLAY_SCALE + (DISPLAY_SCALE / 2);
            M5.Display.fillCircle(draw_x, draw_y, 2, TFT_MAGENTA); 
            M5.Display.printf("Cen(X,Y): %02d,%02d\n", centroid_x_px, centroid_y_px);
        } else {
            M5.Display.printf("Cen(X,Y): --,--\n");
        }

        // Display text stats
        static int prev_sec = ~0u;
        int sec = millis() / 1000;
        if (prev_sec != sec) {
            prev_sec = sec;
            M5.Display.setCursor(0, 8);
            M5.Display.printf("avg: %5.2f\n", temp_data.getAverageTemperature());
            M5.Display.printf("med: %5.2f\n", temp_data.getMedianTemperature());
            M5.Display.printf("high:%5.2f\n", temp_data.getHighestTemperature());
            M5.Display.printf("low: %5.2f\n", temp_data.getLowestTemperature());
            M5.Display.printf("fps: %02d\n", frame_count);
            M5.Display.printf("TH:%5.1fC\n", THRESHOLD_TEMP); 
            frame_count = 0;
        }
    } 
    
    // --- PID Control (Using calculated centroid_x_px) ---
    int right_motor_command = 90;
    int left_motor_command = 90;

    if (centroid_x_px != -1) {
        float pid_adjustment = computePID(centroid_x_px);
        right_motor_command = constrain((int)(BASE_SPEED + pid_adjustment), 0, 180);
        left_motor_command = constrain((int)(BASE_SPEED - pid_adjustment), 0, 180);
    } 

    if (right_motor_command < 100 && right_motor_command > 80 && left_motor_command < 100 && left_motor_command > 80){
        moveForward = true;
    }
    else{
        moveForward = false;
    }

    // 2. --- Switch Bus and Apply Array Motor Commands (Output) ---
    if (moveForward){
        // Serial.println("Dead Zone");
        if(isForwardMode && targetAquired > 0){
            currentServoStates[0] = 135;
            currentServoStates[4] = 45;
        }
        else{
            currentServoStates[0] = 90;
            currentServoStates[4] = 90;
        }
    }
    else{
        // Serial.println("Tracking");
        currentServoStates[0] = 180 - (uint8_t)right_motor_command;
        currentServoStates[4] = (uint8_t)left_motor_command;
    }

    setServoBus(); 
    drive.setServoAnglesFromArray(currentServoStates);
    
    // 3. --- Final Display Update and Loop Control ---
    M5.Display.setCursor(0, 0);
    M5.Display.printf("btn: %02x\n", thermal2.getButtonRawState());
    M5.Display.endWrite(); // Finish fast display updates

    // Optional: Add PID status to display for tuning
    M5.Display.setCursor(0, 140);
    M5.Display.printf("Kp:%.1f E:%.1f A:%.1f", Kp, error, (float)right_motor_command - 90.0);

    // Minor delay to stabilize the loop frequency
    delay(10); 
}