#include "refs.h"
#include "protocol.h"
#include "temp.h"
#include "ad5940.h"
#include <string.h>
#include <stdio.h>
#include <math.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

extern const uint32_t lp_rtia_map[];
extern const float lp_rtia_ohms[];

static int32_t read_adc_code(void)
{
    AD5940_INTCClrFlag(AFEINTSRC_SINC2RDY);
    AD5940_AFECtrlS(AFECTRL_ADCPWR, bTRUE);
    AD5940_Delay10us(25);
    AD5940_AFECtrlS(AFECTRL_ADCCNV, bTRUE);

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

    AD5940_AFECtrlS(AFECTRL_ADCCNV | AFECTRL_ADCPWR, bFALSE);
    AD5940_INTCClrFlag(AFEINTSRC_SINC2RDY);

    uint32_t raw = AD5940_ReadAfeResult(AFERESULT_SINC2);
    return (raw & (1UL << 15)) ? (int32_t)(raw | 0xFFFF0000UL) : (int32_t)raw;
}

#define ADC_OK          0
#define ADC_CLIPPED     1
#define ADC_OSCILLATING 2
#define CLIP_THRESHOLD  30000
#define PROBE_SAMPLES   8

static void init_lp_for_probe(uint32_t rtia_reg)
{
    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.LpBandgapEn  = bTRUE;
    ref.LpRefBufEn   = bTRUE;
    ref.LpRefBoostEn = bTRUE;
    AD5940_REFCfgS(&ref);

    AD5940_AFEPwrBW(AFEPWR_LP, AFEBW_250KHZ);

    LPLoopCfg_Type lp;
    AD5940_StructInit(&lp, sizeof(lp));
    lp.LpDacCfg.LpdacSel      = LPDAC0;
    lp.LpDacCfg.LpDacSrc      = LPDACSRC_MMR;
    lp.LpDacCfg.LpDacVzeroMux = LPDACVZERO_6BIT;
    lp.LpDacCfg.LpDacVbiasMux = LPDACVBIAS_12BIT;
    lp.LpDacCfg.LpDacSW       = LPDACSW_VZERO2LPTIA | LPDACSW_VBIAS2LPPA;
    lp.LpDacCfg.LpDacRef      = LPDACREF_2P5;
    lp.LpDacCfg.DataRst       = bFALSE;
    lp.LpDacCfg.PowerEn       = bTRUE;
    lp.LpDacCfg.DacData6Bit   = 26;
    lp.LpDacCfg.DacData12Bit  = (uint16_t)((1093.75f - 200.0f) / 0.537f + 0.5f);

    lp.LpAmpCfg.LpAmpSel    = LPAMP0;
    lp.LpAmpCfg.LpAmpPwrMod = LPAMPPWR_BOOST3;
    lp.LpAmpCfg.LpPaPwrEn   = bTRUE;
    lp.LpAmpCfg.LpTiaPwrEn  = bTRUE;
    lp.LpAmpCfg.LpTiaRf     = LPTIARF_SHORT;
    lp.LpAmpCfg.LpTiaRload  = LPTIARLOAD_SHORT;
    lp.LpAmpCfg.LpTiaRtia   = rtia_reg;
    lp.LpAmpCfg.LpTiaSW     = LPTIASW(2) | LPTIASW(4) | LPTIASW(5) |
                                LPTIASW(7) | LPTIASW(12);
    AD5940_LPLoopCfgS(&lp);

    ADCBaseCfg_Type adc;
    adc.ADCMuxP = ADCMUXP_LPTIA0_P;
    adc.ADCMuxN = ADCMUXN_LPTIA0_N;
    adc.ADCPga  = ADCPGA_1P5;
    AD5940_ADCBaseCfgS(&adc);

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

    AD5940_INTCCfg(AFEINTC_0, AFEINTSRC_SINC2RDY, bTRUE);
    AD5940_INTCCfg(AFEINTC_1, AFEINTSRC_SINC2RDY, 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);
}

