#include "track_ir_algo.h"

#define TRACK_IR_P_GAIN 0.45f
#define TRACK_IR_MAX_YAW_CORR 18.0f
#define TRACK_IR_LOST_TURN_RPM 10.0f
#define TRACK_IR_LOST_SPEED_SCALE 35U

static int16_t s_last_deviation = 0;

static int16_t track_ir_calc_deviation(uint8_t line_mask)
{
    int16_t weighted_sum = 0;
    uint8_t count = 0;

    if ((line_mask & TRACK_IR_H4_BIT) != 0U) {
        weighted_sum += -30;
        count++;
    }
    if ((line_mask & TRACK_IR_H3_BIT) != 0U) {
        weighted_sum += -10;
        count++;
    }
    if ((line_mask & TRACK_IR_H2_BIT) != 0U) {
        weighted_sum += 10;
        count++;
    }
    if ((line_mask & TRACK_IR_H1_BIT) != 0U) {
        weighted_sum += 30;
        count++;
    }

    if (count == 0U) {
        return 0;
    }

    return (int16_t)(weighted_sum / (int16_t)count);
}

void track_ir_algo_init(void)
{
    s_last_deviation = 0;
}

void track_ir_algo_reset(void)
{
    s_last_deviation = 0;
}

void track_ir_algo_step(uint8_t line_mask, track_ir_ctrl_output_t *out)
{
    track_ir_state_t state;
    int16_t deviation = 0;
    float yaw = 0.0f;
    uint8_t speed_scale = 100U;

    if (out == NULL) {
        return;
    }

    if (line_mask == (TRACK_IR_H1_BIT | TRACK_IR_H2_BIT | TRACK_IR_H3_BIT | TRACK_IR_H4_BIT)) {
        state = TRACK_IR_STATE_LOST;
    } else if (line_mask == 0U) {
        state = TRACK_IR_STATE_CENTER;
    } else {
        uint8_t left_mask = line_mask & (TRACK_IR_H4_BIT | TRACK_IR_H3_BIT);
        uint8_t right_mask = line_mask & (TRACK_IR_H2_BIT | TRACK_IR_H1_BIT);

        if ((left_mask != 0U) && (right_mask == 0U)) {
            state = TRACK_IR_STATE_LEFT;
        } else if ((right_mask != 0U) && (left_mask == 0U)) {
            state = TRACK_IR_STATE_RIGHT;
        } else {
            state = TRACK_IR_STATE_CROSS;
        }
    }

    deviation = track_ir_calc_deviation(line_mask);

    if (state == TRACK_IR_STATE_LOST) {
        speed_scale = TRACK_IR_LOST_SPEED_SCALE;
        if (s_last_deviation < 0) {
            yaw = TRACK_IR_LOST_TURN_RPM;
        } else if (s_last_deviation > 0) {
            yaw = -TRACK_IR_LOST_TURN_RPM;
        } else {
            yaw = 0.0f;
        }
    } else if (state == TRACK_IR_STATE_CROSS) {
        yaw = 0.0f; // 十字路口默认直行
    } else {
        /* 车体旋转正方向与传感器偏差方向相反，因此取负号做闭环修正。 */
        yaw = -TRACK_IR_P_GAIN * (float)deviation;
        if (yaw > TRACK_IR_MAX_YAW_CORR) {
            yaw = TRACK_IR_MAX_YAW_CORR;
        }
        if (yaw < -TRACK_IR_MAX_YAW_CORR) {
            yaw = -TRACK_IR_MAX_YAW_CORR;
        }
    }

    if (deviation != 0) {
        s_last_deviation = deviation;
    }

    out->yaw_correction_rpm = yaw;
    out->speed_scale_percent = speed_scale;
    out->deviation = deviation;
    out->state = state;
}