/*
 * This program implements the 
 * Sliding Windowed Infinite Fourier Transform (SWIFT)
 * and the slightly lower noise version, alpha-SWIFT
 * 
 * SWIFT is then used to build a phase-vocoder pitch shifter
 * 
 * Logan L. Grado, Matthew D. Johnson, and Theoden I. Netoff
 *  Digital Object Identifier 10.1109/MSP.2017.2718039
 *  https://ieeexplore.ieee.org/document/8026592
 * 
 * A bank of digital, complex, resonaters each with the appropriate window
 * copmputes a complete fourier transform at each input sample time.
 * Sample rate in 16 KHz.

 *  - GPIO 28 ADC input 
 * in this example, input cam from a laptop through a high-pass biased to Vdd/2
 * 
 * -- SPI for DAC
 * #define PIN_CS   5
 * define PIN_SCK  6
 * define PIN_MOSI 7
 * define SPI_PORT spi0
 *
 * 
 */

// ==========================================
// === VGA graphics libraryphase shift
// ==========================================
//#include "vga16_graphics_v3.h"
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <stdint.h>

// pico specific
#include "pico/stdlib.h"
#include "hardware/pio.h"
#include "hardware/dma.h"
#include "hardware/adc.h"
#include "hardware/spi.h"
#include "hardware/vreg.h"
#include "hardware/clocks.h"

//
// ==========================================
// === protothreads globals
// ==========================================
#include "hardware/sync.h"
#include "hardware/timer.h"
#include "pico/multicore.h"
#include "string.h"
// protothreads header
#include "pt_cornell_rp2040_v1_4.h"

// ==========================================
// === set up timer ISR  used in this pgm
// ==========================================
// === timer alarm ========================
// !! modifiying alarm zero trashes the cpu 
//        and causes LED  4 long - 4 short
// !! DO NOT USE alarm 0
// This low-level setup is ocnsiderably faster to execute
// than the hogh-level callback

#define ALARM_NUM 1
#define ALARM_IRQ TIMER0_IRQ_1
// ISR interval will be 50 uSec

volatile int alarm_period = 50 ; //62 ;
float Fs = 20000 ; //16000 ;
//
// the actual ISR
void compute_sample(void);
//
static void alarm_irq(void) {
    // mark ISR entry
    gpio_put(15,1);
    // Clear the alarm irq
    hw_clear_bits(&timer_hw->intr, 1u << ALARM_NUM);
    // arm the next interrupt
    // Write the lower 32 bits of the target time to the alarm to arm it
    timer_hw->alarm[ALARM_NUM] = timer_hw->timerawl + alarm_period ;
    //
    // The actual computation
    compute_sample();

    // mark ISR exit
    gpio_put(15,0);
}

// set up the timer alarm ISR
static void alarm_in_us(uint32_t delay_us) {
    // Enable the interrupt for our alarm (the timer outputs 4 alarm irqs)
    hw_set_bits(&timer_hw->inte, 1u << ALARM_NUM);
    // Set irq handler for alarm irq
    irq_set_exclusive_handler(ALARM_IRQ, alarm_irq);
    // Enable the alarm irq
    irq_set_enabled(ALARM_IRQ, true);
    // Enable interrupt in block and at processor
    // Alarm is only 32 bits 
    uint64_t target = timer_hw->timerawl + delay_us;
    // Write the lower 32 bits of the target time to the alarm which
    // will arm it
    timer_hw->alarm[ALARM_NUM] = (uint32_t) target;   
}

// ===========================================
// ADC setup 
// ===========================================

float adc_sample ;
float adc_mean ;
void ADC_setup(void){
  adc_init();
  adc_gpio_init(28);
  // p26 is ADC input 0 - gpio 28 is 2
  adc_select_input(2);
  // free run
  adc_run(1);
  // result is in adc_hw->result
}

// ========================================
// === spi setup 
// =======================================
//SPI configurations
#define PIN_CS   5
#define PIN_SCK  6
#define PIN_MOSI 7
#define SPI_PORT spi0

// constant to tell SPI DAC what to do
// prepend to each 12-bit sample
#define DAC_config_chan_A 0b0011000000000000
// B-channel, 1x, active
#define DAC_config_chan_B 0b1011000000000000
uint16_t DAC_data ; 
void spi_setup(void){
    // Initialize SPI channel (channel, baud rate set to 20MHz)
    // connected to spi DAC
    spi_init(SPI_PORT, 20000000) ;
    // Format (channel, data bits per transfer, polarity, phase, order)
    spi_set_format(SPI_PORT, 16, 0, 0, 0);
    // Map SPI signals to GPIO ports
    //gpio_set_function(PIN_MISO, GPIO_FUNC_SPI);
    gpio_set_function(PIN_SCK, GPIO_FUNC_SPI);
    gpio_set_function(PIN_MOSI, GPIO_FUNC_SPI);
    gpio_set_function(PIN_CS, GPIO_FUNC_SPI) ;
}
//

