#include "balance.h"
#include "slope_estimator.h"
#include "config.h"
#include <math.h>

void balance_init(balance_t *b) {
    b->state = BALANCE_DISARMED;
    b->pitch_deg = 0.0f;
    b->pitch_rate = 0.0f;
    b->integral = 0.0f;
    b->prev_error = 0.0f;
    b->motor_cmd = 0;

    /* Default PID — conservative starting point */
    b->kp = PID_KP;
    b->ki = PID_KI;
    b->kd = PID_KD;
    b->setpoint = 0.0f;

    /* Safety defaults from config */
    b->max_tilt = MAX_TILT_DEG;
    b->max_speed = MAX_SPEED_LIMIT;

    /* Slope estimator */
    slope_estimator_init(&b->slope);
}

void balance_update(balance_t *b, const IMUData *imu, float dt) {
    if (dt <= 0.0f || dt > 0.1f) return;  /* Sanity check dt */

    /* Consume fused angle from mpu6000 complementary filter */
    b->pitch_deg  = imu->pitch;
    b->pitch_rate = imu->pitch_rate;

    /* Advance slope estimator (no-op when not armed) */
    slope_estimator_update(&b->slope, b->pitch_deg, dt);

    /* Safety: tilt cutoff */
    if (b->state == BALANCE_ARMED) {
        if (fabsf(b->pitch_deg) > b->max_tilt) {
            b->state = BALANCE_TILT_FAULT;
            b->motor_cmd = 0;
            b->integral = 0.0f;
            slope_estimator_reset(&b->slope);
            return;
        }
    }

    if (b->state != BALANCE_ARMED) {
        b->motor_cmd = 0;
        b->integral = 0.0f;
        b->prev_error = 0.0f;
        return;
    }

    /* PID — subtract slope estimate so controller balances on incline */
    float tilt_corrected = b->pitch_deg - slope_estimator_get_deg(&b->slope);
    float error = b->setpoint - tilt_corrected;

    /* Proportional */
    float p_term = b->kp * error;

    /* Integral with anti-windup */
    b->integral += error * dt;
    if (b->integral > PID_INTEGRAL_MAX) b->integral = PID_INTEGRAL_MAX;
    if (b->integral < -PID_INTEGRAL_MAX) b->integral = -PID_INTEGRAL_MAX;
    float i_term = b->ki * b->integral;

    /* Derivative on measurement (avoids setpoint kick) */
    float d_term = b->kd * (-b->pitch_rate);

    /* Sum and clamp */
    float output = p_term + i_term + d_term;
    if (output > (float)b->max_speed) output = (float)b->max_speed;
    if (output < -(float)b->max_speed) output = -(float)b->max_speed;

    b->motor_cmd = (int16_t)output;
    b->prev_error = error;
}

void balance_arm(balance_t *b) {
    if (b->state == BALANCE_DISARMED) {
        /* Only arm if roughly upright */
        if (fabsf(b->pitch_deg) < 10.0f) {
            b->integral = 0.0f;
            b->prev_error = 0.0f;
            b->motor_cmd = 0;
            b->state = BALANCE_ARMED;
        }
    }
}

void balance_disarm(balance_t *b) {
    b->state = BALANCE_DISARMED;
    b->motor_cmd = 0;
    b->integral = 0.0f;
    slope_estimator_reset(&b->slope);
}

void balance_park(balance_t *b) {
    /* Suspend balancing from ARMED state only — keeps robot stationary on flat ground */
    if (b->state == BALANCE_ARMED) {
        b->state     = BALANCE_PARKED;
        b->motor_cmd = 0;
        b->integral  = 0.0f;
        b->prev_error = 0.0f;
        slope_estimator_reset(&b->slope);
    }
}

void balance_unpark(balance_t *b) {
    /* Quick re-arm from PARKED — only if pitch is safe (< 20 deg) */
    if (b->state == BALANCE_PARKED) {
        if (fabsf(b->pitch_deg) < 20.0f) {
            b->motor_cmd  = 0;
            b->prev_error = 0.0f;
            b->state      = BALANCE_ARMED;
        }
        /* If pitch too large, stay PARKED — caller checks resulting state */
    }
}
