#include "motor.h"

#include <stdbool.h>
#include <stdint.h>

#include "tim.h"

static uint16_t s_last_count[4] = {0U, 0U, 0U, 0U};
static volatile float s_current_rpm[4] = {0.0f, 0.0f, 0.0f, 0.0f};
static volatile int32_t s_odom_acc_ticks[4] = {0, 0, 0, 0};

/* 速度估计滚动窗口：降低 5ms 单拍增量测速的量化噪声 */
static int32_t s_speed_acc_ticks[4] = {0, 0, 0, 0};
static uint8_t s_speed_window_count = 0U;
static float s_speed_window_dt = 0.0f;

static inline void Motor_WriteBridge(TIM_HandleTypeDef *htim,
                                     uint32_t channel_in1,
                                     uint32_t channel_in2,
                                     uint16_t pwm_in1,
                                     uint16_t pwm_in2);
static void Motor_PrimeEncoderCounters(void);
static uint16_t Motor_ClampAbsToPwm(int32_t value);
static int16_t Motor_SaturateToI16(int32_t value);

static inline void Motor_WriteBridge(TIM_HandleTypeDef *htim,
                                     uint32_t channel_in1,
                                     uint32_t channel_in2,
                                     uint16_t pwm_in1,
                                     uint16_t pwm_in2)
{
    __HAL_TIM_SET_COMPARE(htim, channel_in1, pwm_in1);
    __HAL_TIM_SET_COMPARE(htim, channel_in2, pwm_in2);
}

static uint16_t Motor_ClampAbsToPwm(int32_t value)
{
    if (value > PWM_LIMIT)
    {
        value = PWM_LIMIT;
    }
    else if (value < -PWM_LIMIT)
    {
        value = -PWM_LIMIT;
    }

    if (value < 0)
    {
        value = -value;
    }

    return (uint16_t)value;
}

static int16_t Motor_SaturateToI16(int32_t value)
{
    if (value > 32767)
    {
        return 32767;
    }
    if (value < -32768)
    {
        return -32768;
    }
    return (int16_t)value;
}

static void Motor_PrimeEncoderCounters(void)
{
    s_last_count[MOTOR_FL] = (uint16_t)__HAL_TIM_GET_COUNTER(&htim5);
    s_last_count[MOTOR_RL] = (uint16_t)__HAL_TIM_GET_COUNTER(&htim3);
    s_last_count[MOTOR_FR] = (uint16_t)__HAL_TIM_GET_COUNTER(&htim4);
    s_last_count[MOTOR_RR] = (uint16_t)__HAL_TIM_GET_COUNTER(&htim1);
}

void Motor_Init(void)
{
    uint32_t primask = __get_PRIMASK();

    if (HAL_TIM_PWM_Start(&htim2, TIM_CHANNEL_3) != HAL_OK) { Error_Handler(); }
    if (HAL_TIM_PWM_Start(&htim2, TIM_CHANNEL_4) != HAL_OK) { Error_Handler(); }
    if (HAL_TIM_PWM_Start(&htim9, TIM_CHANNEL_1) != HAL_OK) { Error_Handler(); }
    if (HAL_TIM_PWM_Start(&htim9, TIM_CHANNEL_2) != HAL_OK) { Error_Handler(); }
    if (HAL_TIM_PWM_Start(&htim8, TIM_CHANNEL_1) != HAL_OK) { Error_Handler(); }
    if (HAL_TIM_PWM_Start(&htim8, TIM_CHANNEL_2) != HAL_OK) { Error_Handler(); }
    if (HAL_TIM_PWM_Start(&htim8, TIM_CHANNEL_3) != HAL_OK) { Error_Handler(); }
    if (HAL_TIM_PWM_Start(&htim8, TIM_CHANNEL_4) != HAL_OK) { Error_Handler(); }

    if (HAL_TIM_Encoder_Start(&htim5, TIM_CHANNEL_ALL) != HAL_OK) { Error_Handler(); }
    if (HAL_TIM_Encoder_Start(&htim3, TIM_CHANNEL_ALL) != HAL_OK) { Error_Handler(); }
    if (HAL_TIM_Encoder_Start(&htim4, TIM_CHANNEL_ALL) != HAL_OK) { Error_Handler(); }
    if (HAL_TIM_Encoder_Start(&htim1, TIM_CHANNEL_ALL) != HAL_OK) { Error_Handler(); }

    Motor_Brake_All();
    Motor_PrimeEncoderCounters();

    __disable_irq();
    s_current_rpm[MOTOR_FL] = 0.0f;
    s_current_rpm[MOTOR_RL] = 0.0f;
    s_current_rpm[MOTOR_FR] = 0.0f;
    s_current_rpm[MOTOR_RR] = 0.0f;

    s_odom_acc_ticks[MOTOR_FL] = 0;
    s_odom_acc_ticks[MOTOR_RL] = 0;
    s_odom_acc_ticks[MOTOR_FR] = 0;
    s_odom_acc_ticks[MOTOR_RR] = 0;

    s_speed_acc_ticks[MOTOR_FL] = 0;
    s_speed_acc_ticks[MOTOR_RL] = 0;
    s_speed_acc_ticks[MOTOR_FR] = 0;
    s_speed_acc_ticks[MOTOR_RR] = 0;
    s_speed_window_count = 0U;
    s_speed_window_dt = 0.0f;
    __set_PRIMASK(primask);
}

