#include <unistd.h>
#include <stdbool.h>
#include <math.h>

#include <arm_math.h>

#include "freq_meas.h"
#include "sr_global.h"
#include "dsss_demod.h"
#include "simulation.h"

#include "generated/dsss_gold_code.h"
#include "generated/dsss_butter_filter.h"

/* Generated CWT wavelet LUT */
extern const float * const dsss_cwt_wavelet_table;

struct iir_biquad cwt_filter_bq[DSSS_FILTER_CLEN] = {DSSS_FILTER_COEFF};

float gold_correlate_step(const size_t ncode, const float a[DSSS_CORRELATION_LENGTH], size_t offx, bool debug);
float cwt_convolve_step(const float v[DSSS_WAVELET_LUT_SIZE], size_t offx);
float run_iir(const float x, const int order, const struct iir_biquad q[order], struct iir_biquad_state st[order]);
float run_biquad(float x, const struct iir_biquad *const q, struct iir_biquad_state *const restrict st);
void debug_print_vector(const char *name, size_t len, const float *data, size_t stride, bool index, bool debug);

#ifdef SIMULATION
void debug_print_vector(const char *name, size_t len, const float *data, size_t stride, bool index, bool debug) {
    if (!debug)
        return;

    if (index) {
        DEBUG_PRINTN("    %16s  [", "");
        for (size_t i=0; i<len; i++)
            DEBUG_PRINTN("%8d  ", i);
        DEBUG_PRINTN("]\n");
    }

    DEBUG_PRINTN("    %16s: [", name);
    for (size_t i=0; i<len; i++)
        DEBUG_PRINTN("%8.5f, ", data[i*stride]);
    DEBUG_PRINTN("]\n");
}
#else 
void debug_print_vector(const char *name, size_t len, const float *data, size_t stride, bool index, bool debug) {}
#endif

