#include "eis.h"
#include <math.h>
#include <string.h>
#include <stdio.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

/* resolved hardware state */
static struct {
    EISConfig cfg;
    float sys_clk;
    float rcal_ohms;
    uint32_t rcal_sw_d, rcal_sw_p, rcal_sw_n, rcal_sw_t;
    uint32_t dut_sw_d, dut_sw_p, dut_sw_n, dut_sw_t;
    uint32_t dut_mux_vp, dut_mux_vn;
    uint32_t rtia_reg;
    uint32_t dertia_reg;
} ctx;

static const uint32_t rtia_map[] = {
    [RTIA_200]     = HSTIARTIA_200,
    [RTIA_1K]      = HSTIARTIA_1K,
    [RTIA_5K]      = HSTIARTIA_5K,
    [RTIA_10K]     = HSTIARTIA_10K,
    [RTIA_20K]     = HSTIARTIA_20K,
    [RTIA_40K]     = HSTIARTIA_40K,
    [RTIA_80K]     = HSTIARTIA_80K,
    [RTIA_160K]    = HSTIARTIA_160K,
    [RTIA_EXT_DE0] = HSTIARTIA_OPEN,
};

static void resolve_config(void)
{
    /* RTIA */
    ctx.rtia_reg   = rtia_map[ctx.cfg.rtia];
    ctx.dertia_reg = (ctx.cfg.rtia == RTIA_EXT_DE0) ? HSTIADERTIA_TODE : HSTIADERTIA_OPEN;

    /* RCAL */
    switch (ctx.cfg.rcal) {
    case RCAL_200R:
        ctx.rcal_ohms = 200.0f;
        ctx.rcal_sw_d = SWD_RCAL0;
        ctx.rcal_sw_p = SWP_RCAL0;
        ctx.rcal_sw_n = SWN_RCAL1;
        ctx.rcal_sw_t = SWT_RCAL1 | SWT_TRTIA;
        break;
    default: /* RCAL_3K */
        ctx.rcal_ohms = 3000.0f;
        ctx.rcal_sw_d = SWD_RCAL0;
        ctx.rcal_sw_p = SWP_RCAL0;
        ctx.rcal_sw_n = SWN_AIN0;
        ctx.rcal_sw_t = SWT_AIN0 | SWT_TRTIA;
        break;
    }

    /* DUT electrode routing */
    switch (ctx.cfg.electrode) {
    case ELEC_3WIRE:
        ctx.dut_sw_d   = SWD_CE0;
        ctx.dut_sw_p   = SWP_RE0;
        ctx.dut_sw_n   = SWN_SE0;
        ctx.dut_sw_t   = SWT_SE0LOAD | SWT_TRTIA;
        ctx.dut_mux_vp = ADCMUXP_P_NODE;
        ctx.dut_mux_vn = ADCMUXN_N_NODE;
        break;
    default: /* ELEC_4WIRE */
        ctx.dut_sw_d   = SWD_AIN3;
        ctx.dut_sw_p   = SWP_AIN3;
        ctx.dut_sw_n   = SWN_AIN0;
        ctx.dut_sw_t   = SWT_AIN0 | SWT_TRTIA;
        ctx.dut_mux_vp = ADCMUXP_AIN2;
        ctx.dut_mux_vn = ADCMUXN_AIN1;
        break;
    }
}

static void apply_hsloop(void)
{
    HSLoopCfg_Type hs;
    AD5940_StructInit(&hs, sizeof(hs));
    hs.HsDacCfg.ExcitBufGain    = EXCITBUFGAIN_2;
    hs.HsDacCfg.HsDacGain       = HSDACGAIN_0P2;
    hs.HsDacCfg.HsDacUpdateRate = 7;
    hs.HsTiaCfg.HstiaBias       = HSTIABIAS_1P1;
    hs.HsTiaCfg.HstiaRtiaSel    = ctx.rtia_reg;
    hs.HsTiaCfg.HstiaCtia       = 16;
    hs.HsTiaCfg.HstiaDeRtia     = ctx.dertia_reg;
    hs.HsTiaCfg.HstiaDeRload    = HSTIADERLOAD_OPEN;
    hs.HsTiaCfg.HstiaDe1Rtia    = HSTIADERTIA_OPEN;
    hs.HsTiaCfg.HstiaDe1Rload   = HSTIADERLOAD_OPEN;
    hs.HsTiaCfg.DiodeClose      = bFALSE;
    hs.WgCfg.WgType              = WGTYPE_SIN;
    hs.WgCfg.GainCalEn           = bTRUE;
    hs.WgCfg.OffsetCalEn         = bTRUE;
    hs.WgCfg.SinCfg.SinAmplitudeWord = ctx.cfg.excit_amp;
    hs.WgCfg.SinCfg.SinFreqWord      = 0;
    hs.WgCfg.SinCfg.SinOffsetWord    = 0;
    hs.WgCfg.SinCfg.SinPhaseWord     = 0;
    hs.SWMatCfg.Dswitch = ctx.rcal_sw_d;
    hs.SWMatCfg.Pswitch = ctx.rcal_sw_p;
    hs.SWMatCfg.Nswitch = ctx.rcal_sw_n;
    hs.SWMatCfg.Tswitch = ctx.rcal_sw_t;
    AD5940_HSLoopCfgS(&hs);

    if (ctx.cfg.rtia == RTIA_EXT_DE0)
        AD5940_WriteReg(REG_AFE_DE0RESCON, 0x97);
}

