use defmt::{info, error};

use embassy_stm32::spi::Error;
use embassy_stm32::exti::ExtiInput;
use embassy_sync::blocking_mutex::raw::ThreadModeRawMutex;
use embassy_sync::channel::Channel;
use embassy_sync::mutex::Mutex;
use embassy_time::Timer;

use static_cell::StaticCell;

use heapless::Vec;

use num::complex::Complex;

use crate::ad5940::*;
use crate::ad5940_registers::*;

use bioz_icd_rs::{SingleImpedanceOutput, IcdDftNum, ImpedanceInitError, MeasurementPointSet, SweepImpedanceOutput};
use crate::icd_mapping::IntoDftnum;

pub static IMPEDANCE_CHANNEL_SINGLE: Channel<ThreadModeRawMutex, SingleImpedanceOutput, 50> = Channel::new();
pub static IMPEDANCE_CHANNEL_SWEEP: Channel<ThreadModeRawMutex, SweepImpedanceOutput, 5> = Channel::new();

pub type ImpedanceSetupType = Mutex<ThreadModeRawMutex, ImpedanceSetup>;
pub static IMPEDANCE_SETUP: StaticCell<ImpedanceSetupType> = StaticCell::new();

pub const RCAL: f32 = 1008.0; // Calibration resistor in Ohm

#[derive(PartialEq, Copy, Clone)]
pub enum RunningMode {
    None,
    SingleFrequency2Lead,
    SingleFrequency4Lead,
    SweepFrequency2Lead(MeasurementPointSet),
    SweepFrequency4Lead(MeasurementPointSet),
}

pub enum RtiaCalibrated {
    None,
    Single(RtiaCalibrationResult),
    Vec8(Vec<RtiaCalibrationResult, 8>),
    Vec18(Vec<RtiaCalibrationResult, 18>),
}

#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RtiaCalibrationResult {
    pub magnitude: f32,
    pub phase: f32,
}

pub struct ImpedanceSetup {
    ad5940: AD5940,
    dsp_config: Option<DspConfig>,
    pub running_mode: RunningMode,
    pub rtia_calibrated: RtiaCalibrated,
}

impl ImpedanceSetup {
    pub fn new(ad5940: AD5940) -> Self {
        ImpedanceSetup { ad5940, dsp_config: None, running_mode: RunningMode::None, rtia_calibrated: RtiaCalibrated::None }
    }
        
    pub async fn init(&mut self) -> Result<(), Error> {
        // AFECON:
        self.ad5940.afecon(
            AFECON::DACBUFEN
            | AFECON::DACREFEN
            | AFECON::SINC2EN
            | AFECON::TIAEN
            | AFECON::INAMPEN
            | AFECON::EXBUFEN
            | AFECON::DACEN,
            true,
        )
        .await;

        // Set CLK configuration
        let clk_config = ClkConfig::default()
            .adcclkdiv(ADCCLKDIV::DIV1) // ADCCLK = 32MHz
            .sysclkdiv(SYSCLKDIV::DIV2) // SYSCLK = 16MHz
            .clk32mhzen(CLK32MHZEN::MHz32);

        self.ad5940.apply_clk_config(&clk_config).await.unwrap();

        // Configure LP DAC and TIA
        self.common_mode_output_enable(false).await.unwrap();

        // Set DSP configuration
        let mut dsp_config = DspConfig::default();
        dsp_config
            .gnpgain(GNPGAIN::Gain1)
            .adc_mux_n(MUXSELN::HsTiaNeg)
            .adc_mux_p(MUXSELP::HsTiaPos)
            .ctiacon(CTIACON::C32)
            .rtiacon(RTIACON::R1k)
            .sinc3osr(SINC3OSR::R4)
            .sinc2osr(SINC2OSR::R178)
            .adcsamplerate(ADCSAMPLERATE::R1_6MHz)
            .dftin_sel(DFTINSEL::GainOffset)
            .dftnum(DFTNUM::Num4096)
            .hanning(true)
            .set_clks(16_000_000, 32_000_000); // Check clk_config: In this case SYSCLK = 16MHz and ADCCLK = 32MHz

        self.ad5940.apply_dsp_config(&dsp_config).await.unwrap();
        self.dsp_config = Some(dsp_config);

        // Set SRAM configuration (cmd and data sram)
        let mut sram_config = SramConfig::default();
        sram_config
            .datafifosrcsel(DATAFIFOSRCSEL::DFT)
            .datafifoen(DATAFIFOEN::Normal)
            .data_size(DATA_MEM_SEL::Size2kB)
            .cmd_mode(CMDMEMMDE::MemoryMode)
            .cmd_size(CMD_MEM_SEL::Size2kB);
        self.ad5940.apply_sram_config(sram_config).await.unwrap();

        // WGCON: set sinus output
        let config_wgcon = WGCON::TYPESEL_SIN.bits();
        self.ad5940.write_reg(Register::WGCON, config_wgcon).await.unwrap();

        Ok(())
    }

