// ============================================================================
// WaveForge Expose — L0/L1 Feature Emitter  (implementation)
//
// Compute strategy — no IQ available on bare Heltec V4:
//
//   Spectrum proxy:  The 128-bin normalised RSSI array produced by
//   capture_rf_spectrum() is treated as a power-envelope signal x[k].
//   We compute its 2nd- and 4th-order sample cumulants across the 128 bins.
//   While these are not true IQ cumulants, they still discriminate modulation
//   families (e.g. FHSS bursts vs CW vs wideband noise) and are exactly what
//   the server-side L1 estimator (l1_estimator.py) needs when no SDR is
//   attached.  The imaginary parts of C20 and C40 are zero for a real-valued
//   signal — they are included in the wire format so the server can handle
//   both SDR and proxy sensors with a single code path.
//
//   Rolling RSSI/SNR window: last PACKET_BUF_SIZE observations per band,
//   statically allocated.  Provides mean, std, inter-arrival jitter, and
//   duty-cycle occupancy.
//
// Static memory only — no heap, no STL.
// ============================================================================

#include "wfx_features.h"

#include <cstring>
#include <cstdlib>    // abs (integer)
#include <cmath>      // sqrtf, fabsf

#ifdef UNIT_TEST
  // ---- Host-test shims ----
  // mesh_publish_raw is forward-declared; the test TU provides the stub.
  // config.h constants we need inline here:
  #ifndef DEFAULT_SENSOR_ID
  #define DEFAULT_SENSOR_ID 0x00000001u
  #endif
  #ifndef MSG_TYPE_WFX_FEATURES
  #define MSG_TYPE_WFX_FEATURES 0x10u
  #endif
  #include <cstdint>
  #include <cstddef>
  // Forward declaration — resolved by the stub in test_wfx.cpp
  bool mesh_publish_raw(const uint8_t* data, size_t len);
#else
  #include "mesh_publish.h"   // mesh_publish_raw()
  #include "config.h"         // sensor constants, MSG_TYPE_*
#endif

// ---------------------------------------------------------------------------
// Internal constants
// ---------------------------------------------------------------------------
#define PACKET_BUF_SIZE  64u   // rolling observation window depth per band
#define NUM_BANDS        2u    // 0 = 2.4 GHz WiFi, 1 = 915 MHz LoRa

// WFX packet layout (67 bytes total):
//   [0]        msg_type  = 0x10
//   [1..4]     sensor_id (big-endian)
//   [5..64]    WfxFeatureVec  (60 bytes, packed)
//   [65..66]   CRC-16/CCITT over bytes [0..64]
#define WFX_PACKET_SIZE  67u
#define WFX_HDR_SIZE     5u    // msg_type + sensor_id
#define WFX_CRC_OFFSET   65u
#ifndef MSG_TYPE_WFX_FEATURES
#define MSG_TYPE_WFX_FEATURES  0x10u
#endif

// ---------------------------------------------------------------------------
// Per-band rolling observation buffer
// ---------------------------------------------------------------------------
struct PacketObs {
    int16_t  rssi;    // dBm
    int16_t  snr;     // dB (tenths or raw, as supplied by SX1262)
    uint32_t ts_ms;   // millis() at reception
};

struct BandBuf {
    PacketObs obs[PACKET_BUF_SIZE];
    uint32_t  head;    // next write index (wraps)
    uint32_t  count;   // total inserts (clamped to PACKET_BUF_SIZE for window)
};

// ---------------------------------------------------------------------------
// Module state — fully static
// ---------------------------------------------------------------------------
static BandBuf  s_bands[NUM_BANDS];
static uint32_t s_sensor_id = DEFAULT_SENSOR_ID;

