/* app */
#include "robot_def.h"
#include "robot_cmd.h"
/* module */
#include "controller.h"
#include "ins_task.h"
#include "message_center.h"
#include "general_def.h"
#include "kalman_filter.h"
#include "master_machine.h"
/* bsp */
#include "bsp_dwt.h"
#include "bsp_log.h"

/* ---- 模式 ---- */
static Vision_Mode_e current_mode = VISION_MODE_DEBUG;

/* ---- 视觉追踪 PID (外环) ---- */
static PIDInstance vision_yaw_pid;
static PIDInstance vision_pitch_pid;

/* 视觉测量值卡尔曼滤波器 */
static KalmanFilter_t x_measure_kf;
static KalmanFilter_t y_measure_kf;
static float filtered_x_measure;
static float filtered_y_measure;

/* 内环角度参考累计值 */
static float gimbal_yaw_ref;
static float gimbal_pitch_ref;

/* 手动模式角度速率 */
static float yaw_degree_velocity = 0;
static float pitch_degree_velocity = 0;

/* ---- pub-sub ---- */
static Publisher_t *gimbal_cmd_pub;
static Subscriber_t *gimbal_feed_sub;
static Publisher_t *motor_enable_pub;
static Publisher_t *motor_pid_pub;
static Publisher_t *feedforward_pub;
static Gimbal_Ctrl_Cmd_s gimbal_cmd_send;
static Gimbal_Upload_Data_s gimbal_fetch_data;

/* ---- 激光 ---- */
static void laser_on(void)
{
    HAL_GPIO_WritePin(POWER_5V_GPIO_Port, POWER_5V_Pin, GPIO_PIN_SET);
}

static void laser_off(void)
{
    HAL_GPIO_WritePin(POWER_5V_GPIO_Port, POWER_5V_Pin, GPIO_PIN_RESET);
}

/* ---- 视觉 PID 和卡尔曼滤波器重置 ---- */
static void VisionPIDReset(void)
{
    vision_yaw_pid.Iout      = 0;
    vision_yaw_pid.ITerm     = 0;
    vision_yaw_pid.Last_Err  = 0;
    vision_yaw_pid.Err       = 0;
    vision_yaw_pid.Pout      = 0;
    vision_yaw_pid.Dout      = 0;
    vision_yaw_pid.Output    = 0;
    vision_yaw_pid.Last_Dout = 0;

    vision_pitch_pid.Iout      = 0;
    vision_pitch_pid.ITerm     = 0;
    vision_pitch_pid.Last_Err  = 0;
    vision_pitch_pid.Err       = 0;
    vision_pitch_pid.Pout      = 0;
    vision_pitch_pid.Dout      = 0;
    vision_pitch_pid.Output    = 0;
    vision_pitch_pid.Last_Dout = 0;

    memset(x_measure_kf.xhat_data,   0, x_measure_kf.xhatSize   * sizeof(float));
    memset(x_measure_kf.P_data,      0, x_measure_kf.xhatSize   * x_measure_kf.xhatSize * sizeof(float));
    memset(y_measure_kf.xhat_data, 0, y_measure_kf.xhatSize * sizeof(float));
    memset(y_measure_kf.P_data,    0, y_measure_kf.xhatSize * y_measure_kf.xhatSize * sizeof(float));

    float P_diag[4] = {5.0f, 0.0f, 0.0f, 5.0f};
    memcpy(x_measure_kf.P_data,   P_diag, sizeof(P_diag));
    memcpy(y_measure_kf.P_data, P_diag, sizeof(P_diag));

    filtered_x_measure   = 0;
    filtered_y_measure = 0;
}

