/*
 * test_pid_schedule.c -- host-side unit tests for pid_schedule (Issue #550)
 *
 * Build:
 *   gcc -I /tmp/stub_hal -I include -DTEST_HOST -lm \
 *       -o test_pid_schedule test/test_pid_schedule.c
 *
 * Run:
 *   ./test_pid_schedule
 */

/* ---- Minimal HAL stub (no hardware) ---- */
#ifndef STM32F7XX_HAL_H
#define STM32F7XX_HAL_H
#include <stdint.h>
#include <stdbool.h>
typedef enum { HAL_OK = 0 } HAL_StatusTypeDef;
typedef struct { uint32_t TypeErase; uint32_t Sector; uint32_t NbSectors; uint32_t VoltageRange; } FLASH_EraseInitTypeDef;
#define FLASH_TYPEERASE_SECTORS 0
#define FLASH_SECTOR_7          7
#define VOLTAGE_RANGE_3         3
#define FLASH_TYPEPROGRAM_WORD  0
static inline HAL_StatusTypeDef HAL_FLASH_Unlock(void) { return HAL_OK; }
static inline HAL_StatusTypeDef HAL_FLASH_Lock(void)   { return HAL_OK; }
static inline HAL_StatusTypeDef HAL_FLASHEx_Erase(FLASH_EraseInitTypeDef *e, uint32_t *err) { (void)e; *err = 0xFFFFFFFFUL; return HAL_OK; }
static inline HAL_StatusTypeDef HAL_FLASH_Program(uint32_t t, uint32_t addr, uint64_t data) { (void)t; (void)addr; (void)data; return HAL_OK; }
static inline uint32_t HAL_GetTick(void) { return 0; }
#endif

/* ---- Block flash/jlink/balance headers (not needed) ---- */
/* pid_flash.h is included via pid_schedule.h -- stub flash functions */

/* Forward-declare stubs for pid_flash functions (used by pid_schedule.c) */
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <stdio.h>
#include <math.h>

/* Minimal pid_sched_entry_t and pid_flash stubs before pulling in schedule */
#define PID_FLASH_H  /* prevent pid_flash.h from being re-included */

/* Replicate types from pid_flash.h */
#define PID_SCHED_MAX_BANDS     6u
#define PID_SCHED_FLASH_ADDR    0x0807FF40UL
#define PID_SCHED_MAGIC         0x534C5402UL
#define PID_FLASH_STORE_ADDR    0x0807FFC0UL
#define PID_FLASH_MAGIC         0x534C5401UL
#define PID_FLASH_SECTOR        7
#define PID_FLASH_SECTOR_VOLTAGE 3

typedef struct __attribute__((packed)) {
    float speed_mps;
    float kp;
    float ki;
    float kd;
} pid_sched_entry_t;

typedef struct __attribute__((packed)) {
    uint32_t         magic;
    uint8_t          num_bands;
    uint8_t          flags;
    uint8_t          _pad0[2];
    pid_sched_entry_t bands[PID_SCHED_MAX_BANDS];
    uint8_t          _pad1[24];
} pid_sched_flash_t;

typedef struct __attribute__((packed)) {
    uint32_t magic;
    float    kp;
    float    ki;
    float    kd;
    uint8_t  _pad[48];
} pid_flash_t;

/* Stub flash storage (simulated in RAM) */
static pid_sched_flash_t g_sched_flash;
static pid_flash_t       g_pid_flash;
static bool              g_sched_flash_valid = false;
static bool              g_pid_flash_valid   = false;

bool pid_flash_load(float *kp, float *ki, float *kd)
{
    if (!g_pid_flash_valid || g_pid_flash.magic != PID_FLASH_MAGIC) return false;
    *kp = g_pid_flash.kp; *ki = g_pid_flash.ki; *kd = g_pid_flash.kd;
    return true;
}

bool pid_flash_save(float kp, float ki, float kd)
{
    g_pid_flash.magic = PID_FLASH_MAGIC;
    g_pid_flash.kp = kp; g_pid_flash.ki = ki; g_pid_flash.kd = kd;
    g_pid_flash_valid = true;
    return true;
}