/* ---------- public ---------- */

void eis_default_config(EISConfig *cfg)
{
    memset(cfg, 0, sizeof(*cfg));
    cfg->freq_start_hz    = 1000.0f;
    cfg->freq_stop_hz     = 200000.0f;
    cfg->points_per_decade = 10;
    cfg->rtia      = RTIA_5K;
    cfg->rcal      = RCAL_3K;
    cfg->electrode = ELEC_4WIRE;
    cfg->pga       = ADCPGA_1P5;
    cfg->excit_amp = 500;
}

uint32_t eis_calc_num_points(const EISConfig *cfg)
{
    if (cfg->freq_stop_hz <= cfg->freq_start_hz || cfg->points_per_decade == 0)
        return 1;
    float decades = log10f(cfg->freq_stop_hz / cfg->freq_start_hz);
    uint32_t n = (uint32_t)(decades * cfg->points_per_decade + 0.5f) + 1;
    if (n > EIS_MAX_POINTS) n = EIS_MAX_POINTS;
    if (n < 2) n = 2;
    return n;
}

void eis_init(const EISConfig *cfg)
{
    memcpy(&ctx.cfg, cfg, sizeof(EISConfig));
    ctx.sys_clk = 16000000.0f;
    resolve_config();

    CLKCfg_Type clk;
    memset(&clk, 0, sizeof(clk));
    clk.HFOSCEn        = bTRUE;
    clk.HfOSC32MHzMode = bFALSE;
    clk.SysClkSrc       = SYSCLKSRC_HFOSC;
    clk.ADCCLkSrc       = ADCCLKSRC_HFOSC;
    clk.SysClkDiv       = SYSCLKDIV_1;
    clk.ADCClkDiv       = ADCCLKDIV_1;
    clk.LFOSCEn         = bTRUE;
    clk.HFXTALEn        = bFALSE;
    AD5940_CLKCfg(&clk);

    AFERefCfg_Type ref;
    AD5940_StructInit(&ref, sizeof(ref));
    ref.HpBandgapEn = bTRUE;
    ref.Hp1V1BuffEn = bTRUE;
    ref.Hp1V8BuffEn = bTRUE;
    ref.HSDACRefEn  = bTRUE;
    ref.LpBandgapEn = bFALSE;
    ref.LpRefBufEn  = bFALSE;
    AD5940_REFCfgS(&ref);

    AD5940_AFEPwrBW(AFEPWR_LP, AFEBW_250KHZ);

    AD5940_INTCCfg(AFEINTC_0, AFEINTSRC_DFTRDY, bTRUE);
    AD5940_INTCCfg(AFEINTC_1, AFEINTSRC_DFTRDY, bTRUE);
    AD5940_INTCClrFlag(AFEINTSRC_ALLINT);

    AGPIOCfg_Type gpio;
    AD5940_StructInit(&gpio, sizeof(gpio));
    gpio.FuncSet     = GP0_INT;
    gpio.OutputEnSet = AGPIO_Pin0;
    AD5940_AGPIOCfg(&gpio);

    AD5940_WriteReg(REG_AFE_FIFOCON, 0);

    SEQCfg_Type seq;
    seq.SeqMemSize   = SEQMEMSIZE_4KB;
    seq.SeqBreakEn   = bFALSE;
    seq.SeqIgnoreEn  = bFALSE;
    seq.SeqCntCRCClr = bFALSE;
    seq.SeqEnable    = bTRUE;
    seq.SeqWrTimer   = 0;
    AD5940_SEQCfg(&seq);

    apply_hsloop();

    ADCBaseCfg_Type adc;
    adc.ADCMuxP = ADCMUXP_P_NODE;
    adc.ADCMuxN = ADCMUXN_N_NODE;
    adc.ADCPga  = cfg->pga;
    AD5940_ADCBaseCfgS(&adc);

    AD5940_AFECtrlS(AFECTRL_ALL, bFALSE);
}