/* ---- 模式切换 ---- */
static void SwitchToAutoMode(void)
{
    if (current_mode == VISION_MODE_AUTO)
        return;

    gimbal_yaw_ref   = gimbal_fetch_data.gimbal_imu_data->YawTotalAngle;
    gimbal_pitch_ref = gimbal_fetch_data.gimbal_imu_data->Pitch;

    vision_data.auto_updated = 0;
    VisionPIDReset();

    vision_data.feedforward_updated = 0;

    /* 从调试模式切回时重新使能电机 */
    if (current_mode == VISION_MODE_DEBUG)
    {
        Motor_Enable_s msg = { .enable = 1 };
        PubPushMessage(motor_enable_pub, &msg);
    }

    current_mode = VISION_MODE_AUTO;
    LOGINFO("[CMD] switched to AUTO mode");
}

static void SwitchToManualMode(void)
{
    if (current_mode == VISION_MODE_MANUAL)
        return;

    // gimbal_yaw_ref   = gimbal_fetch_data.gimbal_imu_data->YawTotalAngle;
    // gimbal_pitch_ref = gimbal_fetch_data.gimbal_imu_data->Pitch;

    vision_data.manual_updated = 0;
    yaw_degree_velocity = 0;
    pitch_degree_velocity = 0;

    vision_data.feedforward_updated = 0;

    /* 从调试模式切回时重新使能电机 */
    if (current_mode == VISION_MODE_DEBUG)
    {
        Motor_Enable_s msg = { .enable = 1 };
        PubPushMessage(motor_enable_pub, &msg);
    }

    current_mode = VISION_MODE_MANUAL;
    LOGINFO("[CMD] switched to MANUAL mode");
}

static void SwitchToDebugMode(void)
{
    if (current_mode == VISION_MODE_DEBUG)
        return;

    vision_data.debug_updated = 0;

    current_mode = VISION_MODE_DEBUG;

    /* 失能 yaw 和 pitch 电机 */
    Motor_Enable_s msg = { .enable = 0 };
    PubPushMessage(motor_enable_pub, &msg);

    LOGINFO("[CMD] switched to DEBUG mode, motors disabled");
}

/* ---- 命令处理 ---- */
static void ProcessCommand(void)
{
    if (!vision_data.cmd_pending)
        return;

    vision_data.cmd_pending = 0;

    switch (vision_data.cmd)
    {
    case VISION_CMD_SWITCH_AUTO:
        SwitchToAutoMode();
        break;
    case VISION_CMD_SWITCH_MANUAL:
        SwitchToManualMode();
        break;
    case VISION_CMD_SWITCH_DEBUG:
        SwitchToDebugMode();
        break;
    case VISION_CMD_LASER_ON:
        laser_on();
        break;
    case VISION_CMD_LASER_OFF:
        laser_off();
        break;
    default:
        break;
    }
}

/* ---- KF + PID 初始化 ---- */
static void VisionKalmanFilterInit()
{
    Kalman_Filter_Init(&x_measure_kf, 2, 0, 1);
    static float P_init[4] = {5.0f, 0.0f, 0.0f, 5.0f};
    memcpy(x_measure_kf.P_data, P_init, sizeof(P_init));
    static float F_init[4] = {1.0f, 0.005f, 0.0f, 1.0f};
    memcpy(x_measure_kf.F_data, F_init, sizeof(F_init));
    static float Q_init[4] = {0.1f, 0.0f, 0.0f, 1.0f};
    memcpy(x_measure_kf.Q_data, Q_init, sizeof(Q_init));
    static uint8_t measurement_map[1] = {1};
    static float measurement_degree[1] = {1.0f};
    static float mat_R_diagonal[1] = {1.0f};
    static float state_min_variance[2] = {0.001f, 0.001f};
    x_measure_kf.UseAutoAdjustment = 1;
    memcpy(x_measure_kf.MeasurementMap,        measurement_map,       sizeof(measurement_map));
    memcpy(x_measure_kf.MeasurementDegree,     measurement_degree,    sizeof(measurement_degree));
    memcpy(x_measure_kf.MatR_DiagonalElements, mat_R_diagonal,        sizeof(mat_R_diagonal));
    memcpy(x_measure_kf.StateMinVariance,      state_min_variance,    sizeof(state_min_variance));

    Kalman_Filter_Init(&y_measure_kf, 2, 0, 1);
    memcpy(y_measure_kf.P_data, P_init, sizeof(P_init));
    memcpy(y_measure_kf.F_data, F_init, sizeof(F_init));
    memcpy(y_measure_kf.Q_data, Q_init, sizeof(Q_init));
    y_measure_kf.UseAutoAdjustment = 1;
    memcpy(y_measure_kf.MeasurementMap,        measurement_map,       sizeof(measurement_map));
    memcpy(y_measure_kf.MeasurementDegree,     measurement_degree,    sizeof(measurement_degree));
    memcpy(y_measure_kf.MatR_DiagonalElements, mat_R_diagonal,        sizeof(mat_R_diagonal));
    memcpy(y_measure_kf.StateMinVariance,      state_min_variance,    sizeof(state_min_variance));
}