void Set_Motor_Output(Motor_ID_t id, int16_t control_out)
{
    int32_t hw_command = (int32_t)control_out;
    bool forward;
    uint16_t pwm_val;

    if ((id == MOTOR_FL) || (id == MOTOR_RL))
    {
        hw_command = -hw_command;
    }

    forward = (hw_command < 0);
    pwm_val = Motor_ClampAbsToPwm(hw_command);

    switch (id)
    {
        case MOTOR_FL:
            if (forward)
            {
                Motor_WriteBridge(&htim2, TIM_CHANNEL_3, TIM_CHANNEL_4, pwm_val, 0U);
            }
            else
            {
                Motor_WriteBridge(&htim2, TIM_CHANNEL_3, TIM_CHANNEL_4, 0U, pwm_val);
            }
            break;

        case MOTOR_RL:
            if (forward)
            {
                Motor_WriteBridge(&htim9, TIM_CHANNEL_1, TIM_CHANNEL_2, pwm_val, 0U);
            }
            else
            {
                Motor_WriteBridge(&htim9, TIM_CHANNEL_1, TIM_CHANNEL_2, 0U, pwm_val);
            }
            break;

        case MOTOR_FR:
            if (forward)
            {
                Motor_WriteBridge(&htim8, TIM_CHANNEL_1, TIM_CHANNEL_2, pwm_val, 0U);
            }
            else
            {
                Motor_WriteBridge(&htim8, TIM_CHANNEL_1, TIM_CHANNEL_2, 0U, pwm_val);
            }
            break;

        case MOTOR_RR:
            if (forward)
            {
                Motor_WriteBridge(&htim8, TIM_CHANNEL_3, TIM_CHANNEL_4, pwm_val, 0U);
            }
            else
            {
                Motor_WriteBridge(&htim8, TIM_CHANNEL_3, TIM_CHANNEL_4, 0U, pwm_val);
            }
            break;

        default:
            break;
    }
}

void Motor_Brake_All(void)
{
    Motor_WriteBridge(&htim2, TIM_CHANNEL_3, TIM_CHANNEL_4, 0U, 0U);
    Motor_WriteBridge(&htim9, TIM_CHANNEL_1, TIM_CHANNEL_2, 0U, 0U);
    Motor_WriteBridge(&htim8, TIM_CHANNEL_1, TIM_CHANNEL_2, 0U, 0U);
    Motor_WriteBridge(&htim8, TIM_CHANNEL_3, TIM_CHANNEL_4, 0U, 0U);
}