void eis_reconfigure(const EISConfig *cfg)
{
    memcpy(&ctx.cfg, cfg, sizeof(EISConfig));
    resolve_config();
}

/* ---------- internal helpers ---------- */

static void configure_freq(float freq_hz)
{
    FreqParams_Type fp = AD5940_GetFreqParameters(freq_hz);

    if (fp.HighPwrMode) {
        fp.DftSrc      = DFTSRC_ADCRAW;
        fp.ADCSinc3Osr = ADCSINC3OSR_2;
        fp.ADCSinc2Osr = 0;
        fp.DftNum      = DFTNUM_16384;
    }

    AD5940_WriteReg(REG_AFE_WGFCW,
                    AD5940_WGFreqWordCal(freq_hz, ctx.sys_clk));

    ADCFilterCfg_Type filt;
    AD5940_StructInit(&filt, sizeof(filt));
    filt.ADCSinc3Osr         = fp.ADCSinc3Osr;
    filt.ADCSinc2Osr         = fp.ADCSinc2Osr;
    filt.ADCAvgNum           = ADCAVGNUM_16;
    filt.ADCRate             = ADCRATE_800KHZ;
    filt.BpNotch             = bTRUE;
    filt.BpSinc3             = bFALSE;
    filt.Sinc2NotchEnable    = bTRUE;
    filt.Sinc3ClkEnable      = bTRUE;
    filt.Sinc2NotchClkEnable = bTRUE;
    filt.DFTClkEnable        = bTRUE;
    filt.WGClkEnable         = bTRUE;
    AD5940_ADCFilterCfgS(&filt);

    DFTCfg_Type dft;
    dft.DftNum   = fp.DftNum;
    dft.DftSrc   = fp.DftSrc;
    dft.HanWinEn = bTRUE;
    AD5940_DFTCfgS(&dft);
}

static int32_t sign_extend_18(uint32_t v)
{
    return (v & (1UL << 17)) ? (int32_t)(v | 0xFFFC0000UL) : (int32_t)v;
}

static void dft_measure(uint32_t mux_p, uint32_t mux_n, iImpCar_Type *out)
{
    AD5940_AFECtrlS(AFECTRL_ADCCNV | AFECTRL_DFT, bFALSE);
    AD5940_WriteReg(REG_AFE_FIFOCON, 0);
    AD5940_ReadAfeResult(AFERESULT_DFTREAL);
    AD5940_ReadAfeResult(AFERESULT_DFTIMAGE);
    AD5940_INTCClrFlag(AFEINTSRC_DFTRDY);

    AD5940_ADCMuxCfgS(mux_p, mux_n);

    AD5940_AFECtrlS(AFECTRL_WG | AFECTRL_ADCPWR, bTRUE);
    AD5940_Delay10us(25);

    AD5940_ClrMCUIntFlag();
    AD5940_AFECtrlS(AFECTRL_ADCCNV | AFECTRL_DFT, bTRUE);

    while (!AD5940_GetMCUIntFlag())
        vTaskDelay(1);

    AD5940_AFECtrlS(AFECTRL_ADCCNV | AFECTRL_DFT |
                    AFECTRL_WG | AFECTRL_ADCPWR, bFALSE);
    AD5940_INTCClrFlag(AFEINTSRC_DFTRDY);

    out->Real  = sign_extend_18(AD5940_ReadAfeResult(AFERESULT_DFTREAL));
    out->Image = sign_extend_18(AD5940_ReadAfeResult(AFERESULT_DFTIMAGE));
    out->Image = -out->Image;
}

