// File: EIS/main.c
#include <stdio.h>
#include <string.h>
#include <stdlib.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_RESET       9
#define PIN_AD_INTERRUPT 29 

#define SPI_PORT        spi0
#define SPI_BAUDRATE_HIGH 500000 

// --------------------------------------------------------------------------
// Switch Matrix Bit Definitions
// --------------------------------------------------------------------------
#ifndef SWT_SE0
#define SWT_SE0         (1 << 6)   // T7
#endif
#ifndef SWT_T9
#define SWT_T9          (1 << 8)   // T9
#endif
#ifndef SWT_TR1
#define SWT_TR1         (1 << 11)  // TR1
#endif

// DSWFULLCON
#ifndef SWD_AIN2
#define SWD_AIN2        (1 << 2)   // D3
#endif
#ifndef SWD_CE0
#define SWD_CE0         (1 << 4)   // D5
#endif

// TSWFULLCON
#ifndef SWT_AIN3
#define SWT_AIN3        (1 << 3)   // T4
#endif

// --------------------------------------------------------------------------
// AD5940 Library Interface
// --------------------------------------------------------------------------

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

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

void AD5940_RstClr(void) { gpio_put(PIN_RESET, 0); }
void AD5940_RstSet(void) { gpio_put(PIN_RESET, 1); }
void AD5940_Delay10us(uint32_t time) { sleep_us(time * 10); }

uint32_t AD5940_GetMCUIntFlag(void) {
    uint32_t flag = AD5940_ReadReg(REG_INTC_INTCFLAG0);
    return (flag != 0);
}

void AD5940_ReadWriteNBytes(unsigned char *pSendBuffer, unsigned char *pRecvBuffer, unsigned long length) {
    if (pRecvBuffer == NULL) {
        if (pSendBuffer != NULL) {
            spi_write_blocking(SPI_PORT, pSendBuffer, length);
        }
    } else {
        if (pSendBuffer == NULL) {
            uint8_t dummy = 0x00;
            for (unsigned long i = 0; i < length; i++) {
                spi_write_read_blocking(SPI_PORT, &dummy, &pRecvBuffer[i], 1);
            }
        } else {
            spi_write_read_blocking(SPI_PORT, pSendBuffer, pRecvBuffer, length);
        }
    }
}

// --------------------------------------------------------------------------
// Helpers
// --------------------------------------------------------------------------

void AD5940_PulseReset(void) {
    AD5940_RstClr();
    sleep_ms(10);
    AD5940_RstSet();
    sleep_ms(50); // Increased wait for boot stability
}

// Helper to map Register Macro to Actual Sample Count for calculations
uint32_t GetDftSamples(uint32_t dft_num_macro) {
    if (dft_num_macro == DFTNUM_256) return 256;
    if (dft_num_macro == DFTNUM_512) return 512;
    if (dft_num_macro == DFTNUM_1024) return 1024;
    if (dft_num_macro == DFTNUM_2048) return 2048;
    if (dft_num_macro == DFTNUM_4096) return 4096;
    if (dft_num_macro == DFTNUM_8192) return 8192;
    if (dft_num_macro == DFTNUM_16384) return 16384;
    return 2048; // Safe default
}

// Structure to hold calculated configuration
typedef struct {
    uint32_t DftNum;
    uint32_t DftSrc;
    uint32_t Sinc3Osr;
    uint32_t Sinc2Osr;
    float    SampleRate; // Estimated
} FilterConfig_Type;

