// File: main.c
#include <stdio.h>
#include <math.h>
#include "pico/stdlib.h"
#include "hardware/spi.h"
#include "hardware/gpio.h"
#include "ad5940.h"

// Pin Definitions
#define PIN_MISO 0
#define PIN_CS   1
#define PIN_SCK  2
#define PIN_MOSI 3
#define PIN_RST  9
#define PIN_INT  29

#define RCAL_VALUE 100.0f

// ---------------------------------------------------------------------------
// Platform Interface Implementation (Required by AD5940 Lib)
// ---------------------------------------------------------------------------

void AD5940_CsClr(void) {
    gpio_put(PIN_CS, 0);
}

void AD5940_CsSet(void) {
    gpio_put(PIN_CS, 1);
}

void AD5940_RstClr(void) {
    gpio_put(PIN_RST, 0);
}

void AD5940_RstSet(void) {
    gpio_put(PIN_RST, 1);
}

void AD5940_Delay10us(uint32_t time) {
    sleep_us(time * 10);
}

void AD5940_ReadWriteNBytes(unsigned char *pSendBuffer, unsigned char *pRecvBuff, unsigned long length) {
    spi_write_read_blocking(spi0, pSendBuffer, pRecvBuff, length);
}

// ---------------------------------------------------------------------------
// Application Logic
// ---------------------------------------------------------------------------

void setup_pins(void) {
    // SPI0 at 1MHz
    spi_init(spi0, 1000000);
    gpio_set_function(PIN_MISO, GPIO_FUNC_SPI);
    gpio_set_function(PIN_SCK, GPIO_FUNC_SPI);
    gpio_set_function(PIN_MOSI, GPIO_FUNC_SPI);

    // CS
    gpio_init(PIN_CS);
    gpio_set_dir(PIN_CS, GPIO_OUT);
    gpio_put(PIN_CS, 1);

    // RST
    gpio_init(PIN_RST);
    gpio_set_dir(PIN_RST, GPIO_OUT);
    gpio_put(PIN_RST, 1);

    // INT
    gpio_init(PIN_INT);
    gpio_set_dir(PIN_INT, GPIO_IN);
    gpio_pull_up(PIN_INT);
}

void configure_afe(void) {
    HSDACCfg_Type hsdac_cfg;
    HSTIACfg_Type hsrtia_cfg;
    HSLoopCfg_Type hsloop_cfg;

    // Reset AFE control
    AD5940_AFECtrlS(AFECTRL_ALL, bFALSE);

    // Configure High Speed DAC
    hsdac_cfg.ExcitBufGain = EXCITBUFGAIN_2;
    hsdac_cfg.HsDacGain = HSDACGAIN_1;
    hsdac_cfg.HsDacUpdateRate = 0x1B;
    AD5940_HSDacCfgS(&hsdac_cfg);

    // Configure High Speed TIA
    hsrtia_cfg.HstiaDeRtia = HSTIADERTIA_1K;
    hsrtia_cfg.HstiaCtia = 31; // 31pF
    hsrtia_cfg.HstiaRtiaSel = HSTIARTIA_1K;
    hsrtia_cfg.HstiaBias = HSTIABIAS_1P1;
    hsrtia_cfg.DiodeClose = bFALSE;
    hsrtia_cfg.HstiaDeRload = HSTIADERLOAD_OPEN;
    AD5940_HSTIACfgS(&hsrtia_cfg);

    // Configure High Speed Loop
    hsloop_cfg.HsDacCfg = hsdac_cfg;
    hsloop_cfg.HsTiaCfg = hsrtia_cfg;
    
    hsloop_cfg.SWMatCfg.Dswitch = SWD_OPEN;
    hsloop_cfg.SWMatCfg.Pswitch = SWP_OPEN;
    hsloop_cfg.SWMatCfg.Nswitch = SWN_OPEN;
    hsloop_cfg.SWMatCfg.Tswitch = SWT_OPEN;

    hsloop_cfg.WgCfg.WgType = WGTYPE_SIN;
    hsloop_cfg.WgCfg.GainCalEn = bFALSE;
    hsloop_cfg.WgCfg.OffsetCalEn = bFALSE;
    hsloop_cfg.WgCfg.SinCfg.SinFreqWord = AD5940_WGFreqWordCal(1000.0, 16000000.0);
    hsloop_cfg.WgCfg.SinCfg.SinAmplitudeWord = 2047;
    hsloop_cfg.WgCfg.SinCfg.SinOffsetWord = 0;
    hsloop_cfg.WgCfg.SinCfg.SinPhaseWord = 0;

    AD5940_HSLoopCfgS(&hsloop_cfg);

    // Enable Power
    AD5940_AFECtrlS(AFECTRL_ADCPWR | AFECTRL_HSTIAPWR | AFECTRL_HSDACPWR | AFECTRL_HPREFPWR, bTRUE);
    
    // Enable Interrupts
    AD5940_WriteReg(REG_INTC_INTCSEL0, BITM_INTC_INTCSEL0_INTSEL14);
}