/* RTIA calibration: 2 DFTs through current RCAL switch config */
static fImpCar_Type measure_rtia(void)
{
    iImpCar_Type v_rcal, v_raw;
    dft_measure(ADCMUXP_P_NODE, ADCMUXN_N_NODE, &v_rcal);
    dft_measure(ADCMUXP_HSTIA_P, ADCMUXN_HSTIA_N, &v_raw);
    v_raw.Real  = -v_raw.Real;
    v_raw.Image = -v_raw.Image;
    fImpCar_Type rtia = AD5940_ComplexDivInt(&v_raw, &v_rcal);
    rtia.Real  *= ctx.rcal_ohms;
    rtia.Image *= ctx.rcal_ohms;
    return rtia;
}

/* ---------- measurement ---------- */

int eis_measure_point(float freq_hz, EISPoint *out)
{
    configure_freq(freq_hz);

    SWMatrixCfg_Type sw;
    iImpCar_Type v_tia, v_sense;

    /* switch to RCAL before power-up */
    sw.Dswitch = ctx.rcal_sw_d;
    sw.Pswitch = ctx.rcal_sw_p;
    sw.Nswitch = ctx.rcal_sw_n;
    sw.Tswitch = ctx.rcal_sw_t;
    AD5940_SWMatrixCfgS(&sw);

    AD5940_AFECtrlS(AFECTRL_HPREFPWR | AFECTRL_HSTIAPWR | AFECTRL_INAMPPWR |
                    AFECTRL_EXTBUFPWR | AFECTRL_DACREFPWR | AFECTRL_HSDACPWR |
                    AFECTRL_SINC2NOTCH, bTRUE);

    /* RCAL before */
    fImpCar_Type rtia_before = measure_rtia();

    /* DUT forward */
    sw.Dswitch = ctx.dut_sw_d;
    sw.Pswitch = ctx.dut_sw_p;
    sw.Nswitch = ctx.dut_sw_n;
    sw.Tswitch = ctx.dut_sw_t;
    AD5940_SWMatrixCfgS(&sw);
    AD5940_Delay10us(50);

    dft_measure(ADCMUXP_HSTIA_P, ADCMUXN_HSTIA_N, &v_tia);
    v_tia.Real  = -v_tia.Real;
    v_tia.Image = -v_tia.Image;
    dft_measure(ctx.dut_mux_vp, ctx.dut_mux_vn, &v_sense);

    iImpCar_Type v_tia_fwd   = v_tia;
    iImpCar_Type v_sense_fwd = v_sense;

    /* RCAL after */
    sw.Dswitch = ctx.rcal_sw_d;
    sw.Pswitch = ctx.rcal_sw_p;
    sw.Nswitch = ctx.rcal_sw_n;
    sw.Tswitch = ctx.rcal_sw_t;
    AD5940_SWMatrixCfgS(&sw);
    AD5940_Delay10us(50);

    fImpCar_Type rtia_after = measure_rtia();

    /* DUT reverse (DUT first, then RCAL) */
    sw.Dswitch = ctx.dut_sw_d;
    sw.Pswitch = ctx.dut_sw_p;
    sw.Nswitch = ctx.dut_sw_n;
    sw.Tswitch = ctx.dut_sw_t;
    AD5940_SWMatrixCfgS(&sw);
    AD5940_Delay10us(50);

    dft_measure(ADCMUXP_HSTIA_P, ADCMUXN_HSTIA_N, &v_tia);
    v_tia.Real  = -v_tia.Real;
    v_tia.Image = -v_tia.Image;
    dft_measure(ctx.dut_mux_vp, ctx.dut_mux_vn, &v_sense);

    iImpCar_Type v_tia_rev   = v_tia;
    iImpCar_Type v_sense_rev = v_sense;

    /* RCAL reverse */
    sw.Dswitch = ctx.rcal_sw_d;
    sw.Pswitch = ctx.rcal_sw_p;
    sw.Nswitch = ctx.rcal_sw_n;
    sw.Tswitch = ctx.rcal_sw_t;
    AD5940_SWMatrixCfgS(&sw);
    AD5940_Delay10us(50);

    fImpCar_Type rtia_rev = measure_rtia();

    /* power down, open switches */
    AD5940_AFECtrlS(AFECTRL_WG | AFECTRL_ADCPWR | AFECTRL_ADCCNV |
                    AFECTRL_DFT | AFECTRL_SINC2NOTCH | AFECTRL_HSDACPWR |
                    AFECTRL_HSTIAPWR | AFECTRL_INAMPPWR |
                    AFECTRL_EXTBUFPWR, bFALSE);

    sw.Dswitch = SWD_OPEN;
    sw.Pswitch = SWP_OPEN;
    sw.Nswitch = SWN_OPEN;
    sw.Tswitch = SWT_OPEN;
    AD5940_SWMatrixCfgS(&sw);

    /* forward Z using averaged RTIA bracket */
    fImpCar_Type rtia_avg = {
        .Real  = (rtia_before.Real + rtia_after.Real) * 0.5f,
        .Image = (rtia_before.Image + rtia_after.Image) * 0.5f,
    };
    fImpCar_Type fs_fwd = { (float)v_sense_fwd.Real, (float)v_sense_fwd.Image };
    fImpCar_Type ft_fwd = { (float)v_tia_fwd.Real,   (float)v_tia_fwd.Image };
    fImpCar_Type num    = AD5940_ComplexMulFloat(&fs_fwd, &rtia_avg);
    fImpCar_Type z_fwd  = AD5940_ComplexDivFloat(&num, &ft_fwd);

    /* reverse Z using RTIA from RCAL measured after DUT */
    fImpCar_Type fs_rev = { (float)v_sense_rev.Real, (float)v_sense_rev.Image };
    fImpCar_Type ft_rev = { (float)v_tia_rev.Real,   (float)v_tia_rev.Image };
    num                 = AD5940_ComplexMulFloat(&fs_rev, &rtia_rev);
    fImpCar_Type z_rev  = AD5940_ComplexDivFloat(&num, &ft_rev);

    float mag_fwd = AD5940_ComplexMag(&z_fwd);
    float mag_rev = AD5940_ComplexMag(&z_rev);

    out->freq_hz        = freq_hz;
    out->z_real         = z_fwd.Real;
    out->z_imag         = z_fwd.Image;
    out->mag_ohms       = mag_fwd;
    out->phase_deg      = AD5940_ComplexPhase(&z_fwd) * (float)(180.0 / M_PI);
    out->rtia_mag_before = AD5940_ComplexMag(&rtia_before);
    out->rtia_mag_after  = AD5940_ComplexMag(&rtia_after);
    out->rev_mag        = mag_rev;
    out->rev_phase      = AD5940_ComplexPhase(&z_rev) * (float)(180.0 / M_PI);
    out->pct_err        = (mag_fwd > 0.0f)
                          ? fabsf(mag_fwd - mag_rev) / mag_fwd * 100.0f : 0.0f;

    return 0;
}