    pub async fn calibrate_rtia(&mut self, frequency: u32, wait_time: u32) -> Result<RtiaCalibrationResult, Error> {
        /* CALIBRATION METHOD
        1) Measure the complex voltage V_Rcal across the calibration DUT (Rcal).
        2) Measure the complex voltage V_Rtia across Rtia [HSTIA_P (output) - HSTIA_N].
        3) Note Rtia carries the same current as Rcal; I_Rtia = I_exc = I_Rcal
        4) Implement the equation: Rtia = V_Rtia / I_Rtia 
            --> Rtia = (V_Rtia / V_Rcal) * Rcal
        */

        // self.ad5940.hsr_calibrate_rtia(wait_time).await;
        self.ad5940.wgfcw(frequency).await;
        let wg_amplitude = 2047; // 2047 is the maximum amplitude for a 12-bit DAC --> 1.62V peak-to-peak
        self.ad5940.write_reg(Register::WGAMPLITUDE, wg_amplitude).await.unwrap();

        // Switches for Rcal
        let switch_config = SwitchConfig::default()
            .t9con(T9CON::T9Closed)
            .tmuxcon(TMUXCON::TR1Closed)
            .nmuxcon(NMUXCON::NR1Closed)
            .pmuxcon(PMUXCON::PR0Closed)
            .dmuxcon(DMUXCON::DR0Closed);
        self.ad5940.apply_switch_config(switch_config).await.unwrap();

        // Voltage measurement Rcal
        self.dsp_config.as_mut().unwrap()
            .adc_mux_n(MUXSELN::ExciNNode)
            .adc_mux_p(MUXSELP::ExciPNode);
        self.ad5940.apply_dsp_config(self.dsp_config.as_ref().unwrap()).await.unwrap();

        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON::ADCEN, true).await;
        // self.ad5940.sequencer_wait(16*10).await; // 10 us based on SYSCLK = 16MHz
        Timer::after_micros(10).await;
        self.ad5940.afecon(AFECON::ADCCONVEN | AFECON::DFTEN, true).await;
        // self.ad5940.sequencer_wait(wait_time).await; // Determined above
        // self.ad5940.sequencer_wait(16*20).await; // 20 us based on SYSCLK = 16MHz
        Timer::after_micros((wait_time/16) as u64).await;
        Timer::after_micros(20).await;
        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON:: ADCEN | AFECON::ADCCONVEN | AFECON::DFTEN, false).await;

        let real_raw = self.ad5940.read_reg(Register::DFTREAL).await.unwrap();
        let imag_raw = self.ad5940.read_reg(Register::DFTIMAG).await.unwrap();
        let mut rcal = Complex::new(
            sign_extend_18bit(real_raw) as f32,
            sign_extend_18bit(imag_raw) as f32
        );    

        // Voltage measurement Rtia
        self.dsp_config.as_mut().unwrap()
            .adc_mux_n(MUXSELN::HsTiaNeg)
            .adc_mux_p(MUXSELP::HsTiaPos);
        self.ad5940.apply_dsp_config(self.dsp_config.as_ref().unwrap()).await.unwrap();

        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON::ADCEN, true).await;
        // self.ad5940.sequencer_wait(16*10).await; // 10 us based on SYSCLK = 16MHz
        Timer::after_micros(10).await;
        self.ad5940.afecon(AFECON::ADCCONVEN | AFECON::DFTEN, true).await;
        // self.ad5940.sequencer_wait(wait_time).await; // Determined above
        // self.ad5940.sequencer_wait(16*20).await; // 20 us based on SYSCLK = 16MHz
        Timer::after_micros((wait_time/16) as u64).await;
        Timer::after_micros(20).await;
        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON:: ADCEN | AFECON::ADCCONVEN | AFECON::DFTEN, false).await;

        let real_raw = self.ad5940.read_reg(Register::DFTREAL).await.unwrap();
        let imag_raw = self.ad5940.read_reg(Register::DFTIMAG).await.unwrap();
        let mut rtia = Complex::new(
            sign_extend_18bit(real_raw) as f32,
            sign_extend_18bit(imag_raw) as f32
        );

        // Current is measured inverted in rtia
        rtia.re *= -1.0;
        rtia.im *= -1.0;

        // Impedance imaginary part is measured inverted
        rcal.im *= -1.0;
        rtia.im *= -1.0;

        // Rtia = (V_Rtia / V_Rcal) * Rcal
        let mut temp = rtia/rcal;
        temp = temp.scale(RCAL);

        // Calibration result
        let calibration_result = RtiaCalibrationResult {
            magnitude: temp.norm(), // Magnitude in Ohm
            phase: temp.arg(), // Phase in radians
        };

        Ok(calibration_result)
    }

    async fn calibrate_rtia_multiple<const N: usize>(&mut self, number_of_points: MeasurementPointSet) -> Result<Vec<RtiaCalibrationResult, N>, ImpedanceInitError> {
        let mut results = Vec::<RtiaCalibrationResult, N>::new();
        
        // Set DFT number based on the frequency
        for &frequency in number_of_points.values() {
            // Determine wait time
            let selected_dft_num = match frequency {
                f if f < 10 => DFTNUM::Num16384,
                f if f < 100 => DFTNUM::Num16384,
                f if f < 250 => DFTNUM::Num8192,
                f if f < 1000 => DFTNUM::Num4096,
                f if f < 2500 => DFTNUM::Num2048,
                f if f < 10000 => DFTNUM::Num1024,
                f if f < 25000 => DFTNUM::Num512,
                f if f < 100000 => DFTNUM::Num256,
                _ => DFTNUM::Num128,
            };

            let wait_time: u32;
            
            if let Some(dsp_config) = &mut self.dsp_config {
                dsp_config
                    .dftnum(selected_dft_num);
                
                // Update DFTNUM
                let mut current =  self.ad5940.read_reg(Register::DFTCON).await.unwrap();
                current = DFTNUM::apply(current, selected_dft_num as u32);
                self.ad5940.write_reg(Register::DFTCON, current).await.unwrap();

                wait_time = self.ad5940.sequencer_calculate_wait_time(&dsp_config).await.unwrap();

                // Calibrate Rtia at the given frequency
                let calibration_result = self.calibrate_rtia(frequency, wait_time).await.unwrap();

                results.push(calibration_result).unwrap();
            } else {
                error!("DSP configuration not set, cannot calculate wait time");
                return Err(ImpedanceInitError::DSPNotSet);
            }
        }

        Ok(results)
    }

    pub async fn common_mode_output_enable(&mut self, enable: bool) -> Result<(), Error> {
        // Configure LP DAC and TIA
        let mut lp_config = LpConfig::default();

        if enable {
            lp_config
                .data_reset(true)
                .power_enable(true)
                .data_6bit(31) // Mid-scale for 6-bit DAC
                .common_mode_enable(true)
                .tia_switches_enabled(1 << 5 | 1 << 7 | 1 << 9 | 1 << 13)
                .tia_enable(true)
                .filter_resistor(TIARF::R20k);
            
        } else {
            lp_config
                .power_enable(false)
                .common_mode_enable(false)
                .tia_enable(false);
        }

        self.ad5940.apply_lp_config(&lp_config).await.unwrap();

        Ok(())
    }

    pub async fn init_single_frequency_measurement_2_lead(&mut self, frequency: u32, dft_number: IcdDftNum) -> Result<f32, ImpedanceInitError> {
        // Configure LP DAC and TIA
        self.common_mode_output_enable(false).await.unwrap();

        // Reset FIFO
        self.ad5940.clear_and_enable_fifo().await.unwrap();

        // Set DFT number
        self.dsp_config.as_mut().unwrap().dftnum(dft_number.into_dftnum());
        self.ad5940.apply_dsp_config(self.dsp_config.as_ref().unwrap()).await.unwrap();

        // Configure GPIOs
        self.ad5940.write_reg(Register::GP0CON, 0b10 << 4 |  0b10 << 2 | 0b10).await.unwrap();
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b111).await.unwrap();

        // Calculate wait time between measurement start and stop
        let wait_time;
        let sinus_periods_per_dft;
        if let Some(dsp_config) = &self.dsp_config {
            wait_time = self.ad5940.sequencer_calculate_wait_time(dsp_config).await.unwrap();
            sinus_periods_per_dft = wait_time as f32 / dsp_config.fsys.unwrap() as f32 * frequency as f32;
            // info!("Sinus periods per DFT: {}", sinus_periods_per_dft);
        } else {
            error!("DSP configuration not set, cannot calculate wait time");
            return Err(ImpedanceInitError::DSPNotSet);
        }

        // Configure sequencer
        self.ad5940.sequencer_enable(true).await;

        self.ad5940.wgfcw(frequency).await;
        let wg_amplitude = 2047; // 2047 is the maximum amplitude for a 12-bit DAC --> 1.62V peak-to-peak
        self.ad5940.write_reg(Register::WGAMPLITUDE, wg_amplitude).await.unwrap();

        // Rcal
        let switch_config = SwitchConfig::default()
            .t9con(T9CON::T9Closed)
            .tmuxcon(TMUXCON::TR1Closed)
            .nmuxcon(NMUXCON::NR1Closed)
            .pmuxcon(PMUXCON::PR0Closed)
            .dmuxcon(DMUXCON::DR0Closed);
        self.ad5940.apply_switch_config(switch_config).await.unwrap();
        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON::ADCEN, true).await;
        self.ad5940.sequencer_wait(16*10).await; // 10 us based on SYSCLK = 16MHz
        self.ad5940.afecon(AFECON::ADCCONVEN | AFECON::DFTEN, true).await;
        self.ad5940.sequencer_wait(wait_time).await; // Determined above
        self.ad5940.sequencer_wait(16*20).await; // 20 us based on SYSCLK = 16MHz
        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON:: ADCEN | AFECON::ADCCONVEN | AFECON::DFTEN, false).await;

        // Rz
        let switch_config = SwitchConfig::default()
            .t9con(T9CON::T9Closed)
            .tmuxcon(TMUXCON::T2Closed)
            .nmuxcon(NMUXCON::N2Closed)
            .pmuxcon(PMUXCON::P11Closed)
            .dmuxcon(DMUXCON::D5Closed);
        self.ad5940.apply_switch_config(switch_config).await.unwrap();
        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON::ADCEN, true).await;
        self.ad5940.sequencer_wait(16*10).await; // 10 us based on SYSCLK = 16MHz
        self.ad5940.afecon(AFECON::ADCCONVEN | AFECON::DFTEN, true).await;
        self.ad5940.sequencer_wait(wait_time).await; // Determined above
        self.ad5940.sequencer_wait(16*20).await; // 20 us based on SYSCLK = 16MHz
        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON:: ADCEN | AFECON::ADCCONVEN | AFECON::DFTEN, false).await;

        // Toggle leds
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b010).await.unwrap();
        self.ad5940.sequencer_wait(16 * 1_000).await; // 1ms based on SYSCLK = 16MHz
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b111).await.unwrap();

        self.ad5940.sequencer_enable(false).await;

        // Write sequence to SRAM
        let start_address = 0;
        self.ad5940.sequencer_cmd_write(start_address).await;

        self.ad5940.sequencer_info_configure(0, self.ad5940.seq_len, start_address).await;

        self.start_measurement().await;

        Ok(sinus_periods_per_dft)
    }

    pub async fn init_single_frequency_measurement_4_lead(&mut self, frequency: u32, dft_number: IcdDftNum) -> Result<f32, ImpedanceInitError> {
        // Configure LP DAC and TIA
        self.common_mode_output_enable(true).await.unwrap();

        // Reset FIFO
        self.ad5940.clear_and_enable_fifo().await.unwrap();

        // Set DFT number
        self.dsp_config.as_mut().unwrap().dftnum(dft_number.into_dftnum());
        self.ad5940.apply_dsp_config(self.dsp_config.as_ref().unwrap()).await.unwrap();

        // Configure GPIOs
        self.ad5940.write_reg(Register::GP0CON, 0b10 << 4 |  0b10 << 2 | 0b10).await.unwrap();
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b111).await.unwrap();

        // Calculate wait time between measurement start and stop
        let wait_time;
        let sinus_periods_per_dft;
        if let Some(dsp_config) = &self.dsp_config {
            wait_time = self.ad5940.sequencer_calculate_wait_time(dsp_config).await.unwrap();
            sinus_periods_per_dft = wait_time as f32 / dsp_config.fsys.unwrap() as f32 * frequency as f32;
            // info!("Sinus periods per DFT: {}", sinus_periods_per_dft);
        } else {
            error!("DSP configuration not set, cannot calculate wait time");
            return Err(ImpedanceInitError::DSPNotSet);
        }

        // Calibrate Rtia
        match self.calibrate_rtia(frequency, wait_time).await {
            Ok(calibration_result) => {
                self.rtia_calibrated = RtiaCalibrated::Single(calibration_result);
            },
            Err(e) => {
                error!("Rtia calibration failed: {:?}", e);
            }
        }

        // Configure switches
        let switch_config = SwitchConfig::default()
            .t9con(T9CON::T9Closed)
            .tmuxcon(TMUXCON::T2Closed)
            .nmuxcon(NMUXCON::N2Closed)
            .pmuxcon(PMUXCON::P11Closed)
            .dmuxcon(DMUXCON::D5Closed);
        self.ad5940.apply_switch_config(switch_config).await.unwrap();

        // Configure sequencer
        self.ad5940.sequencer_enable(true).await;

        self.ad5940.wgfcw(frequency).await;
        let wg_amplitude = 2047; // 2047 is the maximum amplitude for a 12-bit DAC --> 1.62V peak-to-peak
        self.ad5940.write_reg(Register::WGAMPLITUDE, wg_amplitude).await.unwrap();

        // Voltage measurement
        self.dsp_config.as_mut().unwrap()
            .adc_mux_n(MUXSELN::AIN3)
            .adc_mux_p(MUXSELP::AIN2);
        // self.dsp_config.as_mut().unwrap()
            // .adc_mux_n(MUXSELN::AIN1)
            // .adc_mux_p(MUXSELP::CE0);
        self.ad5940.apply_dsp_config(self.dsp_config.as_ref().unwrap()).await.unwrap();

        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON::ADCEN, true).await;
        self.ad5940.sequencer_wait(16*10).await; // 10 us based on SYSCLK = 16MHz
        self.ad5940.afecon(AFECON::ADCCONVEN | AFECON::DFTEN, true).await;
        self.ad5940.sequencer_wait(wait_time).await; // Determined above
        self.ad5940.sequencer_wait(16*20).await; // 20 us based on SYSCLK = 16MHz
        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON:: ADCEN | AFECON::ADCCONVEN | AFECON::DFTEN, false).await;

        // Current measurement
        self.dsp_config.as_mut().unwrap()
            .adc_mux_n(MUXSELN::HsTiaNeg)
            .adc_mux_p(MUXSELP::HsTiaPos);
        self.ad5940.apply_dsp_config(self.dsp_config.as_ref().unwrap()).await.unwrap();

        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON::ADCEN, true).await;
        self.ad5940.sequencer_wait(16*10).await; // 10 us based on SYSCLK = 16MHz
        self.ad5940.afecon(AFECON::ADCCONVEN | AFECON::DFTEN, true).await;
        self.ad5940.sequencer_wait(wait_time).await; // Determined above
        self.ad5940.sequencer_wait(16*20).await; // 20 us based on SYSCLK = 16MHz
        self.ad5940.afecon(AFECON::WAVEGENEN | AFECON:: ADCEN | AFECON::ADCCONVEN | AFECON::DFTEN, false).await;

        // Toggle leds
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b010).await.unwrap();
        self.ad5940.sequencer_wait(16 * 1_000).await; // 1ms based on SYSCLK = 16MHz
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b111).await.unwrap();

        self.ad5940.sequencer_enable(false).await;

        // Write sequence to SRAM
        let start_address = 0;
        self.ad5940.sequencer_cmd_write(start_address).await;

        self.ad5940.sequencer_info_configure(0, self.ad5940.seq_len, start_address).await;

        self.start_measurement().await;

        Ok(sinus_periods_per_dft)
    }

    pub async fn init_sweep_frequency_measurement_2_lead<const N: usize>(&mut self, number_of_points: MeasurementPointSet) -> Result<heapless::Vec<f32, N>, ImpedanceInitError> {
        // Configure LP DAC and TIA
        self.common_mode_output_enable(false).await.unwrap();

        // Create vector to store the periods per DFT for each frequency
        let mut periods_per_dft_vec = heapless::Vec::<f32, N>::new();

        // Reset FIFO
        self.ad5940.clear_and_enable_fifo().await.unwrap();

        // Configure GPIOs
        self.ad5940.write_reg(Register::GP0CON, 0b10 << 4 |  0b10 << 2 | 0b10).await.unwrap();
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b111).await.unwrap();

        // Configure sequencer
        self.ad5940.sequencer_enable(true).await;

        // Set DFT number based on the frequency
        for &frequency in number_of_points.values() {    
            // Determine wait time
            let selected_dft_num = match frequency {
                f if f < 10 => DFTNUM::Num16384,
                f if f < 100 => DFTNUM::Num16384,
                f if f < 250 => DFTNUM::Num8192,
                f if f < 1000 => DFTNUM::Num4096,
                f if f < 2500 => DFTNUM::Num2048,
                f if f < 10000 => DFTNUM::Num1024,
                f if f < 25000 => DFTNUM::Num512,
                f if f < 100000 => DFTNUM::Num256,
                _ => DFTNUM::Num128,
            };

            let wait_time: u32;
            
            if let Some(dsp_config) = &mut self.dsp_config {
                dsp_config
                    .dftnum(selected_dft_num);
                
                // Update DFTNUM
                let mut current =  self.ad5940.read_reg(Register::DFTCON).await.unwrap();
                current = DFTNUM::apply(current, selected_dft_num as u32);
                self.ad5940.write_reg(Register::DFTCON, current).await.unwrap();

                wait_time = self.ad5940.sequencer_calculate_wait_time(&dsp_config).await.unwrap();
                let periods_per_dft = wait_time as f32 / dsp_config.fsys.unwrap() as f32 * frequency as f32;
                
                periods_per_dft_vec.push(periods_per_dft).unwrap();
                // info!("{}Hz: Sinus periods per DFT: {}", frequency, periods_per_dft);

            } else {
                error!("DSP configuration not set, cannot calculate wait time");
                return Err(ImpedanceInitError::DSPNotSet);
            }
            
            // Set frequency    
            self.ad5940.wgfcw(frequency as u32).await;
            let wg_amplitude = 2047; // 2047 is the maximum amplitude for a 12-bit DAC --> 1.62V peak-to-peak
            self.ad5940.write_reg(Register::WGAMPLITUDE, wg_amplitude).await.unwrap();

            // Rcal
            let switch_config = SwitchConfig::default()
                .t9con(T9CON::T9Closed)
                .tmuxcon(TMUXCON::TR1Closed)
                .nmuxcon(NMUXCON::NR1Closed)
                .pmuxcon(PMUXCON::PR0Closed)
                .dmuxcon(DMUXCON::DR0Closed);
            self.ad5940.apply_switch_config(switch_config).await.unwrap();
            self.ad5940.afecon(AFECON::WAVEGENEN | AFECON::ADCEN, true).await;
            self.ad5940.sequencer_wait(16*10).await; // 10 us based on SYSCLK = 16MHz
            self.ad5940.afecon(AFECON::ADCCONVEN | AFECON::DFTEN, true).await;
            self.ad5940.sequencer_wait(wait_time).await; // Determined above
            self.ad5940.sequencer_wait(16*20).await; // 20 us based on SYSCLK = 16MHz
            self.ad5940.afecon(AFECON::WAVEGENEN | AFECON:: ADCEN | AFECON::ADCCONVEN | AFECON::DFTEN, false).await;

            // Rz
            let switch_config = SwitchConfig::default()
                .t9con(T9CON::T9Closed)
                .tmuxcon(TMUXCON::T2Closed)
                .nmuxcon(NMUXCON::N2Closed)
                .pmuxcon(PMUXCON::P11Closed)
                .dmuxcon(DMUXCON::D5Closed);
            self.ad5940.apply_switch_config(switch_config).await.unwrap();
            self.ad5940.afecon(AFECON::WAVEGENEN | AFECON::ADCEN, true).await;
            self.ad5940.sequencer_wait(16*10).await; // 10 us based on SYSCLK = 16MHz
            self.ad5940.afecon(AFECON::ADCCONVEN | AFECON::DFTEN, true).await;
            self.ad5940.sequencer_wait(wait_time).await; // Determined above
            self.ad5940.sequencer_wait(16*20).await; // 20 us based on SYSCLK = 16MHz
            self.ad5940.afecon(AFECON::WAVEGENEN | AFECON:: ADCEN | AFECON::ADCCONVEN | AFECON::DFTEN, false).await;
        }

        // Toggle leds
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b010).await.unwrap();
        self.ad5940.sequencer_wait(16 * 1_000).await; // 1ms based on SYSCLK = 16MHz
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b111).await.unwrap();

        self.ad5940.sequencer_enable(false).await;

        // Write sequence to SRAM
        let start_address = 0;
        self.ad5940.sequencer_cmd_write(start_address).await;

        self.ad5940.sequencer_info_configure(0, self.ad5940.seq_len, start_address).await;

        self.start_measurement().await;

        Ok(periods_per_dft_vec)
    }

    pub async fn init_sweep_frequency_measurement_4_lead<const N: usize>(&mut self, number_of_points: MeasurementPointSet) -> Result<heapless::Vec<f32, N>, ImpedanceInitError> {
        // Configure LP DAC and TIA
        self.common_mode_output_enable(true).await.unwrap();

        // Create vector to store the periods per DFT for each frequency
        let mut periods_per_dft_vec = heapless::Vec::<f32, N>::new();

        // Calibrate Rtia
        match number_of_points {
            MeasurementPointSet::Eight => {
                let results = self.calibrate_rtia_multiple::<8>(number_of_points).await.unwrap();
                self.rtia_calibrated = RtiaCalibrated::Vec8(results);
            },
            MeasurementPointSet::Eighteen => {
                let results = self.calibrate_rtia_multiple::<18>(number_of_points).await.unwrap();
                self.rtia_calibrated = RtiaCalibrated::Vec18(results);
            },
        }

        // Reset FIFO
        self.ad5940.clear_and_enable_fifo().await.unwrap();

        // Configure GPIOs
        self.ad5940.write_reg(Register::GP0CON, 0b10 << 4 |  0b10 << 2 | 0b10).await.unwrap();
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b111).await.unwrap();

        // Configure sequencer
        self.ad5940.sequencer_enable(true).await;

        // Set DFT number based on the frequency
        for &frequency in number_of_points.values() {    
            // Determine wait time
            let selected_dft_num = match frequency {
                f if f < 10 => DFTNUM::Num16384,
                f if f < 100 => DFTNUM::Num16384,
                f if f < 250 => DFTNUM::Num8192,
                f if f < 1000 => DFTNUM::Num4096,
                f if f < 2500 => DFTNUM::Num2048,
                f if f < 10000 => DFTNUM::Num1024,
                f if f < 25000 => DFTNUM::Num512,
                f if f < 100000 => DFTNUM::Num256,
                _ => DFTNUM::Num128,
            };

            let wait_time: u32;
            
            if let Some(dsp_config) = &mut self.dsp_config {
                dsp_config
                    .dftnum(selected_dft_num);
                
                // Update DFTNUM
                let mut current =  self.ad5940.read_reg(Register::DFTCON).await.unwrap();
                current = DFTNUM::apply(current, selected_dft_num as u32);
                self.ad5940.write_reg(Register::DFTCON, current).await.unwrap();

                wait_time = self.ad5940.sequencer_calculate_wait_time(&dsp_config).await.unwrap();
                let periods_per_dft = wait_time as f32 / dsp_config.fsys.unwrap() as f32 * frequency as f32;
                
                periods_per_dft_vec.push(periods_per_dft).unwrap();
                // info!("{}Hz: Sinus periods per DFT: {}", frequency, periods_per_dft);
            } else {
                error!("DSP configuration not set, cannot calculate wait time");
                return Err(ImpedanceInitError::DSPNotSet);
            }
            
            // Configure switches
            let switch_config = SwitchConfig::default()
                .t9con(T9CON::T9Closed)
                .tmuxcon(TMUXCON::T2Closed)
                .nmuxcon(NMUXCON::N2Closed)
                .pmuxcon(PMUXCON::P11Closed)
                .dmuxcon(DMUXCON::D5Closed);
            self.ad5940.apply_switch_config(switch_config).await.unwrap();
            
            // Set frequency    
            self.ad5940.wgfcw(frequency as u32).await;
            let wg_amplitude = 2047; // 2047 is the maximum amplitude for a 12-bit DAC --> 1.62V peak-to-peak
            self.ad5940.write_reg(Register::WGAMPLITUDE, wg_amplitude).await.unwrap();

            // Voltage measurement
            self.dsp_config.as_mut().unwrap()
                .adc_mux_n(MUXSELN::AIN3)
                .adc_mux_p(MUXSELP::AIN2);
            // self.dsp_config.as_mut().unwrap()
                // .adc_mux_n(MUXSELN::AIN1)
                // .adc_mux_p(MUXSELP::CE0);
            self.ad5940.apply_dsp_config(self.dsp_config.as_ref().unwrap()).await.unwrap();

            self.ad5940.afecon(AFECON::WAVEGENEN | AFECON::ADCEN, true).await;
            self.ad5940.sequencer_wait(16*10).await; // 10 us based on SYSCLK = 16MHz
            self.ad5940.afecon(AFECON::ADCCONVEN | AFECON::DFTEN, true).await;
            self.ad5940.sequencer_wait(wait_time).await; // Determined above
            self.ad5940.sequencer_wait(16*20).await; // 20 us based on SYSCLK = 16MHz
            self.ad5940.afecon(AFECON::WAVEGENEN | AFECON:: ADCEN | AFECON::ADCCONVEN | AFECON::DFTEN, false).await;

            // Current measurement
            self.dsp_config.as_mut().unwrap()
                .adc_mux_n(MUXSELN::HsTiaNeg)
                .adc_mux_p(MUXSELP::HsTiaPos);
            self.ad5940.apply_dsp_config(self.dsp_config.as_ref().unwrap()).await.unwrap();

            self.ad5940.afecon(AFECON::WAVEGENEN | AFECON::ADCEN, true).await;
            self.ad5940.sequencer_wait(16*10).await; // 10 us based on SYSCLK = 16MHz
            self.ad5940.afecon(AFECON::ADCCONVEN | AFECON::DFTEN, true).await;
            self.ad5940.sequencer_wait(wait_time).await; // Determined above
            self.ad5940.sequencer_wait(16*20).await; // 20 us based on SYSCLK = 16MHz
            self.ad5940.afecon(AFECON::WAVEGENEN | AFECON:: ADCEN | AFECON::ADCCONVEN | AFECON::DFTEN, false).await;
        }

        // Toggle leds
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b010).await.unwrap();
        self.ad5940.sequencer_wait(16 * 1_000).await; // 1ms based on SYSCLK = 16MHz
        self.ad5940.write_reg(Register::SYNCEXTDEVICE, 0b111).await.unwrap();

        self.ad5940.sequencer_enable(false).await;

        // Write sequence to SRAM
        let start_address = 0;
        self.ad5940.sequencer_cmd_write(start_address).await;

        self.ad5940.sequencer_info_configure(0, self.ad5940.seq_len, start_address).await;

        self.start_measurement().await;

        Ok(periods_per_dft_vec)
    }

    pub async fn start_measurement(&mut self) {
        self.ad5940.sequencer_trigger(0).await;
    }

    pub async fn get_fifo_count(&mut self) ->  Result<u32, Error> {
        self.ad5940.get_fifo_count().await
    }

    pub async fn read_fifo(&mut self, data: &mut [u32]) -> Result<(), Error> {
        self.ad5940.read_fifo(data).await
    }
}