static void shutdown_lp(void)
{
    LPLoopCfg_Type lp;
    AD5940_StructInit(&lp, sizeof(lp));
    lp.LpDacCfg.LpdacSel      = LPDAC0;
    lp.LpDacCfg.LpDacSrc      = LPDACSRC_MMR;
    lp.LpDacCfg.LpDacVzeroMux = LPDACVZERO_6BIT;
    lp.LpDacCfg.LpDacVbiasMux = LPDACVBIAS_12BIT;
    lp.LpDacCfg.LpDacSW       = 0;
    lp.LpDacCfg.LpDacRef      = LPDACREF_2P5;
    lp.LpDacCfg.PowerEn       = bFALSE;
    lp.LpDacCfg.DataRst       = bFALSE;
    lp.LpDacCfg.DacData6Bit   = 0;
    lp.LpDacCfg.DacData12Bit  = 0;
    lp.LpAmpCfg.LpAmpSel      = LPAMP0;
    lp.LpAmpCfg.LpAmpPwrMod   = LPAMPPWR_NORM;
    lp.LpAmpCfg.LpPaPwrEn     = bFALSE;
    lp.LpAmpCfg.LpTiaPwrEn    = bFALSE;
    lp.LpAmpCfg.LpTiaRf       = LPTIARF_OPEN;
    lp.LpAmpCfg.LpTiaRload    = LPTIARLOAD_SHORT;
    lp.LpAmpCfg.LpTiaRtia     = LPTIARTIA_OPEN;
    lp.LpAmpCfg.LpTiaSW       = 0;
    AD5940_LPLoopCfgS(&lp);

    SWMatrixCfg_Type sw = { SWD_OPEN, SWP_OPEN, SWN_OPEN, SWT_OPEN };
    AD5940_SWMatrixCfgS(&sw);
}

static int check_adc_clipping(uint8_t lp_rtia_idx)
{
    init_lp_for_probe(lp_rtia_map[lp_rtia_idx]);
    vTaskDelay(pdMS_TO_TICKS(50));

    int clip_count = 0;
    int sign_changes = 0;
    int any_high = 0;
    int32_t prev_sign = 0;

    for (int i = 0; i < PROBE_SAMPLES; i++) {
        int32_t code = read_adc_code();
        int32_t abs_code = code < 0 ? -code : code;

        if (abs_code > CLIP_THRESHOLD) {
            clip_count++;
            any_high = 1;
        }

        int32_t cur_sign = (code > 0) ? 1 : ((code < 0) ? -1 : 0);
        if (i > 0 && cur_sign != 0 && prev_sign != 0 && cur_sign != prev_sign)
            sign_changes++;
        if (cur_sign != 0) prev_sign = cur_sign;
    }

    shutdown_lp();
    AD5940_AFECtrlS(AFECTRL_ALL, bFALSE);

    if (sign_changes >= 3 && any_high)
        return ADC_OSCILLATING;
    if (clip_count > PROBE_SAMPLES / 2)
        return ADC_CLIPPED;
    return ADC_OK;
}

static void lp_find_valid_range(LpRtiaRange *range)
{
    int low_clip = check_adc_clipping(0);
    int high_clip = check_adc_clipping(LP_RTIA_512K);

    uint8_t lo = 0, hi = LP_RTIA_512K;

    if (low_clip != ADC_OK) {
        uint8_t a = 0, b = LP_RTIA_512K;
        while (a < b) {
            uint8_t m = (a + b) / 2;
            if (check_adc_clipping(m) != ADC_OK)
                a = m + 1;
            else
                b = m;
        }
        lo = a;
    }

    if (high_clip != ADC_OK) {
        uint8_t a = lo, b = LP_RTIA_512K;
        while (a < b) {
            uint8_t m = (a + b + 1) / 2;
            if (check_adc_clipping(m) != ADC_OK)
                a = m;
            else
                b = m - 1;
        }
        hi = a;
    }

    range->low_idx = lo;
    range->high_idx = hi;
    range->valid = (lo <= hi) ? 1 : 0;
}