// Calculates optimal DFT and Filter settings based on frequency
// Optimized for High Speed Loop (HFOSC 16MHz) only
void AD5940_GetFilterCfg(float freq, FilterConfig_Type *pCfg) {
    // Strategy: 
    // For > 50Hz, use SINC3 (Fast, ~200kHz)
    // For <= 50Hz, use SINC2 path (Slower, ~9kHz) to capture full waves
    
    if (freq >= 50.0f) {
        pCfg->DftSrc = DFTSRC_SINC3;
        pCfg->Sinc3Osr = ADCSINC3OSR_4;
        pCfg->Sinc2Osr = ADCSINC2OSR_22; // Don't care for SINC3 src
        pCfg->SampleRate = 200000.0f; // 800k / 4

        // Adjust DFT size for speed vs accuracy at high freq
        if (freq > 40000.0f)      pCfg->DftNum = DFTNUM_2048;
        else if (freq > 1000.0f)  pCfg->DftNum = DFTNUM_4096;
        else                      pCfg->DftNum = DFTNUM_8192;
    } 
    else {
        // Low Frequency High Speed (e.g. 10Hz)
        // Use SINC2 filter to slow down data rate
        // Rate = 800k / (4 * 22) = 9090 Hz
        pCfg->DftSrc = DFTSRC_SINC2NOTCH;
        pCfg->Sinc3Osr = ADCSINC3OSR_4;
        pCfg->Sinc2Osr = ADCSINC2OSR_22; 
        pCfg->SampleRate = 9090.0f;

        // 8192 samples @ 9kHz = ~0.9 seconds
        // Covers 9 cycles of 10Hz. Perfect.
        pCfg->DftNum = DFTNUM_8192; 
    }
}

// --------------------------------------------------------------------------
// Measurement Functions
// --------------------------------------------------------------------------

static uint32_t AppBuff[512]; 

// Internal Helper for Calibration
static AD5940Err RunCalibration(void)
{
    ADCPGACal_Type pga_cal;
    
    pga_cal.AdcClkFreq = 16000000.0;
    pga_cal.SysClkFreq = 16000000.0;
    pga_cal.ADCPga = ADCPGA_1;
    pga_cal.ADCSinc2Osr = ADCSINC2OSR_1333;
    pga_cal.ADCSinc3Osr = ADCSINC3OSR_4;
    pga_cal.TimeOut10us = 100000; 
    pga_cal.VRef1p82 = 1.82; 
    pga_cal.VRef1p11 = 1.11;
    pga_cal.PGACalType = PGACALTYPE_OFFSETGAIN;

    return AD5940_ADCPGACal(&pga_cal);
}

// Command Handler for "CAL"
static void AppADCPgaCal(void)
{
    printf("LOG,Preparing for Calibration...\n");
    // Ensure clean state before CAL
    AD5940_PulseReset();
    AD5940_Initialize();
    AD5940_WriteReg(REG_INTC_INTCSEL0, 0xFFFFFFFF);
    AD5940_WriteReg(REG_INTC_INTCCLR, 0xFFFFFFFF);

    AD5940Err err = RunCalibration();
    if(err != AD5940ERR_OK) printf("LOG,ADC Cal Failed: %d\n", err);
    else printf("LOG,ADC Cal Passed.\n");
    
    // Cleanup: Turn off everything after CAL to leave clean slate
    AD5940_AFECtrlS(AFECTRL_ALL, bFALSE);
}