bool pid_flash_load_schedule(pid_sched_entry_t *out_entries, uint8_t *out_n)
{
    if (!g_sched_flash_valid || g_sched_flash.magic != PID_SCHED_MAGIC) return false;
    if (g_sched_flash.num_bands == 0 || g_sched_flash.num_bands > PID_SCHED_MAX_BANDS) return false;
    memcpy(out_entries, g_sched_flash.bands, g_sched_flash.num_bands * sizeof(pid_sched_entry_t));
    *out_n = g_sched_flash.num_bands;
    return true;
}

bool pid_flash_save_all(float kp_s, float ki_s, float kd_s,
                        const pid_sched_entry_t *entries, uint8_t num_bands)
{
    if (num_bands == 0 || num_bands > PID_SCHED_MAX_BANDS) return false;
    g_sched_flash.magic = PID_SCHED_MAGIC;
    g_sched_flash.num_bands = num_bands;
    memcpy(g_sched_flash.bands, entries, num_bands * sizeof(pid_sched_entry_t));
    g_sched_flash_valid = true;
    g_pid_flash.magic = PID_FLASH_MAGIC;
    g_pid_flash.kp = kp_s; g_pid_flash.ki = ki_s; g_pid_flash.kd = kd_s;
    g_pid_flash_valid = true;
    return true;
}

/* Stub mpu6000.h and balance.h so pid_schedule.h doesn't pull in hardware types */
#define MPU6000_H
typedef struct { float ax, ay, az, gx, gy, gz, pitch, pitch_rate; } IMUData;

#define BALANCE_H
typedef enum { BALANCE_DISARMED=0, BALANCE_ARMED, BALANCE_TILT_FAULT } balance_state_t;
typedef struct {
    balance_state_t state;
    float pitch_deg, pitch_rate;
    float integral, prev_error;
    int16_t motor_cmd;
    float kp, ki, kd, setpoint, max_tilt;
    int16_t max_speed;
} balance_t;

/* Include the implementation directly */
#include "../src/pid_schedule.c"

/* ============================================================
 * Test framework
 * ============================================================ */
static int g_pass = 0, g_fail = 0;

#define ASSERT(cond, msg) do { \
    if (cond) { g_pass++; } \
    else { g_fail++; printf("FAIL [%s:%d] %s\n", __FILE__, __LINE__, msg); } \
} while (0)

#define ASSERT_NEAR(a, b, eps, msg) ASSERT(fabsf((a)-(b)) < (eps), msg)

static void reset_flash(void)
{
    g_sched_flash_valid = false;
    g_pid_flash_valid   = false;
    memset(&g_sched_flash, 0xFF, sizeof(g_sched_flash));
    memset(&g_pid_flash,   0xFF, sizeof(g_pid_flash));
}

/* ============================================================
 * Tests
 * ============================================================ */

static void test_init_loads_default_when_flash_empty(void)
{
    reset_flash();
    pid_schedule_init();
    ASSERT(pid_schedule_get_num_bands() == 3u, "default 3 bands");
    pid_sched_entry_t tbl[PID_SCHED_MAX_BANDS];
    uint8_t n;
    pid_schedule_get_table(tbl, &n);
    ASSERT(n == 3u, "get_table returns 3");
    ASSERT_NEAR(tbl[0].speed_mps, 0.00f, 1e-5f, "band0 speed=0.00");
    ASSERT_NEAR(tbl[0].kp, 40.0f, 1e-4f, "band0 kp=40");
    ASSERT_NEAR(tbl[2].speed_mps, 0.80f, 1e-5f, "band2 speed=0.80");
    ASSERT_NEAR(tbl[2].kp, 28.0f, 1e-4f, "band2 kp=28");
}

static void test_init_loads_from_flash_when_valid(void)
{
    reset_flash();
    pid_sched_entry_t entries[2] = {
        { .speed_mps = 0.0f, .kp = 10.0f, .ki = 0.5f, .kd = 0.2f },
        { .speed_mps = 1.0f, .kp = 20.0f, .ki = 0.8f, .kd = 0.4f },
    };
    pid_flash_save_all(1.0f, 0.1f, 0.1f, entries, 2u);
    pid_schedule_init();
    ASSERT(pid_schedule_get_num_bands() == 2u, "init loads 2 bands from flash");
    pid_sched_entry_t tbl[PID_SCHED_MAX_BANDS];
    uint8_t n;
    pid_schedule_get_table(tbl, &n);
    ASSERT_NEAR(tbl[1].kp, 20.0f, 1e-4f, "flash band1 kp=20");
}