float perform_dft(void) {
    iImpCar_Type dft_res;
    fImpCar_Type f_res;
    
    AD5940_AFECtrlS(AFECTRL_ADCCNV | AFECTRL_DFT, bTRUE);
    
    // Wait for interrupt or timeout
    uint32_t timeout = 50000;
    while(gpio_get(PIN_INT) && timeout--) {
        sleep_us(10);
    }
    
    dft_res.Real = (int32_t)AD5940_ReadReg(REG_AFE_DFTREAL);
    dft_res.Image = (int32_t)AD5940_ReadReg(REG_AFE_DFTIMAG);
    
    AD5940_AFECtrlS(AFECTRL_ADCCNV | AFECTRL_DFT, bFALSE);
    AD5940_WriteReg(REG_INTC_INTCFLAG0, BITM_INTC_INTCFLAG0_FLAG14); 

    f_res.Real = (float)dft_res.Real;
    f_res.Image = (float)dft_res.Image;
    
    return AD5940_ComplexMag(&f_res);
}

void set_sw_matrix(uint32_t lp_sw, uint32_t hs_sw) {
    AD5940_WriteReg(REG_AFE_LPTIASW0, lp_sw);
    AD5940_WriteReg(REG_AFE_HSTIACON, hs_sw);
}

void measure_at_freq(float freq) {
    FreqParams_Type freq_params = AD5940_GetFreqParameters(freq);
    
    // 1. Update HSDAC Update Rate based on frequency
    uint32_t hsdac_rate = 0x1B;
    if(freq >= 80000.0f) {
        hsdac_rate = 0x07;
    }
    
    uint32_t hsdaccon = AD5940_ReadReg(REG_AFE_HSDACCON);
    hsdaccon &= ~BITM_AFE_HSDACCON_RATE;
    hsdaccon |= (hsdac_rate << BITP_AFE_HSDACCON_RATE);
    AD5940_WriteReg(REG_AFE_HSDACCON, hsdaccon);

    // 2. Update Waveform Generator Frequency
    uint32_t freq_word = AD5940_WGFreqWordCal(freq, 16000000.0f);
    AD5940_WriteReg(REG_AFE_WGFCW, freq_word);

    // 3. Update Filter Settings
    ADCFilterCfg_Type filter_cfg;
    // Initialize struct to 0 to avoid garbage
    filter_cfg.ADCSinc3Osr = freq_params.ADCSinc3Osr;
    filter_cfg.ADCSinc2Osr = freq_params.ADCSinc2Osr;
    filter_cfg.ADCAvgNum = ADCAVGNUM_16;
    filter_cfg.ADCRate = ADCRATE_800KHZ;
    filter_cfg.BpNotch = bTRUE;
    filter_cfg.BpSinc3 = bFALSE;
    filter_cfg.Sinc2NotchEnable = bTRUE;
    filter_cfg.DFTClkEnable = bTRUE;
    filter_cfg.WGClkEnable = bTRUE;
    filter_cfg.Sinc3ClkEnable = bTRUE;
    filter_cfg.Sinc2NotchClkEnable = bTRUE;
    AD5940_ADCFilterCfgS(&filter_cfg);

    DFTCfg_Type dft_cfg;
    dft_cfg.DftNum = freq_params.DftNum;
    dft_cfg.DftSrc = freq_params.DftSrc;
    dft_cfg.HanWinEn = bTRUE;
    AD5940_DFTCfgS(&dft_cfg);

    // 4. Perform Measurement
    // RCAL
    set_sw_matrix(0, SWP_RCAL0 | SWN_RCAL1); 
    sleep_ms(5); 
    float mag_cal = perform_dft();

    // Z
    set_sw_matrix(0, SWP_CE0 | SWP_RE0 | SWN_SE0 | SWT_DE0); 
    sleep_ms(5);
    float mag_z = perform_dft();

    if (mag_z > 0.01f && mag_cal > 0.01f) {
        float impedance = (mag_cal / mag_z) * RCAL_VALUE;
        printf("DATA,%.2f,%.2f,0,0,0\n", freq, impedance);
    } else {
        printf("DATA,%.2f,0,0,0,0\n", freq);
    }
}

void run_sweep(float start_freq, float end_freq, int steps) {
    float log_start = log10f(start_freq);
    float log_end = log10f(end_freq);
    float step_size = (steps > 1) ? (log_end - log_start) / (steps - 1) : 0;

    printf("START_SWEEP\n");

    for (int i = 0; i < steps; ++i) {
        float freq = powf(10.0f, log_start + (i * step_size));
        measure_at_freq(freq);
    }

    printf("END_SWEEP\n");
}

void system_init(void) {
    setup_pins();
    
    // Hardware Reset AD5940
    AD5940_RstClr();
    sleep_ms(10);
    AD5940_RstSet();
    sleep_ms(10);

    AD5940_Initialize();
    configure_afe();
}

int main() {
    stdio_init_all();
    sleep_ms(2000);
    
    system_init();

    printf("SYSTEM_READY\n");

    while (true) {
        int c = getchar_timeout_us(100);
        if (c == 'm') {
            run_sweep(100.0f, 100000.0f, 50);
        }
        if (c == 'z') {
            system_init();
            printf("RESET_DONE\n");
        }
    }
    return 0;
}