void RobotCMDInit()
{
    gimbal_cmd_pub  = PubRegister("gimbal_cmd", sizeof(Gimbal_Ctrl_Cmd_s));
    gimbal_feed_sub = SubRegister("gimbal_feed", sizeof(Gimbal_Upload_Data_s));
    motor_enable_pub = PubRegister("motor_enable", sizeof(Motor_Enable_s));
    motor_pid_pub    = PubRegister("motor_pid", sizeof(Motor_PID_Tune_s));
    feedforward_pub  = PubRegister("gimbal_feedforward", sizeof(Gimbal_Feedforward_s));

    static PID_Init_Config_s yaw_vpid_config = {
        .Kp = 100.0f, .Ki = 0.0f, .Kd = 0.0f,
        .IntegralLimit = 100.0f,
        .CoefA = 0.5f, .CoefB = 0.1f,
        .Improve = PID_Trapezoid_Intergral | PID_Integral_Limit |
                   PID_Derivative_On_Measurement | PID_ChangingIntegrationRate,
        .MaxOut = 200.0f,
    };
    static PID_Init_Config_s pitch_vpid_config = {
        .Kp = 100.0f, .Ki = 0.0f, .Kd = 0.0f,
        .IntegralLimit = 100.0f,
        .CoefA = 0.5f, .CoefB = 0.1f,
        .Improve = PID_Trapezoid_Intergral | PID_Integral_Limit |
                   PID_Derivative_On_Measurement | PID_ChangingIntegrationRate,
        .MaxOut = 200.0f,
    };
    PIDInit(&vision_yaw_pid, &yaw_vpid_config);
    PIDInit(&vision_pitch_pid, &pitch_vpid_config);
    VisionKalmanFilterInit();

    Vision_Uart_Init(&huart1);
    laser_off();
}

/* ---- 自动 / 视觉追踪模式 ---- */
static void VisionTracking()
{
    if (vision_data.auto_updated)
    {
        vision_data.auto_updated = 0;

        x_measure_kf.MeasuredVector[0] = vision_data.measure_x;
        Kalman_Filter_Update(&x_measure_kf);
        filtered_x_measure = x_measure_kf.FilteredValue[0];

        y_measure_kf.MeasuredVector[0] = vision_data.measure_y;
        Kalman_Filter_Update(&y_measure_kf);
        filtered_y_measure = y_measure_kf.FilteredValue[0];
        float ref_x = vision_data.ref_x;
        float ref_y = vision_data.ref_y;
        #if VISION_PID_OUTPUT_YAW_REVERSE == 1
        float yaw_degree_velocity   = -PIDCalculate(&vision_yaw_pid, filtered_x_measure, ref_x);
        #else
        float yaw_degree_velocity   = PIDCalculate(&vision_yaw_pid, filtered_x_measure, ref_x);
        #endif
        #if VISION_PID_OUTPUT_PITCH_REVERSE == 1
        float pitch_degree_velocity = -PIDCalculate(&vision_pitch_pid, filtered_y_measure, ref_y);
        #else
        float pitch_degree_velocity = PIDCalculate(&vision_pitch_pid, filtered_y_measure, ref_y);
        #endif

        gimbal_yaw_ref   += 0.001f * ROBOT_CMD_DELAY * yaw_degree_velocity;
        gimbal_pitch_ref += 0.001f * ROBOT_CMD_DELAY * pitch_degree_velocity;
    }

    if (gimbal_pitch_ref > 30.0f)
        gimbal_pitch_ref = 30.0f;
    else if (gimbal_pitch_ref < -30.0f)
        gimbal_pitch_ref = -30.0f;

    gimbal_cmd_send.yaw   = gimbal_yaw_ref;
    gimbal_cmd_send.pitch = gimbal_pitch_ref;

    if (vision_data.feedforward_updated)
    {
        vision_data.feedforward_updated = 0;
        Gimbal_Feedforward_s ff = { .feedforward = vision_data.feedforward };
        PubPushMessage(feedforward_pub, &ff);
    }
}