// High Speed Loop (HSTIA + HSDAC) for ALL Frequencies
void MeasureHighSpeed(float freq)
{
    // 1. HARD RESET & Re-Initialize
    AD5940_PulseReset();
    AD5940_Initialize(); 
    AD5940_WriteReg(REG_INTC_INTCSEL0, 0xFFFFFFFF); 
    AD5940_WriteReg(REG_INTC_INTCCLR, 0xFFFFFFFF);
    
    // 2. Calibrate
    // Note: Calibration routines are invasive. They might leave the AFE in weird states.
    AD5940Err calErr = RunCalibration();
    if(calErr != AD5940ERR_OK) {
        printf("LOG,Auto-Calibration Failed: %d\n", calErr);
    }
    
    // CRITICAL: Shut down EVERYTHING after calibration. 
    AD5940_AFECtrlS(AFECTRL_ALL, bFALSE);
    
    // CRITICAL: Disable Sequencer
    AD5940_SEQCtrlS(bFALSE); 
    
    // 3. Clock Configuration
    AD5940_WriteReg(REG_ALLON_OSCKEY, 0xCB14);      
    AD5940_WriteReg(REG_ALLON_OSCCON, 0x0003); // Enable HFOSC
    sleep_ms(2);                                    

    AD5940_WriteReg(REG_AFE_LPMODEKEY, 0xC59D6);    
    AD5940_WriteReg(REG_AFE_LPMODECLKSEL, 0);  // Select HFOSC
    
    spi_set_baudrate(SPI_PORT, SPI_BAUDRATE_HIGH);

    // 4. Switch Matrix Cleanup
    AD5940_WriteReg(REG_AFE_LPTIASW0, 0x0000); 
    AD5940_WriteReg(REG_AFE_LPDACSW0, 0x0000);

    // 5. Configure Loop
    HSLoopCfg_Type hs_loop;
    hs_loop.HsDacCfg.ExcitBufGain = EXCITBUFGAIN_2;
    hs_loop.HsDacCfg.HsDacGain = HSDACGAIN_1;
    hs_loop.HsDacCfg.HsDacUpdateRate = 7; 
    
    if (freq > 1000.0f) {
        hs_loop.HsTiaCfg.HstiaRtiaSel = HSTIARTIA_1K; 
    } else {
        hs_loop.HsTiaCfg.HstiaRtiaSel = HSTIARTIA_5K; 
    }
    
    hs_loop.HsTiaCfg.HstiaBias = HSTIABIAS_1P1;
    hs_loop.HsTiaCfg.HstiaCtia = 16;
    hs_loop.HsTiaCfg.DiodeClose = bFALSE;
    
    AD5940_HSLoopCfgS(&hs_loop);

    // 6. Switch Matrix (S+/S-)
    AD5940_WriteReg(REG_AFE_SWCON, BITM_AFE_SWCON_SWSOURCESEL); 
    AD5940_WriteReg(REG_AFE_DSWFULLCON, SWD_CE0); 
    AD5940_WriteReg(REG_AFE_PSWFULLCON, 0x0000);
    AD5940_WriteReg(REG_AFE_NSWFULLCON, 0x0000);
    AD5940_WriteReg(REG_AFE_TSWFULLCON, SWT_SE0 | SWT_T9);

    // 7. Get Dynamic Configuration
    FilterConfig_Type cfg;
    AD5940_GetFilterCfg(freq, &cfg);

    // 8. ADC Configuration
    ADCBaseCfg_Type adc_cfg;
    memset(&adc_cfg, 0, sizeof(adc_cfg)); 
    adc_cfg.ADCMuxP = ADCMUXP_HSTIA_P; 
    adc_cfg.ADCMuxN = ADCMUXN_HSTIA_N;
    adc_cfg.ADCPga = ADCPGA_1;
    AD5940_ADCBaseCfgS(&adc_cfg);

    // Initialize and Enable Clocks
    ADCFilterCfg_Type filter_cfg;
    memset(&filter_cfg, 0, sizeof(filter_cfg)); 
    
    filter_cfg.ADCSinc3Osr = cfg.Sinc3Osr;
    filter_cfg.ADCSinc2Osr = cfg.Sinc2Osr;
    filter_cfg.ADCAvgNum = ADCAVGNUM_16; 
    filter_cfg.ADCRate = ADCRATE_800KHZ;
    filter_cfg.BpNotch = bTRUE;
    filter_cfg.BpSinc3 = bFALSE;
    filter_cfg.Sinc2NotchEnable = bTRUE;
    
    // Explicitly Enable Clocks
    filter_cfg.DFTClkEnable = bTRUE;
    filter_cfg.WGClkEnable = bTRUE;
    filter_cfg.Sinc3ClkEnable = bTRUE;
    filter_cfg.Sinc2NotchClkEnable = bTRUE;
    
    AD5940_ADCFilterCfgS(&filter_cfg);
    
    // Configure ADC Buffers
    AD5940_WriteReg(REG_AFE_ADCBUFCON, 0x005F3D0F); 

    // 9. Waveform
    AD5940_WriteReg(REG_AFE_WGFCW, AD5940_WGFreqWordCal(freq, 16000000.0));
    AD5940_WriteReg(REG_AFE_WGAMPLITUDE, 2047); 
    AD5940_WriteReg(REG_AFE_WGOFFSET, 0);
    AD5940_WriteReg(REG_AFE_WGPHASE, 0);
    AD5940_WriteReg(REG_AFE_WGCON, WGTYPE_SIN); 

    // 10. DFT
    DFTCfg_Type dft_cfg;
    dft_cfg.DftNum = cfg.DftNum;
    dft_cfg.DftSrc = cfg.DftSrc;
    dft_cfg.HanWinEn = bTRUE;
    AD5940_DFTCfgS(&dft_cfg);

    // 11. Configure FIFO for DFT
    // Ensure we capture DFT results in FIFO
    FIFOCfg_Type fifo_cfg;
    fifo_cfg.FIFOEn = bTRUE;
    fifo_cfg.FIFOMode = FIFOMODE_FIFO;
    fifo_cfg.FIFOSize = FIFOSIZE_4KB; 
    fifo_cfg.FIFOSrc = FIFOSRC_DFT;   
    fifo_cfg.FIFOThresh = 4;          
    AD5940_FIFOCfg(&fifo_cfg);

    // 12. Run
    
    // Power Up Analog Blocks
    AD5940_AFECtrlS(AFECTRL_HSDACPWR|AFECTRL_HSTIAPWR|AFECTRL_ADCPWR|AFECTRL_DACREFPWR|AFECTRL_EXTBUFPWR|AFECTRL_INAMPPWR, bTRUE);
    
    sleep_ms(20); 
    
    // CRITICAL FIX: Force AFECON Clean
    uint32_t current_afecon = AD5940_ReadReg(REG_AFE_AFECON);
    current_afecon &= ~(0x00003000); // Clear Temp bits
    AD5940_WriteReg(REG_AFE_AFECON, current_afecon);

    // Enable ADC Repeat Mode
    AD5940_ADCRepeatCfgS(0xffffffff); 

    // Clear and Enable Interrupts
    AD5940_WriteReg(REG_INTC_INTCCLR, 0xFFFFFFFF);
    AD5940_WriteReg(REG_INTC_INTCSEL0, 0xFFFFFFFF); 

    // Start Conversion
    AD5940_AFECtrlS(AFECTRL_ADCCNV|AFECTRL_DFT|AFECTRL_WG, bTRUE); 

    // Calculate Timeout using Helper
    uint32_t sample_count = GetDftSamples(cfg.DftNum);
    float duration = (float)sample_count / cfg.SampleRate;
    uint32_t timeout_us = (uint32_t)(duration * 1000000.0f) + 2000000; 
    
    // Polling Loop with Exit Condition
    uint32_t start_time = to_us_since_boot(get_absolute_time());
    bool success = false;
    
    // Results
    int32_t real = 0;
    int32_t image = 0;

    while((to_us_since_boot(get_absolute_time()) - start_time) < timeout_us) {
        
        // Polling Strategy: Check FIFO Count instead of Interrupt Flags
        // One DFT result = 2 words (Real + Image). We wait for >= 2 words.
        uint32_t fifo_cnt = (AD5940_ReadReg(REG_AFE_FIFOCNTSTA) & 0x07FF0000) >> 16;
        
        if(fifo_cnt >= 2) {
            // Read result from FIFO
            uint32_t buffer[2]; 
            AD5940_FIFORd(buffer, 2);
            real = (int32_t)buffer[0];
            image = (int32_t)buffer[1];
            
            AD5940_WriteReg(REG_INTC_INTCCLR, 0xFFFFFFFF);
            success = true;
            break;
        }
        sleep_us(100);
    }

    if (!success) {
        uint32_t flags = AD5940_ReadReg(REG_INTC_INTCFLAG0);
        uint32_t afecon_read = AD5940_ReadReg(REG_AFE_AFECON);
        uint32_t fifo_cnt = (AD5940_ReadReg(REG_AFE_FIFOCNTSTA) & 0x07FF0000) >> 16;
        printf("LOG,Error: Measurement Timeout (HS). Exp Duration: %.2fs. Flags: 0x%08X, AFECON: 0x%08X, FIFO: %d\n", duration, flags, afecon_read, fifo_cnt);
    }

    // 13. Process Result
    // Sign Extend 18-bit to 32-bit
    // Store original raw values before processing
    int32_t raw_real = real;
    int32_t raw_image = image;

    if(real & (1<<17)) real |= 0xFFFC0000;
    if(image & (1<<17)) image |= 0xFFFC0000;
    
    float mag = sqrt((float)real*real + (float)image*image);
    float rtia_val = (freq > 1000.0f) ? 1000.0f : 5000.0f;
    float impedance = (mag > 0) ? (2047.0f / mag) * rtia_val : 0;
    float phase = atan2((float)image, (float)real);

    // Updated Output Format: Freq, Impedance, Phase, RawReal, RawImag
    printf("DATA,%.2f,%.4f,%.4f,%d,%d\n", freq, impedance, phase, raw_real, raw_image);
    
    AD5940_AFECtrlS(AFECTRL_ALL, bFALSE);
}