// Task to wait for the interrupt from the AD5940 and read the FIFO
// Depending on the running mode, it will read either a single frequency or multiple frequencies
// Result is sent via channel to the communication task which then sends it via USB to the host
#[embassy_executor::task]
pub async fn impedance_setup_readout_task(mut pin: ExtiInput<'static>, impedance_setup: &'static ImpedanceSetupType) {
    loop {
        // Wait untill sequence is done
        pin.wait_for_rising_edge().await;

        // Lock the impedance setup
        let mut impedance_setup = impedance_setup.lock().await;

        // Trigger the sequencer again
        if impedance_setup.running_mode == RunningMode::SingleFrequency2Lead || impedance_setup.running_mode == RunningMode::SingleFrequency4Lead || impedance_setup.running_mode == RunningMode::SweepFrequency2Lead(MeasurementPointSet::Eight) || impedance_setup.running_mode == RunningMode::SweepFrequency2Lead(MeasurementPointSet::Eighteen) || impedance_setup.running_mode == RunningMode::SweepFrequency4Lead(MeasurementPointSet::Eight) || impedance_setup.running_mode == RunningMode::SweepFrequency4Lead(MeasurementPointSet::Eighteen) {
            impedance_setup.start_measurement().await;
        }

        // Read the FIFO count
        let count = impedance_setup.get_fifo_count().await.unwrap() as usize;
        // info!("FIFOCNTSTA: {}", count);

        match impedance_setup.running_mode {
            RunningMode::None => {
                continue; // Skip processing if not running
            }
            RunningMode::SingleFrequency2Lead => {
                if count >= 4 {
                    let mut data: [u32; 4] = [0; 4];

                    impedance_setup.read_fifo(data.as_mut_slice()).await.unwrap();

                    let result = calculate_impedance_2_lead(data);
                    
                    // Log
                    // info!("Impedance: Magnitude = {} Ω, Phase = {} rad", result.magnitude, result.phase);

                    let data = SingleImpedanceOutput {
                        magnitude: result.magnitude,
                        phase: result.phase,
                    };

                    IMPEDANCE_CHANNEL_SINGLE.try_send(data).ok();
                }
            }
            RunningMode::SingleFrequency4Lead => {
                if count >= 4 {
                    let mut data: [u32; 4] = [0; 4];

                    impedance_setup.read_fifo(data.as_mut_slice()).await.unwrap();
                    match impedance_setup.rtia_calibrated {
                        RtiaCalibrated::Single(ref cal) => {
                            let result = calculate_impedance_4_lead(data, cal);

                            let data = SingleImpedanceOutput {
                                magnitude: result.magnitude,
                                phase: result.phase
                            };

                            IMPEDANCE_CHANNEL_SINGLE.try_send(data).ok();
                        }
                        _ => {
                            error!("Rtia not (correctly) calibrated, cannot compute impedance");
                            continue;
                        },
                    }
                }
            }
            RunningMode::SweepFrequency2Lead(points) => {
                // Each frequency point produces 4 samples (DFT real/imag for Rcal and Rz)
                let required_count = points.len() * 4;
                
                if count >= required_count {
                    // Use stack-allocated array
                    let mut data: [u32; 72] = [0; 72]; // max needed size (18*4=72)
                    let data_slice = &mut data[..required_count];
                    impedance_setup.read_fifo(data_slice).await.unwrap();

                    // Output structure
                    let mut impedance_output = SweepImpedanceOutput {
                        points,
                        magnitudes_8: Vec::new(),
                        phases_8: Vec::new(),
                        magnitudes_18: Vec::new(),
                        phases_18: Vec::new(),
                    };

                    // Take 4 samples per frequency point
                    for chunk in data_slice.chunks(4) {
                        let result = calculate_impedance_2_lead(chunk.try_into().unwrap());

                        match points {
                            MeasurementPointSet::Eight => {
                                impedance_output.magnitudes_8.push(result.magnitude).ok();
                                impedance_output.phases_8.push(result.phase).ok();
                            }
                            MeasurementPointSet::Eighteen => {
                                impedance_output.magnitudes_18.push(result.magnitude).ok();
                                impedance_output.phases_18.push(result.phase).ok();
                            }
                        }
                    }

                    IMPEDANCE_CHANNEL_SWEEP.try_send(impedance_output).ok();
                }
                continue;
            }
            RunningMode::SweepFrequency4Lead(points) => {
                 // Each frequency point produces 4 samples (DFT real/imag for Rcal and Rz)
                let required_count = points.len() * 4;
                
                if count >= required_count {
                    // Use stack-allocated array
                    let mut data: [u32; 72] = [0; 72]; // max needed size (18*4=72)
                    let data_slice = &mut data[..required_count];
                    impedance_setup.read_fifo(data_slice).await.unwrap();

                    // Output structure
                    let mut impedance_output = SweepImpedanceOutput {
                        points,
                        magnitudes_8: Vec::new(),
                        phases_8: Vec::new(),
                        magnitudes_18: Vec::new(),
                        phases_18: Vec::new(),
                    };
                    
                    match (&points, &impedance_setup.rtia_calibrated) {
                        (MeasurementPointSet::Eight, RtiaCalibrated::Vec8(cal_vec)) => {
                            // Take 4 samples per frequency point
                            for (i, chunk) in data_slice.chunks(4).enumerate() {
                                let cal = &cal_vec[i];
                                let result = calculate_impedance_4_lead(chunk.try_into().unwrap(), cal);

                                impedance_output.magnitudes_8.push(result.magnitude).ok();
                                impedance_output.phases_8.push(result.phase).ok();
                            }
                        },
                        (MeasurementPointSet::Eighteen, RtiaCalibrated::Vec18(cal_vec)) => {
                            // Take 4 samples per frequency point
                            for (i, chunk) in data_slice.chunks(4).enumerate() {
                                let cal = &cal_vec[i];
                                let result = calculate_impedance_4_lead(chunk.try_into().unwrap(), cal);

                                impedance_output.magnitudes_18.push(result.magnitude).ok();
                                impedance_output.phases_18.push(result.phase).ok();
                            }
                        },
                        _ => {
                            error!("Rtia not (correctly) calibrated, cannot compute impedance");
                            continue;
                        }
                    }

                    IMPEDANCE_CHANNEL_SWEEP.try_send(impedance_output).ok();
                }
                continue;
              
            }
        }
    }
}