static void test_get_gains_below_first_band(void)
{
    reset_flash();
    pid_schedule_init();  /* default table: 0.0, 0.3, 0.8 m/s */
    float kp, ki, kd;
    pid_schedule_get_gains(0.0f, &kp, &ki, &kd);
    ASSERT_NEAR(kp, 40.0f, 1e-4f, "speed=0 -> kp=40 (clamp low)");
    /* abs(-0.1)=0.1 m/s: between band0(0.0) and band1(0.3), t=1/3 -> kp=40+(35-40)/3 */
    pid_schedule_get_gains(-0.1f, &kp, &ki, &kd);
    ASSERT_NEAR(kp, 40.0f + (35.0f - 40.0f) * (0.1f / 0.3f), 0.01f,
                "speed=-0.1 interpolates via abs(speed)");
}

static void test_get_gains_above_last_band(void)
{
    reset_flash();
    pid_schedule_init();
    float kp, ki, kd;
    pid_schedule_get_gains(2.0f, &kp, &ki, &kd);
    ASSERT_NEAR(kp, 28.0f, 1e-4f, "speed=2.0 -> kp=28 (clamp high)");
}

static void test_get_gains_at_band_boundary(void)
{
    reset_flash();
    pid_schedule_init();
    float kp, ki, kd;
    pid_schedule_get_gains(0.30f, &kp, &ki, &kd);
    ASSERT_NEAR(kp, 35.0f, 1e-4f, "speed=0.30 exactly -> kp=35");
    pid_schedule_get_gains(0.80f, &kp, &ki, &kd);
    ASSERT_NEAR(kp, 28.0f, 1e-4f, "speed=0.80 exactly -> kp=28");
}

static void test_interpolation_midpoint(void)
{
    reset_flash();
    pid_schedule_init();
    /* Between band0 (0.0,kp=40) and band1 (0.3,kp=35): at t=0.5 -> kp=37.5 */
    float kp, ki, kd;
    pid_schedule_get_gains(0.15f, &kp, &ki, &kd);
    ASSERT_NEAR(kp, 37.5f, 0.01f, "interp midpoint kp=37.5");
    /* Between band1 (0.3,kp=35) and band2 (0.8,kp=28): at t=0.2 -> 35+(28-35)*0.2=33.6 */
    pid_schedule_get_gains(0.40f, &kp, &ki, &kd);
    float expected = 35.0f + (28.0f - 35.0f) * ((0.40f - 0.30f) / (0.80f - 0.30f));
    ASSERT_NEAR(kp, expected, 0.01f, "interp band1->2 kp");
}

static void test_interpolation_ki_kd(void)
{
    reset_flash();
    pid_schedule_init();
    float kp, ki, kd;
    pid_schedule_get_gains(0.15f, &kp, &ki, &kd);
    /* ki: band0=1.5, band1=1.0, t=0.5 -> 1.25 */
    ASSERT_NEAR(ki, 1.25f, 0.01f, "interp midpoint ki=1.25");
    /* kd: band0=1.2, band1=1.0, t=0.5 -> 1.1 */
    ASSERT_NEAR(kd, 1.1f,  0.01f, "interp midpoint kd=1.1");
}

static void test_set_table_and_sort(void)
{
    pid_sched_entry_t tbl[3] = {
        { .speed_mps = 0.8f, .kp = 5.0f, .ki = 0.1f, .kd = 0.1f },
        { .speed_mps = 0.0f, .kp = 9.0f, .ki = 0.3f, .kd = 0.3f },
        { .speed_mps = 0.4f, .kp = 7.0f, .ki = 0.2f, .kd = 0.2f },
    };
    pid_schedule_set_table(tbl, 3u);
    ASSERT(pid_schedule_get_num_bands() == 3u, "set_table 3 bands");
    pid_sched_entry_t out[PID_SCHED_MAX_BANDS];
    uint8_t n;
    pid_schedule_get_table(out, &n);
    /* After sort: 0.0, 0.4, 0.8 */
    ASSERT_NEAR(out[0].speed_mps, 0.0f, 1e-5f, "sorted[0]=0.0");
    ASSERT_NEAR(out[1].speed_mps, 0.4f, 1e-5f, "sorted[1]=0.4");
    ASSERT_NEAR(out[2].speed_mps, 0.8f, 1e-5f, "sorted[2]=0.8");
}

