#include "pid_schedule.h"
#include "pid_flash.h"
#include <string.h>
#include <math.h>   /* fabsf */

/* ---- Default 3-band table ---- */
static const pid_sched_entry_t k_default_table[3] = {
    { .speed_mps = 0.00f, .kp = 40.0f, .ki = 1.5f, .kd = 1.2f },
    { .speed_mps = 0.30f, .kp = 35.0f, .ki = 1.0f, .kd = 1.0f },
    { .speed_mps = 0.80f, .kp = 28.0f, .ki = 0.5f, .kd = 0.8f },
};

/* ---- Active table ---- */
static pid_sched_entry_t s_bands[PID_SCHED_MAX_BANDS];
static uint8_t           s_num_bands    = 0u;
static uint8_t           s_active_band  = 0u;   /* lower-bracket index of last call */
static uint8_t           s_prev_band    = 0xFFu; /* sentinel: forces integrator reset on first apply */

/* ---- sort helper (insertion sort — table is small, ≤6 entries) ---- */
static void sort_bands(void)
{
    for (uint8_t i = 1u; i < s_num_bands; i++) {
        pid_sched_entry_t key = s_bands[i];
        int8_t j = (int8_t)(i - 1u);
        while (j >= 0 && s_bands[j].speed_mps > key.speed_mps) {
            s_bands[j + 1] = s_bands[j];
            j--;
        }
        s_bands[j + 1] = key;
    }
}

/* ---- pid_schedule_init() ---- */
void pid_schedule_init(void)
{
    pid_sched_entry_t tmp[PID_SCHED_MAX_BANDS];
    uint8_t           n = 0u;

    if (pid_flash_load_schedule(tmp, &n)) {
        /* Validate entries minimally */
        bool ok = true;
        for (uint8_t i = 0u; i < n; i++) {
            if (tmp[i].kp < 0.0f || tmp[i].kp > 500.0f ||
                tmp[i].ki < 0.0f || tmp[i].ki > 50.0f  ||
                tmp[i].kd < 0.0f || tmp[i].kd > 50.0f  ||
                tmp[i].speed_mps < 0.0f) {
                ok = false;
                break;
            }
        }
        if (ok) {
            memcpy(s_bands, tmp, n * sizeof(pid_sched_entry_t));
            s_num_bands = n;
            sort_bands();
            s_active_band = 0u;
            return;
        }
    }

    /* Fall back to built-in default */
    memcpy(s_bands, k_default_table, sizeof(k_default_table));
    s_num_bands   = 3u;
    s_active_band = 0u;
    s_prev_band   = 0xFFu;
}

/* ---- pid_schedule_get_gains() ---- */
void pid_schedule_get_gains(float speed_mps, float *kp, float *ki, float *kd)
{
    float spd = fabsf(speed_mps);

    if (s_num_bands == 0u) {
        *kp = k_default_table[0].kp;
        *ki = k_default_table[0].ki;
        *kd = k_default_table[0].kd;
        return;
    }

    /* Clamp below first entry */
    if (spd <= s_bands[0].speed_mps) {
        s_active_band = 0u;
        *kp = s_bands[0].kp;
        *ki = s_bands[0].ki;
        *kd = s_bands[0].kd;
        return;
    }

    /* Clamp above last entry */
    uint8_t last = s_num_bands - 1u;
    if (spd >= s_bands[last].speed_mps) {
        s_active_band = last;
        *kp = s_bands[last].kp;
        *ki = s_bands[last].ki;
        *kd = s_bands[last].kd;
        return;
    }

    /* Find bracket [i-1, i] where bands[i-1].speed <= spd < bands[i].speed */
    uint8_t i = 1u;
    while (i < s_num_bands && s_bands[i].speed_mps <= spd) i++;
    /* Now bands[i-1].speed_mps <= spd < bands[i].speed_mps */

    s_active_band = (uint8_t)(i - 1u);

    float dv = s_bands[i].speed_mps - s_bands[i - 1u].speed_mps;
    float t  = (dv > 0.0f) ? (spd - s_bands[i - 1u].speed_mps) / dv : 0.0f;

    *kp = s_bands[i - 1u].kp + t * (s_bands[i].kp - s_bands[i - 1u].kp);
    *ki = s_bands[i - 1u].ki + t * (s_bands[i].ki - s_bands[i - 1u].ki);
    *kd = s_bands[i - 1u].kd + t * (s_bands[i].kd - s_bands[i - 1u].kd);
}

/* ---- pid_schedule_apply() ---- */
void pid_schedule_apply(balance_t *b, float speed_mps)
{
    float kp, ki, kd;
    pid_schedule_get_gains(speed_mps, &kp, &ki, &kd);

    b->kp = kp;
    b->ki = ki;
    b->kd = kd;

    /* Reset integrator on band transition to prevent windup spike */
    if (s_active_band != s_prev_band) {
        b->integral   = 0.0f;
        s_prev_band   = s_active_band;
    }
}

/* ---- pid_schedule_set_table() ---- */
void pid_schedule_set_table(const pid_sched_entry_t *entries, uint8_t n)
{
    if (n == 0u) n = 1u;
    if (n > PID_SCHED_MAX_BANDS) n = PID_SCHED_MAX_BANDS;

    memcpy(s_bands, entries, n * sizeof(pid_sched_entry_t));
    s_num_bands   = n;
    s_active_band = 0u;
    s_prev_band   = 0xFFu;
    sort_bands();
}

/* ---- pid_schedule_get_table() ---- */
void pid_schedule_get_table(pid_sched_entry_t *out_entries, uint8_t *out_n)
{
    memcpy(out_entries, s_bands, s_num_bands * sizeof(pid_sched_entry_t));
    *out_n = s_num_bands;
}

/* ---- pid_schedule_get_num_bands() ---- */
uint8_t pid_schedule_get_num_bands(void)
{
    return s_num_bands;
}

/* ---- pid_schedule_flash_save() ---- */
bool pid_schedule_flash_save(float kp_single, float ki_single, float kd_single)
{
    return pid_flash_save_all(kp_single, ki_single, kd_single,
                              s_bands, s_num_bands);
}

/* ---- pid_schedule_active_band_idx() ---- */
uint8_t pid_schedule_active_band_idx(void)
{
    return s_active_band;
}

/* ---- pid_schedule_get_default_table() ---- */
void pid_schedule_get_default_table(pid_sched_entry_t *out_entries, uint8_t *out_n)
{
    memcpy(out_entries, k_default_table, sizeof(k_default_table));
    *out_n = 3u;
}