void refs_collect(RefStore *store, const EISConfig *cfg)
{
    memset(store, 0, sizeof(*store));

    EISConfig ref_cfg = *cfg;
    ref_cfg.rcal = RCAL_3K;

    for (int r = 0; r < REF_EIS_RTIA_COUNT; r++) {
        ref_cfg.rtia = (EISRtia)r;
        eis_init(&ref_cfg);

        send_ref_frame(REF_MODE_EIS, (uint8_t)r);

        uint32_t n = eis_calc_num_points(&ref_cfg);
        send_sweep_start(n, ref_cfg.freq_start_hz, ref_cfg.freq_stop_hz);

        int got = eis_sweep(store->eis[r].pts, n, send_eis_point);
        store->eis[r].n_points = (uint32_t)got;
        store->eis[r].valid = (got > 0) ? 1 : 0;

        send_sweep_end();
        printf("Ref EIS RTIA %d: %d pts\n", r, got);
    }

    printf("Ref: LP range search (LSV)\n");
    send_ref_frame(REF_MODE_LSV, 0);
    lp_find_valid_range(&store->lsv_range);
    if (store->lsv_range.valid)
        send_ref_lp_range(REF_MODE_LSV, store->lsv_range.low_idx, store->lsv_range.high_idx);
    printf("Ref LSV range: %u-%u\n", store->lsv_range.low_idx, store->lsv_range.high_idx);

    printf("Ref: LP range search (Amp)\n");
    send_ref_frame(REF_MODE_AMP, 0);
    lp_find_valid_range(&store->amp_range);
    if (store->amp_range.valid)
        send_ref_lp_range(REF_MODE_AMP, store->amp_range.low_idx, store->amp_range.high_idx);
    printf("Ref Amp range: %u-%u\n", store->amp_range.low_idx, store->amp_range.high_idx);

    printf("Ref: LP range search (Cl)\n");
    send_ref_frame(REF_MODE_CL, 0);
    lp_find_valid_range(&store->cl_range);
    if (store->cl_range.valid)
        send_ref_lp_range(REF_MODE_CL, store->cl_range.low_idx, store->cl_range.high_idx);
    printf("Ref Cl range: %u-%u\n", store->cl_range.low_idx, store->cl_range.high_idx);

    printf("Ref: pH OCP\n");
    send_ref_frame(REF_MODE_PH, 0);
    PhConfig ph_cfg;
    ph_cfg.stabilize_s = 10.0f;
    ph_cfg.temp_c = temp_get();
    echem_ph_ocp(&ph_cfg, &store->ph_ref);
    store->ph_valid = 1;
    send_ph_result(store->ph_ref.v_ocp_mv, store->ph_ref.ph, store->ph_ref.temp_c);

    store->has_refs = 1;
    send_refs_done();
    printf("Ref collection complete\n");
}

void refs_send(const RefStore *store)
{
    if (!store->has_refs) {
        send_ref_status(0);
        return;
    }
    send_ref_status(1);

    for (int r = 0; r < REF_EIS_RTIA_COUNT; r++) {
        if (!store->eis[r].valid) continue;

        send_ref_frame(REF_MODE_EIS, (uint8_t)r);
        uint32_t n = store->eis[r].n_points;
        send_sweep_start(n, store->eis[r].pts[0].freq_hz,
                         store->eis[r].pts[n - 1].freq_hz);
        for (uint32_t i = 0; i < n; i++)
            send_eis_point((uint16_t)i, &store->eis[r].pts[i]);
        send_sweep_end();
    }

    if (store->lsv_range.valid)
        send_ref_lp_range(REF_MODE_LSV, store->lsv_range.low_idx, store->lsv_range.high_idx);
    if (store->amp_range.valid)
        send_ref_lp_range(REF_MODE_AMP, store->amp_range.low_idx, store->amp_range.high_idx);
    if (store->cl_range.valid)
        send_ref_lp_range(REF_MODE_CL, store->cl_range.low_idx, store->cl_range.high_idx);

    if (store->ph_valid) {
        send_ref_frame(REF_MODE_PH, 0);
        send_ph_result(store->ph_ref.v_ocp_mv, store->ph_ref.ph, store->ph_ref.temp_c);
    }

    send_refs_done();
}

void refs_clear(RefStore *store)
{
    memset(store, 0, sizeof(*store));
    send_ref_status(0);
}