// ===========================================
// dSP definitions 
// ===========================================
// The Sliding Windowed Infinite Fourier Transform -- with alpha-SWIFT
// Logan L. Grado, Matthew D. Johnson, and Theoden I. Netoff
// Digital Object Identifier 10.1109/MSP.2017.2718039
//file:///home/bruce/Downloads/10078055.pdf

// === SWIFT setup
// can be any number, limited by the speed of the cpu
// set to number desired frequencies in the filter bank
#define N_FREQ  48 //201   

// A list of the frequencys to transform
// need not be uniform
float F[N_FREQ][2] ;

// the complex spectral estimate
// frequencies are normailsed and run 0 to pi
// at 10khz sample rate that is 0 to 5kz
float Xr[N_FREQ][2], Xi[N_FREQ][2] ;
float gain[N_FREQ][2] ;
float dk_rate[N_FREQ][2] ;
float rise_rate[N_FREQ][2] ;
float rise_time ;

// additive components of the alpha-SWIFT
float slow_Xr[N_FREQ][2], slow_Xi[N_FREQ][2];
float fast_Xr[N_FREQ][2], fast_Xi[N_FREQ][2] ;

// and the mag of X
float Fmag[N_FREQ][2], lp_Fmag[N_FREQ][2];

// the per-sample complex phase shift at each frequecy
// omega = 2*pi*F/Fs so OHM = cos(omega) + j*sin(omega)
float slow_OHMr[N_FREQ][2], slow_OHMi[N_FREQ][2] ;
float fast_OHMr[N_FREQ][2], fast_OHMi[N_FREQ][2] ;

//  dds to output filter bank content for shifterr
#define twoPi 6.283185f
#define pow2_32 4294967296.0f
// resynthesis
unsigned int  accum1;
unsigned int inc1, inc_temp;
float input_freq = 100.0f ;
float drive_table[256] ;
float drive_sample ;
float phase_diff, lp_phase_diff, last_phase ;
int last_pitch_bin, pitch_bin_valid ;
int use_pitch = false ;

// audio input
float input_sample ;

// dac ouptut signal
float test_out ;

//====================================
// === magnitude approx good to about +/-2%
// see https://en.wikipedia.org/wiki/Alpha_max_plus_beta_min_algorithm
#define min(a,b) (((a)<(b))?(a):(b))
#define max(a,b) (((a)>(b))?(a):(b))
//====================================
float magnitude(float real, float img){
  float mmax, mmin ;
  float c1 = 0.89820f ;
  float c2 = 0.48597f ;
  
  mmin = min(fabsf(real), fabsf(img)) ; 
  mmax = max(fabsf(real), fabsf(img)) ;
  // 
  return max(mmax, ((mmax*c1) + (mmin*c2) )); 
}


// ==================================================
// === sdetup vocoder -- 
// ==================================================
static PT_THREAD (protothread_setup(struct pt *pt)) {
    PT_BEGIN(pt);

      static float dds_last_phase ;
    // =====================================
    // =====================================
    // This is where you set up the desired frequency
    // analysis that you want
    // You MUST SET UP each seqarate resonator
    // --set up F, gain, decay and rise time  arrays
    
    for(int i=0; i<N_FREQ; i++){
      // choose to make a log-spaced spectrum
      // from 90 to 4300 Hz
      if (i==0) F[0][0] = 90.0 ; //100.0f ;
      else  F[i][0] = F[i-1][0] * 1.085 ; // f
      // copy to output array
      F[i][1] = F[i][0] ;

      // set the gain of each frequency
      // this seems to equalize fairly well
      gain[i][0] = (F[i][0]) / 90.0f;
      gain[i][1] = gain[i][0] ;

      // set the decay time constant for each freq
      // sert to 8 times the period of F[i]
      #define dk_ratio 20.0f // 3.0
      dk_rate[i][0] = exp(-1.0f/(dk_ratio * Fs /F[i][0])) ; // samples
      dk_rate[i][1] = dk_rate[i][0] ;

      // set the rise rate consrtant for each freq
      // bigger lowers noise, but attenuates higher freq
      // 0.1 is very small, 20.0 is big
      //#define rise_time 20.0f
      rise_time = 0.2f * dk_ratio * Fs /F[i][0] ;
      rise_rate[i][0] = exp(-1.0f/rise_time) ; // samples
      rise_rate[i][1] = rise_rate[i][0] ;
      // =====================================
      // =====================================
      
      // compute complex frequenc for phasor
      // omega = 2*pi*F/Fs so OHM = cos(omega) + j*sin(omega)
      // and premultiply by rate constants
      float omega = twoPi * F[i][0]/Fs ;

      slow_OHMr[i][0] = dk_rate[i][0] * cos(omega) ;
      slow_OHMr[i][1] = slow_OHMr[i][0] ;
      //
      slow_OHMi[i][0] = dk_rate[i][0] * sin(omega); 
      slow_OHMi[i][1] = slow_OHMi[i][0] ;
      //
      fast_OHMr[i][0] = rise_rate[i][0] * cos(omega) ;
      fast_OHMr[i][1] = fast_OHMr[i][0] ;
      //
      fast_OHMi[i][0] = rise_rate[i][0] * sin(omega);  
      fast_OHMi[i][1] = fast_OHMi[i][0] ;    
    }

    // thin inc will be immediately changed by the isr
     inc1 =  (unsigned int) input_freq * pow2_32/ Fs ;

     // ======================================

    ////
    while(true) {
      // just spin
      PT_YIELD_usec(1000) ;
      //
     
   }
   PT_END(pt);
} // sdetup thread