int eis_sweep(EISPoint *out, uint32_t max_points, eis_point_cb_t cb)
{
    uint32_t n = eis_calc_num_points(&ctx.cfg);
    if (n > max_points) n = max_points;

    /* guard: throwaway at start frequency to warm up AFE */
    EISPoint guard;
    eis_measure_point(ctx.cfg.freq_start_hz, &guard);

    SoftSweepCfg_Type sweep;
    sweep.SweepEn     = bTRUE;
    sweep.SweepStart  = ctx.cfg.freq_start_hz;
    sweep.SweepStop   = ctx.cfg.freq_stop_hz;
    sweep.SweepPoints = n;
    sweep.SweepLog    = bTRUE;
    sweep.SweepIndex  = 0;

    printf("\n%10s  %12s  %10s  %12s  %12s  %7s\n",
           "Freq(Hz)", "|Z|(Ohm)", "Phase(deg)", "Re(Ohm)", "Im(Ohm)", "Err%");
    printf("---------------------------------------------------------------------\n");

    eis_measure_point(ctx.cfg.freq_start_hz, &out[0]);
    printf("%10.1f  %12.2f  %10.2f  %12.2f  %12.2f  %6.2f%%\n",
           out[0].freq_hz, out[0].mag_ohms, out[0].phase_deg,
           out[0].z_real, out[0].z_imag, out[0].pct_err);
    if (cb) cb(0, &out[0]);

    for (uint32_t i = 1; i < n; i++) {
        float freq;
        AD5940_SweepNext(&sweep, &freq);
        eis_measure_point(freq, &out[i]);
        printf("%10.1f  %12.2f  %10.2f  %12.2f  %12.2f  %6.2f%%\n",
               out[i].freq_hz, out[i].mag_ohms, out[i].phase_deg,
               out[i].z_real, out[i].z_imag, out[i].pct_err);
        if (cb) cb((uint16_t)i, &out[i]);
    }

    /* guard: throwaway at stop frequency to cap the sweep cleanly */
    eis_measure_point(ctx.cfg.freq_stop_hz, &guard);

    AD5940_AFECtrlS(AFECTRL_ALL, bFALSE);
    return (int)n;
}