// ---------------------------------------------------------------------------
// CRC-16/CCITT (polynomial 0x1021, init 0xFFFF, no reflection)
// Matches the implementation used in band_bridge.cpp.
// ---------------------------------------------------------------------------
static uint16_t wfx_crc16(const uint8_t* data, size_t len) {
    uint16_t crc = 0xFFFFu;
    for (size_t i = 0; i < len; i++) {
        crc ^= (uint16_t)data[i] << 8;
        for (int b = 0; b < 8; b++) {
            crc = (crc & 0x8000u) ? (uint16_t)((crc << 1) ^ 0x1021u)
                                  : (uint16_t)(crc << 1);
        }
    }
    return crc;
}

// ---------------------------------------------------------------------------
// wfx_features_init
// ---------------------------------------------------------------------------
void wfx_features_init() {
    memset(s_bands, 0, sizeof(s_bands));
    // sensor_id: in real firmware this is read from NVS; use compile-time default.
    s_sensor_id = DEFAULT_SENSOR_ID;
}

// ---------------------------------------------------------------------------
// wfx_features_observe_iq  (reserved — no-op on bare ESP32)
// ---------------------------------------------------------------------------
void wfx_features_observe_iq(const float* /*i*/, const float* /*q*/,
                              size_t /*n*/, uint8_t /*band_id*/) {
    // No IQ frontend available on Heltec V4 + SX1262.
    // When a future RadioBox-Nano SDR front-end is attached, implement here.
}

// ---------------------------------------------------------------------------
// wfx_features_observe_rssi_packet
// ---------------------------------------------------------------------------
void wfx_features_observe_rssi_packet(int16_t rssi, int16_t snr,
                                      uint32_t ts_ms, uint8_t band_id) {
    if (band_id >= NUM_BANDS) return;

    BandBuf& b = s_bands[band_id];
    uint32_t idx = b.head % PACKET_BUF_SIZE;
    b.obs[idx].rssi  = rssi;
    b.obs[idx].snr   = snr;
    b.obs[idx].ts_ms = ts_ms;
    b.head++;
    if (b.count < PACKET_BUF_SIZE) b.count++;
}

// ---------------------------------------------------------------------------
// Internal helpers — rolling-window statistics
// ---------------------------------------------------------------------------

// Number of valid entries in a band buffer
static uint32_t buf_valid_count(const BandBuf& b) {
    return (b.count < PACKET_BUF_SIZE) ? b.count : PACKET_BUF_SIZE;
}

// Mean of RSSI over rolling window (float)
static float rssi_mean(const BandBuf& b) {
    uint32_t n = buf_valid_count(b);
    if (n == 0) return 0.0f;
    float sum = 0.0f;
    for (uint32_t i = 0; i < n; i++) {
        sum += (float)b.obs[i].rssi;
    }
    return sum / (float)n;
}

// Std-dev of RSSI over rolling window
static float rssi_std(const BandBuf& b, float mean) {
    uint32_t n = buf_valid_count(b);
    if (n < 2) return 0.0f;
    float acc = 0.0f;
    for (uint32_t i = 0; i < n; i++) {
        float d = (float)b.obs[i].rssi - mean;
        acc += d * d;
    }
    return sqrtf(acc / (float)n);
}

// Mean of SNR over rolling window
static float snr_mean(const BandBuf& b) {
    uint32_t n = buf_valid_count(b);
    if (n == 0) return 0.0f;
    float sum = 0.0f;
    for (uint32_t i = 0; i < n; i++) {
        sum += (float)b.obs[i].snr;
    }
    return sum / (float)n;
}

// Std-dev of SNR
static float snr_std(const BandBuf& b, float mean) {
    uint32_t n = buf_valid_count(b);
    if (n < 2) return 0.0f;
    float acc = 0.0f;
    for (uint32_t i = 0; i < n; i++) {
        float d = (float)b.obs[i].snr - mean;
        acc += d * d;
    }
    return sqrtf(acc / (float)n);
}