static void test_set_table_clamps_n(void)
{
    pid_sched_entry_t big[8];
    memset(big, 0, sizeof(big));
    for (int i = 0; i < 8; i++) big[i].speed_mps = (float)i * 0.1f;
    pid_schedule_set_table(big, 8u);
    ASSERT(pid_schedule_get_num_bands() == PID_SCHED_MAX_BANDS, "clamp to MAX_BANDS");
}

static void test_set_table_min_1(void)
{
    pid_sched_entry_t one = { .speed_mps = 0.5f, .kp = 30.0f, .ki = 1.0f, .kd = 0.8f };
    pid_schedule_set_table(&one, 0u);  /* n=0 clamped to 1 */
    ASSERT(pid_schedule_get_num_bands() == 1u, "min 1 band");
}

static void test_active_band_idx_clamp_low(void)
{
    reset_flash();
    pid_schedule_init();
    float kp, ki, kd;
    pid_schedule_get_gains(0.0f, &kp, &ki, &kd);
    ASSERT(pid_schedule_active_band_idx() == 0u, "active=0 when clamped low");
}

static void test_active_band_idx_interpolating(void)
{
    reset_flash();
    pid_schedule_init();
    float kp, ki, kd;
    pid_schedule_get_gains(0.5f, &kp, &ki, &kd);  /* between band1 and band2 */
    ASSERT(pid_schedule_active_band_idx() == 1u, "active=1 between band1-2");
}

static void test_active_band_idx_clamp_high(void)
{
    reset_flash();
    pid_schedule_init();
    float kp, ki, kd;
    pid_schedule_get_gains(5.0f, &kp, &ki, &kd);
    ASSERT(pid_schedule_active_band_idx() == 2u, "active=2 when clamped high");
}

static void test_apply_writes_gains(void)
{
    reset_flash();
    pid_schedule_init();
    balance_t b;
    memset(&b, 0, sizeof(b));
    pid_schedule_apply(&b, 0.0f);
    ASSERT_NEAR(b.kp, 40.0f, 1e-4f, "apply: kp written");
    ASSERT_NEAR(b.ki, 1.5f,  1e-4f, "apply: ki written");
    ASSERT_NEAR(b.kd, 1.2f,  1e-4f, "apply: kd written");
}

static void test_apply_resets_integral_on_band_change(void)
{
    reset_flash();
    pid_schedule_init();
    balance_t b;
    memset(&b, 0, sizeof(b));
    b.integral = 99.0f;

    /* First call: sets s_prev_band from sentinel -> band 0 (integral reset) */
    pid_schedule_apply(&b, 0.0f);
    ASSERT_NEAR(b.integral, 0.0f, 1e-6f, "apply: integral reset on first call");

    b.integral = 77.0f;
    pid_schedule_apply(&b, 0.0f);  /* same band -- no reset */
    ASSERT_NEAR(b.integral, 77.0f, 1e-6f, "apply: integral preserved same band");

    b.integral = 55.0f;
    pid_schedule_apply(&b, 0.5f);  /* band changes 0->1 -- reset */
    ASSERT_NEAR(b.integral, 0.0f, 1e-6f, "apply: integral reset on band change");
}