/* ---- 手动模式: 上位机角度参考 ---- */
static void ManualMode()
{
    if (vision_data.manual_updated)
    {
        vision_data.manual_updated = 0;
        yaw_degree_velocity   = vision_data.yaw_degree_velocity;
        pitch_degree_velocity = vision_data.pitch_degree_velocity;
    }
    gimbal_yaw_ref   += 0.001f * ROBOT_CMD_DELAY * yaw_degree_velocity;
    gimbal_pitch_ref += 0.001f * ROBOT_CMD_DELAY * pitch_degree_velocity;

    if (gimbal_pitch_ref > 30.0f)
        gimbal_pitch_ref = 30.0f;
    else if (gimbal_pitch_ref < -30.0f)
        gimbal_pitch_ref = -30.0f;

    gimbal_cmd_send.yaw   = gimbal_yaw_ref;
    gimbal_cmd_send.pitch = gimbal_pitch_ref;

    if (vision_data.feedforward_updated)
    {
        vision_data.feedforward_updated = 0;
        Gimbal_Feedforward_s ff = { .feedforward = vision_data.feedforward };
        PubPushMessage(feedforward_pub, &ff);
    }
}

/* ---- 调试模式: 上位机下发 PID 参数调参 ---- */
static void DebugMode()
{
    if (vision_data.debug_updated)
    {
        vision_data.debug_updated = 0;

        /* 直接设置视觉环 PID 参数 */
        vision_yaw_pid.Kp = vision_data.vision_yaw_Kp;
        vision_yaw_pid.Ki = vision_data.vision_yaw_Ki;
        vision_yaw_pid.Kd = vision_data.vision_yaw_Kd;

        vision_pitch_pid.Kp = vision_data.vision_pitch_Kp;
        vision_pitch_pid.Ki = vision_data.vision_pitch_Ki;
        vision_pitch_pid.Kd = vision_data.vision_pitch_Kd;

        /* 发布电机角度环 PID 参数到 gimbal */
        Motor_PID_Tune_s tune = {
            .yaw_Kp   = vision_data.motor_yaw_Kp,
            .yaw_Ki   = vision_data.motor_yaw_Ki,
            .yaw_Kd   = vision_data.motor_yaw_Kd,
            .pitch_Kp = vision_data.motor_pitch_Kp,
            .pitch_Ki = vision_data.motor_pitch_Ki,
            .pitch_Kd = vision_data.motor_pitch_Kd,
        };
        PubPushMessage(motor_pid_pub, &tune);

        LOGINFO("[CMD] PID tuned: v_yaw(%.2f,%.2f,%.2f) v_pitch(%.2f,%.2f,%.2f)",
                vision_data.vision_yaw_Kp, vision_data.vision_yaw_Ki, vision_data.vision_yaw_Kd,
                vision_data.vision_pitch_Kp, vision_data.vision_pitch_Ki, vision_data.vision_pitch_Kd);
    }
}

void RobotCMDTask()
{
    SubGetMessage(gimbal_feed_sub, &gimbal_fetch_data);

    ProcessCommand();

    if (current_mode == VISION_MODE_AUTO)
        VisionTracking();
    else if (current_mode == VISION_MODE_DEBUG)
        DebugMode();
    else
        ManualMode();

    PubPushMessage(gimbal_cmd_pub, (void *)&gimbal_cmd_send);
}