#[derive(Debug)]
pub struct ImpedanceResult {
    pub magnitude: f32,
    pub phase: f32,
}

/// Convert raw 18-bit 2's complement value to signed i32
fn sign_extend_18bit(val: u32) -> i32 {
    ((val << 14) as i32) >> 14
}

/// Calculate magnitude and phase of Rz using Rcal reference
pub fn calculate_impedance_2_lead(data: [u32; 4]) -> ImpedanceResult {
    let mut dft_rcal = Complex::new(
        sign_extend_18bit(data[0]) as f32, 
        sign_extend_18bit(data[1]) as f32
    );
    let mut dft_rz = Complex::new(
        sign_extend_18bit(data[2]) as f32, 
        sign_extend_18bit(data[3]) as f32
    );

    // Current is measured inverted in rtia
    // --> Will cancel out in the impedance calculation
    // dft_rcal.re *= -1.0;
    // dft_rcal.im *= -1.0;

    // dft_rz.re *= -1.0;
    // dft_rz.im *= -1.0;

    // Impedance imaginary part is measured inverted
    dft_rcal.im *= -1.0;
    dft_rz.im *= -1.0;

    let (rcal_mag, rcal_phase) = dft_rcal.to_polar();
    let (rz_mag, rz_phase) = dft_rz.to_polar();

    let magnitude = (rcal_mag / rz_mag) * RCAL;
    let phase = rcal_phase - rz_phase;

    ImpedanceResult { magnitude, phase }
}

/// Calculate magnitude and phase of Rz using the calibrated Rtia
pub fn calculate_impedance_4_lead(data: [u32; 4], rtia: &RtiaCalibrationResult) -> ImpedanceResult {
    let mut dft_volt = Complex::new(
        sign_extend_18bit(data[0]) as f32, 
        sign_extend_18bit(data[1]) as f32
    );
    let mut dft_curr = Complex::new(
        sign_extend_18bit(data[2]) as f32, 
        sign_extend_18bit(data[3]) as f32
    );

    // Current is measured inverted in rtia
    dft_curr.re *= -1.0;
    dft_curr.im *= -1.0;

    // Impedance imaginary part is measured inverted
    dft_volt.im *= -1.0;
    dft_curr.im *= -1.0;

    let (volt_mag, volt_phase) = dft_volt.to_polar();
    let (curr_mag, curr_phase) = dft_curr.to_polar();

    // Calculate impedance using calibrated Rtia
    let magnitude = volt_mag / curr_mag * rtia.magnitude;
    let phase = volt_phase - curr_phase + rtia.phase;

    let data = ImpedanceResult {
        magnitude,
        phase
    };

    data
}