// Inter-arrival jitter: std-dev of successive timestamp differences [ms]
static float ia_jitter(const BandBuf& b) {
    uint32_t n = buf_valid_count(b);
    if (n < 2) return 0.0f;

    // Compute differences (sorted by insertion order via modular indexing)
    float diffs[PACKET_BUF_SIZE - 1];
    uint32_t start = (b.count >= PACKET_BUF_SIZE)
                     ? (b.head % PACKET_BUF_SIZE)      // oldest slot
                     : 0u;
    uint32_t nd = 0;
    float mean_d = 0.0f;
    for (uint32_t i = 1; i < n; i++) {
        uint32_t ia = (start + i - 1) % PACKET_BUF_SIZE;
        uint32_t ib = (start + i)     % PACKET_BUF_SIZE;
        int32_t dt = (int32_t)b.obs[ib].ts_ms - (int32_t)b.obs[ia].ts_ms;
        if (dt < 0) dt = -dt;  // guard against timer rollover
        diffs[nd] = (float)dt;
        mean_d += (float)dt;
        nd++;
    }
    if (nd == 0) return 0.0f;
    mean_d /= (float)nd;
    float var = 0.0f;
    for (uint32_t i = 0; i < nd; i++) {
        float d = diffs[i] - mean_d;
        var += d * d;
    }
    return sqrtf(var / (float)nd);
}

// Occupancy: fraction of window slots that contain a packet.
// Because the buffer is filled by observed arrivals, every valid slot IS a
// packet — so occupancy is min(count, PACKET_BUF_SIZE) / PACKET_BUF_SIZE.
static float occupancy(const BandBuf& b) {
    return (float)buf_valid_count(b) / (float)PACKET_BUF_SIZE;
}

// ---------------------------------------------------------------------------
// Spectrum-proxy cumulants
//
// We treat the 128-bin normalised spectrum as a real-valued discrete signal
// x[k], k=0..N-1.  The 'cumulants' below are the sample central moments of
// x[k] and x[k]^2, matching the cumulant definitions used in l1_estimator.py:
//
//   m1   = E[x]
//   m2   = E[x^2]     (2nd raw moment)
//   C20  = m2 - m1^2  (2nd central moment — variance)  → real; imag = 0
//   C21  = E[|x|^2]   (for real signal: same as m2)
//   m4   = E[x^4]     (4th raw moment)
//   C40  = m4 - 3*m2^2                                 → kurtosis excess
//   C42  = m4 - |C20|^2 - 2*m2^2 + 2*m1^2
//          (simplified for real signal)
//
// These values are meaningful even though we don't have baseband IQ; the
// server L1 estimator uses them as complementary features alongside whatever
// it derives from SDR or spectrum data.
// ---------------------------------------------------------------------------
static void compute_spectrum_cumulants(const float* spectrum, size_t bins,
                                       float& c20_real, float& c20_imag,
                                       float& c21,
                                       float& c40_real, float& c40_imag,
                                       float& c42_out) {
    if (bins == 0) {
        c20_real = c20_imag = c21 = c40_real = c40_imag = c42_out = 0.0f;
        return;
    }

    float N = (float)bins;

    // 1st and 2nd raw moments
    float m1 = 0.0f, m2 = 0.0f, m4 = 0.0f;
    for (size_t k = 0; k < bins; k++) {
        float x = spectrum[k];
        float x2 = x * x;
        m1 += x;
        m2 += x2;
        m4 += x2 * x2;
    }
    m1 /= N;
    m2 /= N;
    m4 /= N;

    // C20 = variance
    float C20 = m2 - m1 * m1;
    c20_real = C20;
    c20_imag = 0.0f;  // real-valued signal — imaginary part is always zero

    // C21 = E[|x|^2] = m2 for real signal
    c21 = m2;

    // C40 = excess kurtosis of x  (= m4 - 3*m2^2 for zero-mean signal;
    //        we include the mean correction as per Swami & Sadler 2000)
    float C40 = m4 - 3.0f * m2 * m2;
    c40_real = C40;
    c40_imag = 0.0f;

    // C42 = m4 - |C20|^2 - 2*m2^2 + 2*m1^2  (exact for real proxy signal)
    c42_out = m4 - (C20 * C20) - 2.0f * m2 * m2 + 2.0f * m1 * m1;
}