#ifdef SIMULATION
void dsss_demod_step(struct dsss_demod_state *st, float new_value, size_t sim_pos, int record_channel) {
    bool debug = (record_channel == -1)
        && (sim_pos > 1000)
        && (sim_pos % DSSS_CORRELATION_LENGTH == DSSS_CORRELATION_LENGTH-1);

    if (debug) DEBUG_PRINT("Iteration %zd: signal=%f", sim_pos, new_value);
#else
void dsss_demod_step(struct dsss_demod_state *st, float new_value) {
#endif

    //const float peak_group_threshold = 0.05 * DSSS_CORRELATION_LENGTH;
    //const float hole_patching_threshold = 0.01 * DSSS_CORRELATION_LENGTH;

    st->signal[st->signal_wpos] = new_value;
    st->signal_wpos = (st->signal_wpos + 1) % ARRAY_LENGTH(st->signal);

    /* use new, incremented wpos for gold_correlate_step as first element of old data in ring buffer */
    for (size_t i=0; i<DSSS_GOLD_CODE_COUNT; i++)
        st->correlation[i][st->correlation_wpos] = gold_correlate_step(i, st->signal, st->signal_wpos, false);

    debug_print_vector("correlation",
            DSSS_GOLD_CODE_COUNT, &st->correlation[0][st->correlation_wpos], DSSS_WAVELET_LUT_SIZE, true, debug);

    st->correlation_wpos = (st->correlation_wpos + 1) % ARRAY_LENGTH(st->correlation[0]);

    float cwt[DSSS_GOLD_CODE_COUNT];
    for (size_t i=0; i<DSSS_GOLD_CODE_COUNT; i++)
        cwt[i] = cwt_convolve_step(st->correlation[i], st->correlation_wpos);

    debug_print_vector("cwt", DSSS_GOLD_CODE_COUNT, cwt, 1, false, debug);

    float avg[DSSS_GOLD_CODE_COUNT];
    for (size_t i=0; i<DSSS_GOLD_CODE_COUNT; i++)
        avg[i] = run_iir(fabs(cwt[i]), ARRAY_LENGTH(cwt_filter_bq), cwt_filter_bq, st->cwt_filter[i].st);

    debug_print_vector("avg", DSSS_GOLD_CODE_COUNT, avg, 1, false, debug);

    if (record_channel != -1)
        DEBUG_PRINTN("%f, %f, %f\n",
                st->correlation[record_channel][st->correlation_wpos], cwt[record_channel], avg[record_channel]);

    float max_val = st->group.max;
    int max_ch = st->group.max_ch;
    int max_idx = st->group.max_idx + 1;
    bool found = false;
    for (size_t i=0; i<DSSS_GOLD_CODE_COUNT; i++) {
        float val = cwt[i] / avg[i];

        if (fabs(val) > DSSS_THESHOLD_FACTOR)
            found = true;

        if (fabs(val) > fabs(max_val)) {
            max_val = val;
            max_ch = i;
            max_idx = st->group.len;
        }
    }

    if (found) {
        /* Continue ongoing group */
        st->group.len++;
        st->group.max = max_val;
        st->group.max_ch = max_ch;
        st->group.max_idx = max_idx;
        return;
    }
    
    if (st->group.len == 0)
        /* We're between groups */
        return;

    /* A group ended. Process result. */
    if (record_channel == -1)
        DEBUG_PRINT("GROUP FOUND: %8d len=%3d max=%f ch=%d offx=%d",
                sim_pos, st->group.len, st->group.max, st->group.max_ch, st->group.max_idx);

    /* reset grouping state */
    st->group.len = 0;
    st->group.max_idx = 0;
    st->group.max_ch = 0;
    st->group.max = 0.0f;
}

float run_iir(const float x, const int order, const struct iir_biquad q[order], struct iir_biquad_state st[order]) {
    float intermediate = x;
    for (int i=0; i<(order+1)/2; i++)
        intermediate = run_biquad(intermediate, &q[i], &st[i]);
    return intermediate;
}

float run_biquad(float x, const struct iir_biquad *const q, struct iir_biquad_state *const restrict st) {
    /* direct form 2, see https://en.wikipedia.org/wiki/Digital_biquad_filter */
    float intermediate = x + st->reg[0] * -q->a[0] + st->reg[1] * -q->a[1];
    float out = intermediate * q->b[0] + st->reg[0] * q->b[1] + st->reg[1] * q->b[2];
    st->reg[1] = st->reg[0];
    st->reg[0] = intermediate;
    return out;
}

float cwt_convolve_step(const float v[DSSS_WAVELET_LUT_SIZE], size_t offx) {
    float sum = 0.0f;
    for (ssize_t j=0; j<DSSS_WAVELET_LUT_SIZE; j++) {
        /* Our wavelet is symmetric so convolution and correlation are identical. Use correlation here for ease of
         * implementation */
        sum += v[(offx + j) % DSSS_WAVELET_LUT_SIZE] * dsss_cwt_wavelet_table[j];
        //DEBUG_PRINT("        j=%d v=%f w=%f", j, v[(offx + j) % DSSS_WAVELET_LUT_SIZE], dsss_cwt_wavelet_table[j]);
    }
    return sum;
}

/* Compute last element of correlation for input [a] and hard-coded gold sequences.
 *
 * This is intened to be used once for each new incoming sample in [a]. It expects [a] to be of length
 * [dsss_correlation_length] and produces the one sample where both the reference sequence and the input fully overlap.
 * This is equivalent to "valid" mode in numpy's terminology[0].
 *
 * [0] https://docs.scipy.org/doc/numpy/reference/generated/numpy.correlate.html
 */
float gold_correlate_step(const size_t ncode, const float a[DSSS_CORRELATION_LENGTH], size_t offx, bool debug) {

    float acc_outer = 0.0f;
    uint8_t table_byte = 0;
    if (debug) DEBUG_PRINTN("Correlate n=%d: ", ncode);
    for (size_t i=0; i<DSSS_GOLD_CODE_LENGTH; i++) {

        if ((i&7) == 0) {
            table_byte = dsss_gold_code_table[ncode][i>>3]; /* Fetch sequence table item */
            if (debug) DEBUG_PRINTN("|");
        }
        int bv = table_byte & (0x80>>(i&7)); /* Extract bit */
        bv = !!bv*2 - 1; /* Map 0, 1 -> -1, 1 */
        if (debug) DEBUG_PRINTN("%s%d\033[0m", bv == 1 ? "\033[92m" : "\033[91m", (bv+1)/2);

        float acc_inner = 0.0f;
        for (size_t j=0; j<DSSS_DECIMATION; j++)
            acc_inner += a[(offx + i*DSSS_DECIMATION + j) % DSSS_CORRELATION_LENGTH]; /* Multiply item */
        //if (debug) DEBUG_PRINTN("%.2f ", acc_inner);
        acc_outer += acc_inner * bv;
    }
    if (debug) DEBUG_PRINTN("\n");
    return acc_outer / DSSS_CORRELATION_LENGTH;
}