// --------------------------------------------------------------------------
// Main
// --------------------------------------------------------------------------

void process_command(char* cmd) {
    char* token = strtok(cmd, " ");
    if (!token) return;

    if (strcmp(token, "CAL") == 0) {
        AppADCPgaCal();
    } 
    else if (strcmp(token, "MEAS") == 0) {
        char* arg1 = strtok(NULL, " ");
        if (arg1) {
            float freq = strtof(arg1, NULL);
            if (freq < 10.0f) {
                freq = 10.0f; 
            }
            MeasureHighSpeed(freq);
        }
    }
    else if (strcmp(token, "ID") == 0) {
        uint32_t chip_id = AD5940_ReadReg(REG_AFECON_CHIPID);
        printf("LOG,Chip ID: 0x%04X\n", chip_id);
    }
}

int main() {
    stdio_init_all();
    sleep_ms(1000); 

    spi_init(SPI_PORT, SPI_BAUDRATE_HIGH);
    spi_set_format(SPI_PORT, 8, SPI_CPOL_0, SPI_CPHA_0, SPI_MSB_FIRST);
    
    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_init(PIN_CS);
    gpio_set_dir(PIN_CS, GPIO_OUT);
    gpio_put(PIN_CS, 1);

    gpio_init(PIN_RESET);
    gpio_set_dir(PIN_RESET, GPIO_OUT);
    gpio_put(PIN_RESET, 1);

    gpio_init(PIN_AD_INTERRUPT);
    gpio_set_dir(PIN_AD_INTERRUPT, GPIO_IN);
    gpio_pull_up(PIN_AD_INTERRUPT);

    // Initial Start
    AD5940_PulseReset();
    AD5940_Initialize();
    AD5940_WriteReg(REG_INTC_INTCSEL0, 0xFFFFFFFF); 
    AD5940_WriteReg(REG_INTC_INTCCLR, 0xFFFFFFFF);
    AD5940_SEQGenInit(AppBuff, 512);
    
    char buffer[64];
    int pos = 0;

    while (true) {
        int c = getchar_timeout_us(1000);
        if (c != PICO_ERROR_TIMEOUT) {
            if (c == '\n' || c == '\r') {
                if (pos > 0) {
                    buffer[pos] = 0;
                    process_command(buffer);
                    pos = 0;
                }
            } else if (pos < 63) {
                buffer[pos++] = (char)c;
            }
        }
    }
}