void Motor_Update_RPM(float dt_s)
{
    uint16_t curr_count_fl;
    uint16_t curr_count_rl;
    uint16_t curr_count_fr;
    uint16_t curr_count_rr;
    int16_t delta[4];
    float rpm_scale;
    int i;

    if (dt_s <= 0.0f)
    {
        return;
    }

    curr_count_fl = (uint16_t)__HAL_TIM_GET_COUNTER(&htim5);
    curr_count_rl = (uint16_t)__HAL_TIM_GET_COUNTER(&htim3);
    curr_count_fr = (uint16_t)__HAL_TIM_GET_COUNTER(&htim4);
    curr_count_rr = (uint16_t)__HAL_TIM_GET_COUNTER(&htim1);

    delta[MOTOR_FL] = (int16_t)(curr_count_fl - s_last_count[MOTOR_FL]);
    delta[MOTOR_RL] = (int16_t)(curr_count_rl - s_last_count[MOTOR_RL]);
    delta[MOTOR_FR] = (int16_t)(curr_count_fr - s_last_count[MOTOR_FR]);
    delta[MOTOR_RR] = (int16_t)(curr_count_rr - s_last_count[MOTOR_RR]);

    s_last_count[MOTOR_FL] = curr_count_fl;
    s_last_count[MOTOR_RL] = curr_count_rl;
    s_last_count[MOTOR_FR] = curr_count_fr;
    s_last_count[MOTOR_RR] = curr_count_rr;

    delta[MOTOR_FL] = (int16_t)(-delta[MOTOR_FL]);
    delta[MOTOR_RL] = (int16_t)(-delta[MOTOR_RL]);

    s_odom_acc_ticks[MOTOR_FL] += delta[MOTOR_FL];
    s_odom_acc_ticks[MOTOR_RL] += delta[MOTOR_RL];
    s_odom_acc_ticks[MOTOR_FR] += delta[MOTOR_FR];
    s_odom_acc_ticks[MOTOR_RR] += delta[MOTOR_RR];

    /* 速度估计窗口累加 */
    s_speed_acc_ticks[MOTOR_FL] += delta[MOTOR_FL];
    s_speed_acc_ticks[MOTOR_RL] += delta[MOTOR_RL];
    s_speed_acc_ticks[MOTOR_FR] += delta[MOTOR_FR];
    s_speed_acc_ticks[MOTOR_RR] += delta[MOTOR_RR];
    s_speed_window_dt += dt_s;
    s_speed_window_count++;

    if (s_speed_window_count >= MOTOR_SPEED_WINDOW_SAMPLES)
    {
        rpm_scale = 60.0f / ((float)PULSES_PER_REV * s_speed_window_dt);

        for (i = 0; i < 4; ++i)
        {
            s_current_rpm[i] = (float)s_speed_acc_ticks[i] * rpm_scale;
            s_speed_acc_ticks[i] = 0;
        }

        s_speed_window_count = 0U;
        s_speed_window_dt = 0.0f;
    }
}

float Get_Motor_RPM(Motor_ID_t id)
{
    if ((id < MOTOR_FL) || (id > MOTOR_RR))
    {
        return 0.0f;
    }

    return s_current_rpm[id];
}

void Motor_Get_And_Clear_Delta_Ticks(int16_t *d_fl, int16_t *d_rl, int16_t *d_fr, int16_t *d_rr)
{
    int32_t fl;
    int32_t rl;
    int32_t fr;
    int32_t rr;
    uint32_t primask;

    if ((d_fl == NULL) || (d_rl == NULL) || (d_fr == NULL) || (d_rr == NULL))
    {
        return;
    }

    primask = __get_PRIMASK();
    __disable_irq();

    fl = s_odom_acc_ticks[MOTOR_FL];
    rl = s_odom_acc_ticks[MOTOR_RL];
    fr = s_odom_acc_ticks[MOTOR_FR];
    rr = s_odom_acc_ticks[MOTOR_RR];

    s_odom_acc_ticks[MOTOR_FL] = 0;
    s_odom_acc_ticks[MOTOR_RL] = 0;
    s_odom_acc_ticks[MOTOR_FR] = 0;
    s_odom_acc_ticks[MOTOR_RR] = 0;

    __set_PRIMASK(primask);

    *d_fl = Motor_SaturateToI16(fl);
    *d_rl = Motor_SaturateToI16(rl);
    *d_fr = Motor_SaturateToI16(fr);
    *d_rr = Motor_SaturateToI16(rr);
}