static void test_flash_save_and_reload(void)
{
    reset_flash();
    pid_sched_entry_t tbl[2] = {
        { .speed_mps = 0.0f, .kp = 15.0f, .ki = 0.6f, .kd = 0.3f },
        { .speed_mps = 0.5f, .kp = 10.0f, .ki = 0.4f, .kd = 0.2f },
    };
    pid_schedule_set_table(tbl, 2u);
    bool ok = pid_schedule_flash_save(25.0f, 1.1f, 0.9f);
    ASSERT(ok, "flash_save returns true");
    ASSERT(g_sched_flash_valid, "flash_save wrote sched record");
    ASSERT(g_pid_flash_valid,   "flash_save wrote pid record");
    ASSERT_NEAR(g_pid_flash.kp, 25.0f, 1e-4f, "pid kp saved");

    /* Now reload */
    pid_schedule_init();
    ASSERT(pid_schedule_get_num_bands() == 2u, "reload 2 bands");
    float kp, ki, kd;
    pid_schedule_get_gains(0.0f, &kp, &ki, &kd);
    ASSERT_NEAR(kp, 15.0f, 1e-4f, "reload kp at speed=0");
}

static void test_get_default_table(void)
{
    pid_sched_entry_t def[PID_SCHED_MAX_BANDS];
    uint8_t n;
    pid_schedule_get_default_table(def, &n);
    ASSERT(n == 3u, "default table has 3 entries");
    ASSERT_NEAR(def[0].kp, 40.0f, 1e-4f, "default[0] kp=40");
    ASSERT_NEAR(def[1].kp, 35.0f, 1e-4f, "default[1] kp=35");
    ASSERT_NEAR(def[2].kp, 28.0f, 1e-4f, "default[2] kp=28");
}

static void test_init_discards_invalid_flash(void)
{
    reset_flash();
    /* Write a valid record but with out-of-range gain */
    pid_sched_entry_t bad[1] = {{ .speed_mps=0.0f, .kp=999.0f, .ki=0.1f, .kd=0.1f }};
    pid_flash_save_all(1.0f, 0.1f, 0.1f, bad, 1u);
    pid_schedule_init();
    /* Should fall back to default */
    ASSERT(pid_schedule_get_num_bands() == 3u, "invalid flash -> default 3 bands");
}

static void test_single_band_clamps_both_ends(void)
{
    pid_sched_entry_t one = { .speed_mps = 0.5f, .kp = 50.0f, .ki = 2.0f, .kd = 1.5f };
    pid_schedule_set_table(&one, 1u);
    float kp, ki, kd;
    pid_schedule_get_gains(0.0f, &kp, &ki, &kd);
    ASSERT_NEAR(kp, 50.0f, 1e-4f, "single band: clamp low -> kp=50");
    pid_schedule_get_gains(9.9f, &kp, &ki, &kd);
    ASSERT_NEAR(kp, 50.0f, 1e-4f, "single band: clamp high -> kp=50");
}

static void test_negative_speed_symmetric(void)
{
    reset_flash();
    pid_schedule_init();
    float kp_fwd, ki_fwd, kd_fwd;
    float kp_rev, ki_rev, kd_rev;
    pid_schedule_get_gains( 0.5f, &kp_fwd, &ki_fwd, &kd_fwd);
    pid_schedule_get_gains(-0.5f, &kp_rev, &ki_rev, &kd_rev);
    ASSERT_NEAR(kp_fwd, kp_rev, 1e-5f, "symmetric: kp same for +/-speed");
    ASSERT_NEAR(ki_fwd, ki_rev, 1e-5f, "symmetric: ki same for +/-speed");
    ASSERT_NEAR(kd_fwd, kd_rev, 1e-5f, "symmetric: kd same for +/-speed");
}

int main(void)
{
    printf("=== test_pid_schedule ===\n");

    test_init_loads_default_when_flash_empty();
    test_init_loads_from_flash_when_valid();
    test_get_gains_below_first_band();
    test_get_gains_above_last_band();
    test_get_gains_at_band_boundary();
    test_interpolation_midpoint();
    test_interpolation_ki_kd();
    test_set_table_and_sort();
    test_set_table_clamps_n();
    test_set_table_min_1();
    test_active_band_idx_clamp_low();
    test_active_band_idx_interpolating();
    test_active_band_idx_clamp_high();
    test_apply_writes_gains();
    test_apply_resets_integral_on_band_change();
    test_flash_save_and_reload();
    test_get_default_table();
    test_init_discards_invalid_flash();
    test_single_band_clamps_both_ends();
    test_negative_speed_symmetric();

    printf("PASSED: %d  FAILED: %d\n", g_pass, g_fail);
    return (g_fail == 0) ? 0 : 1;
}