#include <sys/cdefs.h>
// Sample from the ADC continuously at a particular sample rate
// and then compute an FFT over the data
//
// much of this code is from pico-examples/adc/dma_capture/dma_capture.c
// the rest is written by Alex Wulff (www.AlexWulff.com)

#include <stdio.h>
#include <math.h>

#include "pico/stdlib.h"
#include "hardware/adc.h"
#include "hardware/dma.h"
#include "pico/multicore.h"
#include "kiss_fftr.h"

#include <stdio.h>
#include <stdlib.h>

#include "pico/stdlib.h"
#include "hardware/pio.h"
#include "hardware/clocks.h"
#include "ws2812.pio.h"

#define IS_RGBW false
#define NUM_PIXELS 150

#ifdef PICO_DEFAULT_WS2812_PIN
#define WS2812_PIN PICO_DEFAULT_WS2812_PIN
#else
// default to pin 2 if the board doesn't have a default WS2812 pin defined
#define WS2812_PIN 6
#endif

typedef struct {
    double r;       // a fraction between 0 and 1
    double g;       // a fraction between 0 and 1
    double b;       // a fraction between 0 and 1
} rgb;

static inline void put_pixel(uint32_t pixel_grb) {
    pio_sm_put_blocking(pio0, 0, pixel_grb << 8u);
}

// set this to determine sample rate
// 0     = 500,000 Hz
// 960   = 50,000 Hz
// 9600  = 5,000 Hz
#define CLOCK_DIV 960
#define FSAMP 50000

// Channel 0 is GPIO26
#define CAPTURE_CHANNEL 0
#define LED_PIN 25

// BE CAREFUL: anything over about 9000 here will cause things
// to silently break. The code will compile and upload, but due
// to memory issues nothing will work properly
#define NSAMP 2048 //4096

// globals
dma_channel_config cfg;
uint dma_chan;
float freqs[NSAMP];
float freqsInLog[NSAMP];
float power[NSAMP / 2];
float powerInLog[NUM_PIXELS];

uint8_t cap_buf[NSAMP];
uint8_t cap_res[NSAMP];

uint64_t lastColorChange;

void setup();

void sample(uint8_t *capture_buf);

void startSample();

rgb hsv2rgb(double h, double s, double v);

int main() {


    kiss_fft_scalar fft_in[NSAMP]; // kiss_fft_scalar is a float
    kiss_fft_cpx fft_out[NSAMP];
    kiss_fftr_cfg cfg = kiss_fftr_alloc(NSAMP, false, 0, 0);


    // setup ports and outputs
    setup();

    multicore_reset_core1();
    multicore_launch_core1(startSample);
    printf("Started\n");
    while (1) {
        //printf("Loop\n");
        // get NSAMP samples at FSAMP
        //uint64_t sampleStart = to_us_since_boot(get_absolute_time());
        //sample(cap_buf);
        //printf("Sample time: %llu \n", to_us_since_boot(get_absolute_time())-sampleStart);
        // fill fourier transform input while subtracting DC component
        uint64_t sum = 0;
        for (int i = 0; i < NSAMP; i++) { sum += cap_res[i]; }
        float avg = (float) sum / NSAMP;
        for (int i = 0; i < NSAMP; i++) { fft_in[i] = (float) cap_res[i] - avg; }

        // compute fast fourier transform
        kiss_fftr(cfg, fft_in, fft_out);

        // compute power and calculate max freq component

        // any frequency bin over NSAMP/2 is aliased (nyquist sampling theorum)
        for (int i = 0; i < NSAMP / 2; i++) {
            power[i] = fft_out[i].r * fft_out[i].r + fft_out[i].i * fft_out[i].i;
        }

        float f_max = FSAMP;
        float f_res = f_max / NSAMP;
        for (int i = 0; i < NUM_PIXELS; i++) {
            float lowFreq = freqsInLog[i];//freqs[i]; //freqsInLog[i];
            float highFreq;
            if (i != NUM_PIXELS - 1) {
                highFreq = freqsInLog[i + 1];//freqs[i+1]; //freqsInLog[i + 1];
            } else {
                highFreq = FSAMP / 2;
            }
            int lowInd = lowFreq / f_res;
            int highInd = highFreq / f_res;

            float totalPower = 0;
            for (int j = lowInd; j < highInd + 1; j++) {
                totalPower += power[j];
            }
            //printf("lowInd = %d, highInd = %d, lowFreq = %f, highFreq = %f, freq[lowInd] = %f, freq[highInd] = %f\n", lowInd, highInd, lowFreq, highFreq, freqs[lowInd], freqs[highInd]);
            float div_power = totalPower / (highInd + 1 - lowInd);
            powerInLog[i] = 20 * log(fmax(div_power/100000, 1));//div_power/1000.0; //20 * log(fmax(div_power - 1000000, 1));
        }
        //for (int i = 0; i < NUM_PIXELS - 1; i++) {
        //    printf("Power for freq %f to %f = %f (%f)\n", freqsInLog[i], freqsInLog[i + 1], powerInLog[i], power[i]);
        //}
        //uint64_t timmmme = to_us_since_boot(get_absolute_time()) - lastColorChange;
        //printf("%llu \n", timmmme);
        if (to_us_since_boot(get_absolute_time()) - lastColorChange < 300) {
            sleep_us(300 - (to_us_since_boot(get_absolute_time()) - lastColorChange));
        }

        for (int i = 0; i < NUM_PIXELS; i++) {
            //uint32_t value = (uint32_t)(fmin(255, powerInLog[i]/400.0 * 255));
            rgb color = hsv2rgb(fmin(359, powerInLog[i]/200.0 * 359), 1, fmin(0.5, powerInLog[i]/100.0));
            //printf("power = %f, color = %f, %f, %f\n", powerInLog[i]/200.0, color.r, color.g, color.b);
            uint32_t value = ((uint32_t)(fmin(color.r, 1) * 255) << 8) | ((uint32_t)(fmin(color.g, 1) * 255) << 16) | (uint32_t)(fmin(color.b, 1) * 255);
            //printf("b = %f\n", color.b*255);
            //printf("Color = %lu, pixel = %d\n", value, i);
            put_pixel(value);
            put_pixel(value);
        }
        lastColorChange = to_us_since_boot(get_absolute_time());
        //sleep_ms(50);
    }

    // should never get here
    kiss_fft_free(cfg);
}

