#include "balancer.h"

void Balancer_Init(Balancer *bal)
{
    /* Motors */
    Motor_Init(&bal->motor_l, &htim2, TIM_CHANNEL_1,
               GPIOA, GPIO_PIN_1, GPIOA, GPIO_PIN_2);
    Motor_Init(&bal->motor_r, &htim2, TIM_CHANNEL_2,
               GPIOA, GPIO_PIN_4, GPIOA, GPIO_PIN_5);

    /* Encoders (both share TIM3) */
    Encoder_Init(&bal->enc_l, &htim3);
    Encoder_Init(&bal->enc_r, &htim3);    /* share same timer */

    /* PID: angle loop */
    PID_Init(&bal->pid_angle, 30.0f, 0.5f, 1.5f,
             -1000.0f, 1000.0f);

    /* PID: speed loop */
    PID_Init(&bal->pid_speed, 0.5f, 0.05f, 0.0f,
             -30.0f, 30.0f);   /* output = angle reference in degrees */

    /* Control state */
    bal->angle_ref = 0.0f;
    bal->speed_ref = 0.0f;
    bal->motor_output = 0.0f;
    bal->motor_output_l = 0.0f;
    bal->motor_output_r = 0.0f;

    bal->tick_angle = 0;
    bal->tick_speed = 0;

    bal->calibrated = 0;
    bal->running = 0;
    bal->fault = 0;

    bal->gyro_bias_x = 0;
    bal->gyro_bias_y = 0;
    bal->gyro_bias_z = 0;
}

void Balancer_Calibrate(Balancer *bal)
{
    PID_Reset(&bal->pid_angle);
    PID_Reset(&bal->pid_speed);
    bal->angle_ref = 0.0f;
    bal->motor_output = 0.0f;
    bal->fault = 0;
}

void Balancer_Run(Balancer *bal, float dt_angle, float dt_speed)
{
    /* Speed PID (outer) */
    float avg_speed = (bal->enc_l.speed_filtered +
                       bal->enc_r.speed_filtered) * 0.5f;

    bal->angle_ref = PID_Update(&bal->pid_speed,
                                 bal->speed_ref, avg_speed, dt_speed);

    /* Angle PID (inner) */
    bal->motor_output = PID_Update(&bal->pid_angle,
                                    bal->angle_ref, bal->att.pitch, dt_angle);

    /* Apply to motors */
    bal->motor_output_l = bal->motor_output;
    bal->motor_output_r = bal->motor_output;

    Motor_SetSpeedScaled(&bal->motor_l, (int16_t)bal->motor_output_l);
    Motor_SetSpeedScaled(&bal->motor_r, (int16_t)bal->motor_output_r);
}

void Balancer_EmergencyStop(Balancer *bal)
{
    Motor_Stop(&bal->motor_l);
    Motor_Stop(&bal->motor_r);
    bal->motor_output = 0.0f;
    bal->motor_output_l = 0.0f;
    bal->motor_output_r = 0.0f;
    bal->running = 0;
    bal->fault = 1;
}