// ==================================================
// === user's serial input thread on core 1
// ==================================================
// serial_read an serial_write do not block any thread
// except this one
static PT_THREAD (protothread_serial(struct pt *pt))
{
    PT_BEGIN(pt);
      
      static char *cmd, *arg1, *arg2, *arg3, *arg4;

      while(1) {
        // prompt the user for input
        sprintf(pt_serial_out_buffer, "cmd> ") ;
        // spawn a thread to do the non-blocking serial write
        // to print the pt_serial_out_buffer string
        // this output routine YIELDs between each character printed.
        serial_write ; 
        // wait for user
        serial_read ;
        // tokenize the input string
        //  --- NOTE that strtok is not re-entrant!
        //  ---   Use the reentrant, thread-safe, version if you need to use it
        //   ---  in moe than one thread!
        cmd = strtok(pt_serial_in_buffer, "  ");
        arg1 = strtok(NULL, "  ");
        arg2 = strtok(NULL, "  ");
        arg3 = strtok(NULL, "  ");
        arg4 = strtok(NULL, "  ");

        //wait
        sleep_ms(10);
        
        // 
        if(strcmp(cmd,"vary")==0){
          //sscanf(arg1,"%f", &use_pitch);  
          use_pitch = true ; 
        }

        else if(strcmp(cmd,"fix")==0){
          sscanf(arg1,"%f", &input_freq);  
          use_pitch = false ; 
          inc1 =  (unsigned int) input_freq * pow2_32/ Fs ;
        }

        //
        else{
          printf("Huh?\n\r") ;
        }
        //
        // NEVER exit while
      } // END WHILE(1)
  PT_END(pt);
} // serial thread