rgb hsv2rgb(double h, double s, double v) {
    double      hh, p, q, t, ff;
    long        i;
    rgb        out;

    if(s <= 0.0) {       // < is bogus, just shuts up warnings
        out.r = v;
        out.g = v;
        out.b = v;
        return out;
    }
    hh = h;
    if(hh >= 360.0) hh = 0.0;
    hh /= 60.0;
    i = (long)hh;
    ff = hh - i;
    p = v * (1.0 - s);
    q = v * (1.0 - (s * ff));
    t = v * (1.0 - (s * (1.0 - ff)));

    switch(i) {
        case 0:
            out.r = v;
            out.g = t;
            out.b = p;
            break;
        case 1:
            out.r = q;
            out.g = v;
            out.b = p;
            break;
        case 2:
            out.r = p;
            out.g = v;
            out.b = t;
            break;

        case 3:
            out.r = p;
            out.g = q;
            out.b = v;
            break;
        case 4:
            out.r = t;
            out.g = p;
            out.b = v;
            break;
        case 5:
        default:
            out.r = v;
            out.g = p;
            out.b = q;
            break;
    }
    return out;
}

void sample(uint8_t *capture_buf) {
    adc_fifo_drain();
    adc_run(false);

    dma_channel_configure(dma_chan, &cfg,
                          capture_buf,    // dst
                          &adc_hw->fifo,  // src
                          NSAMP,          // transfer count
                          true            // start immediately
    );

    gpio_put(LED_PIN, 1);
    adc_run(true);
    dma_channel_wait_for_finish_blocking(dma_chan);
    gpio_put(LED_PIN, 0);
}

void startSample() {
    while(1) {
        adc_fifo_drain();
        adc_run(false);

        dma_channel_configure(dma_chan, &cfg,
                              cap_buf,    // dst
                              &adc_hw->fifo,  // src
                              NSAMP,          // transfer count
                              true            // start immediately
        );

        gpio_put(LED_PIN, 1);
        adc_run(true);
        dma_channel_wait_for_finish_blocking(dma_chan);
        gpio_put(LED_PIN, 0);

        memcpy(cap_res, cap_buf, sizeof cap_res);
    }
}

void setup() {
    stdio_init_all();

    // todo get free sm
    PIO pio = pio0;
    int sm = 0;
    uint offset = pio_add_program(pio, &ws2812_program);
    ws2812_program_init(pio, sm, offset, WS2812_PIN, 800000, IS_RGBW);

    //gpio_init(LED_PIN);
    //gpio_set_dir(LED_PIN, GPIO_OUT);

    adc_gpio_init(26 + CAPTURE_CHANNEL);

    adc_init();
    adc_select_input(CAPTURE_CHANNEL);
    adc_fifo_setup(
            true,    // Write each completed conversion to the sample FIFO
            true,    // Enable DMA data request (DREQ)
            1,       // DREQ (and IRQ) asserted when at least 1 sample present
            false,   // We won't see the ERR bit because of 8 bit reads; disable.
            true     // Shift each sample to 8 bits when pushing to FIFO
    );

    // set sample rate
    adc_set_clkdiv(CLOCK_DIV);

    sleep_ms(1000);
    // Set up the DMA to start transferring data as soon as it appears in FIFO
    uint dma_chan = dma_claim_unused_channel(true);
    cfg = dma_channel_get_default_config(dma_chan);

    // Reading from constant address, writing to incrementing byte addresses
    channel_config_set_transfer_data_size(&cfg, DMA_SIZE_8);
    channel_config_set_read_increment(&cfg, false);
    channel_config_set_write_increment(&cfg, true);

    // Pace transfers based on availability of ADC samples
    channel_config_set_dreq(&cfg, DREQ_ADC);

    printf("Starting\n");
    // calculate frequencies of each bin
    float f_max = FSAMP;
    float f_res = f_max / NSAMP;
    for (int i = 0; i < NSAMP; i++) {
        freqs[i] = f_res * i;
    }
    // Or 20kHz?
    float max_power = log(FSAMP / 2.0);
    float min_power = log(80);
    for (int i = 0; i < NUM_PIXELS; i++) {
        freqsInLog[i] = exp(i / (float) NUM_PIXELS * (max_power - min_power) + min_power);
        printf("freqsInLog[%d] = %f.1\n", i, freqsInLog[i]);
    }
}