// ---------------------------------------------------------------------------
// wfx_features_compute
// ---------------------------------------------------------------------------
bool wfx_features_compute(const float* spectrum, size_t bins,
                          WfxFeatureVec* out) {
    if (!out) return false;

    // Use band 0 (2.4 GHz WiFi) for RSSI/jitter stats; band 1 (LoRa) as
    // secondary.  The caller sets out->band_id before publishing.
    const BandBuf& wifi = s_bands[0];
    const BandBuf& lora = s_bands[1];

    // Require at least one observation in either band
    if (buf_valid_count(wifi) == 0 && buf_valid_count(lora) == 0) {
        return false;
    }

    out->sensor_id    = s_sensor_id;

    // timestamp_ms is set by the caller (loop millis()) or here as a best-effort
    // — leave it to the caller to fill after compute() returns.  Default to 0.
    out->timestamp_ms = 0;

    // Cumulants from spectrum proxy — use locals to avoid packed-field alias UB.
    {
        float c20r = 0.0f, c20i = 0.0f, c21v = 0.0f;
        float c40r = 0.0f, c40i = 0.0f, c42v = 0.0f;
        compute_spectrum_cumulants(spectrum, bins,
                                   c20r, c20i, c21v,
                                   c40r, c40i, c42v);
        out->c20_real = c20r;
        out->c20_imag = c20i;
        out->c21      = c21v;
        out->c40_real = c40r;
        out->c40_imag = c40i;
        out->c42      = c42v;
    }

    // RSSI / SNR stats — prefer the band with more observations
    const BandBuf& dominant =
        (buf_valid_count(wifi) >= buf_valid_count(lora)) ? wifi : lora;

    float rm = rssi_mean(dominant);
    out->rssi_mean  = rm;
    out->rssi_std   = rssi_std(dominant, rm);

    float sm = snr_mean(dominant);
    out->snr_mean   = sm;
    out->snr_std    = snr_std(dominant, sm);

    out->ia_jitter_ms = ia_jitter(dominant);
    out->occupancy    = occupancy(dominant);

    out->band_id    = 0;  // caller may override
    out->reserved[0] = out->reserved[1] = out->reserved[2] = 0;

    return true;
}

// ---------------------------------------------------------------------------
// wfx_features_publish
// ---------------------------------------------------------------------------
bool wfx_features_publish(const WfxFeatureVec& v) {
    // Packet layout:
    //  [0]      MSG_TYPE_WFX_FEATURES (0x10)
    //  [1..4]   sensor_id big-endian
    //  [5..60]  WfxFeatureVec (56 bytes, packed)
    //  [61..62] CRC-16 over bytes [0..60]
    //  [63]     reserved

    static uint8_t pkt[WFX_PACKET_SIZE];
    memset(pkt, 0, sizeof(pkt));

    pkt[0] = (uint8_t)MSG_TYPE_WFX_FEATURES;

    // sensor_id big-endian
    pkt[1] = (uint8_t)(v.sensor_id >> 24);
    pkt[2] = (uint8_t)(v.sensor_id >> 16);
    pkt[3] = (uint8_t)(v.sensor_id >>  8);
    pkt[4] = (uint8_t)(v.sensor_id);

    // Copy feature vector (56 bytes)
    memcpy(&pkt[WFX_HDR_SIZE], &v, sizeof(WfxFeatureVec));

    // CRC over header + payload (61 bytes)
    uint16_t crc = wfx_crc16(pkt, WFX_CRC_OFFSET);
    pkt[WFX_CRC_OFFSET]     = (uint8_t)(crc >> 8);
    pkt[WFX_CRC_OFFSET + 1] = (uint8_t)(crc);

    return mesh_publish_raw(pkt, WFX_PACKET_SIZE);
}