// ==================================================
// === dsp ISR -- 
// ==================================================
// 
void compute_sample(void){

    // convert ADC to float
    adc_sample = (float) (adc_hw->result) ;
    
    // low pass to get mean
    // then subtract the mean (highpass)
    // high pass at around 0.1 second, so one Hz
    adc_mean = 0.98f * adc_mean + 0.02f * adc_sample ;
    adc_sample -= adc_mean ;
    input_sample = adc_sample ;

   // zero the output sum
    test_out = 0.0f ;

    // outer llop over input or output filter bank
    for (int j=0; j<2; j++){
      for(int i=2; i<N_FREQ; i++){   //N_FREQ
        float temp ;
        // the slow arm of the swift
        temp = (slow_Xr[i][j] * slow_OHMr[i][j] - slow_Xi[i][j] * slow_OHMi[i][j]) + 
          ((j)? drive_sample : input_sample) ; //-
        slow_Xi[i][j] = (slow_Xr[i][j] * slow_OHMi[i][j] + slow_Xi[i][j] * slow_OHMr[i][j]) ;//-
        slow_Xr[i][j] = temp ;
        //
        // the fast arm of the swift
        temp = (fast_Xr[i][j] * fast_OHMr[i][j] - fast_Xi[i][j] * fast_OHMi[i][j]) + 
          ((j)? drive_sample : input_sample) ;  ; //-
        fast_Xi[i][j] = (fast_Xr[i][j] * fast_OHMi[i][j] + fast_Xi[i][j] * fast_OHMr[i][j]) ;//-
        fast_Xr[i][j] = temp ;
        //
        // combine the two to make X
        Xr[i][j] = slow_Xr[i][j] - fast_Xr[i][j];
        Xi[i][j] = slow_Xi[i][j] - fast_Xi[i][j];     
      } 
    } // end the per-filter loop

    // estimate the exciting freauency
    int pitch_bin = 0 ;
    float max_bin_mag = 0.0f ;
    
    for (int i=5; i<10; i++){
      if (lp_Fmag[i][0] > max_bin_mag){
        max_bin_mag = lp_Fmag[i][0] ;
        pitch_bin = i ;
      }
    }
    if (pitch_bin == last_pitch_bin){
      pitch_bin_valid = true ;
    }
    else {
      pitch_bin_valid = false ;
      last_pitch_bin = pitch_bin ;
    }
    //
    // last_pitch_bin, pitch_bin_valid
    float temp_phase = fmodf( (atan2f(Xi[pitch_bin][0], Xr[pitch_bin][0]) + twoPi), twoPi)  ;
    // actual phase diff of resonator
    phase_diff =  (temp_phase >= last_phase)? temp_phase - last_phase:
                                                            (twoPi - last_phase) + temp_phase ;
    lp_phase_diff = 0.99f * lp_phase_diff + 0.01f * phase_diff ;
    // copy for next pass
    last_phase = temp_phase ;
    // convert phase diff to increment units for next pass
    //inc_temp =  (unsigned int) ( lp_phase_diff  * pow2_32 / twoPi);
    // update frequency at zero phase
    if((temp_phase < 0.1f) && pitch_bin_valid && use_pitch) {
      inc1 =  (unsigned int) ( 0.5f* lp_phase_diff  * pow2_32 / twoPi);
    }

    // now the drive for the output
    accum1 += inc1 ;
    drive_sample = drive_table[accum1>>24] ; 

    // get the input filter magnitude and form
    // input mag times output waveform
    for(int i=2; i<N_FREQ; i++){ 
        lp_Fmag[i][0] = 0.99f * lp_Fmag[i][0] + 0.01f * fabsf(Xr[i][0]) * gain[i][0];
        test_out += lp_Fmag[i][0] * Xr[i][1] * gain[i][1]; // gain 
    }

    // scale may need to change depending on source amp
    // offset converts signed to offset binary
    //DAC_data = DAC_config_chan_A | (((int)(input_sample*0.1f) + 2048) & 0xfff)   ; 
    //spi0_hw->dr = DAC_data ;
    DAC_data = DAC_config_chan_B | (((int)(test_out*2.0e-7f) + 2048) & 0xfff)   ; 
    // NOTE === nonblocking SPI write ===
    spi0_hw->dr = DAC_data ;
    //blocking out: spi_write16_blocking(SPI_PORT, &DAC_data, 1) ; 
} // end fo ISR

// ========================================
// === core 1 main -- started in main below
// ========================================
void core1_main(){ 
  //
  //  === add threads  ====================
  // for core 1
  pt_add_thread(protothread_serial) ;
  pt_add_thread(protothread_setup);
  //
  // === initalize the scheduler ==========
  pt_schedule_start ;
  // NEVER exits
  // ======================================
}

// ========================================
// === core 0 main
// ========================================
int main(){
  // set the clock
  vreg_set_voltage(VREG_VOLTAGE_1_30 );
  //enum vreg_voltage
  // set the clock
  set_sys_clock_khz (300000, false);

  // start the serial i/o
  stdio_init_all() ;
  // announce the threader version on system reset
  printf("\n\rProtothreads RP2040/2350 v1.4 two-core\n\r");

  // the spi for DAC
  spi_setup() ;

  // === congure the ADC =========
  ADC_setup() ;
  
  // dds table for testing
  for (int i=0; i<256; i++) {
    // output filter drive table 
    drive_table[i] = 900.0f*((i<1)? 1.0f : (i<2)?-1.0f : 0.0f) + 200.0f *((float)rand()/RAND_MAX - 0.5f);
    //drive_table[i] = 900.0f*((i<10)? (float)i * 10.0f :  (i-10) * ) 
    //    + 200.0f *((float)rand()/RAND_MAX - 0.5f);
  }

  //
  // fire off interrupt
  alarm_in_us(alarm_period);

  gpio_init(15) ;	
  gpio_set_dir(15, GPIO_OUT) ;
  gpio_init(14) ;	
  gpio_set_dir(14, GPIO_OUT) ;

  // start core 1 threads
  multicore_reset_core1();
  multicore_launch_core1(&core1_main);


  // === config threads ========================
  // for core 0
  //
  //pt_add_thread(protothread_graphics);
  //pt_add_thread(protothread_toggle25);
  //
  // === initalize the scheduler ===============
  pt_schedule_start ;
  // NEVER exits
  // ===========================================
} // end main