#![no_std]
#![no_main]

use defmt::info;
use embassy_executor::Spawner;
use embassy_stm32::exti::ExtiInput;
use embassy_sync::mutex::Mutex;
use embassy_time::{Timer, Duration};
use embassy_futures::{select::select, select::Either};
use embassy_stm32::gpio::{Level, Output, Speed};
use embassy_stm32::{i2c, spi, Config};
use embassy_stm32::time::Hertz;

use {defmt_rtt as _, panic_probe as _};

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

use bioz_icd_rs::ImpedanceOutput;

mod ad5940;
use ad5940::AD5940;

// mod adg2128;
// use adg2128::State;

// mod electrodes;
// use electrodes::{Electrodes, Electrode, AD5940Pin};

mod ad5940_registers;

use embassy_stm32::usb::Driver;
use embassy_stm32::{bind_interrupts, peripherals, usb};

mod communication;
use communication::{init_communication, LED_FREQUENCY_SIGNAL};

use impedance::IMPEDANCE_CHANNEL;

mod impedance;
use impedance::{ImpedanceSetup, ImpedanceSetupType, IMPEDANCE_SETUP};

mod icd_mapping;

bind_interrupts!(struct Irqs {
    USB_DRD_FS => usb::InterruptHandler<peripherals::USB>;
});

#[embassy_executor::main]
async fn main(spawner: Spawner) {

    let mut config = Config::default();
    {
        use embassy_stm32::rcc::*;
        config.rcc.hsi = None;
        config.rcc.hsi48 = Some(Hsi48Config { sync_from_usb: true }); // needed for USB
        config.rcc.hse = Some(Hse {
            freq: Hertz(24_000_000),
            mode: HseMode::Oscillator,
        });
        config.rcc.pll1 = Some(Pll {
            source: PllSource::HSE,
            prediv: PllPreDiv::DIV6,
            mul: PllMul::MUL125,
            divp: Some(PllDiv::DIV2),
            divq: Some(PllDiv::DIV2),
            divr: Some(PllDiv::DIV2),
        });
        config.rcc.sys = Sysclk::PLL1_P;
        config.rcc.mux.usbsel = mux::Usbsel::HSI48;
    }    

    let p = embassy_stm32::init(config);
    info!("Hello World!");

    // let mut led = Output::new(p.PA5, Level::High, Speed::Low);

    // Set up SPI for AD5940
    let cs = Output::new(p.PC9, Level::High, Speed::Low);
    let rst = Output::new(p.PB3, Level::High, Speed::Low);

    let spi = spi::Spi::new_blocking(
        p.SPI1,
        p.PA5, // SCK
        p.PA7, // MOSI
        p.PA6, // MISO
        spi::Config::default()
    );

    let mut ad5940 = AD5940::new(spi, cs, rst);
    ad5940.reset().await.unwrap();
    Timer::after_millis(1).await;
    ad5940.system_init().await.unwrap();

    // ad5940.init_temperature().await.unwrap();
    // ad5940.init_waveform().await.unwrap();
    let mut impedance_setup = ImpedanceSetup::new(ad5940);
    impedance_setup.init().await.unwrap();
    let mut impedance_setup = IMPEDANCE_SETUP.init(Mutex::new(impedance_setup));
    
    // impedance_setup.lock().await.init_single_frequency_measurement().await;

    // // Set up I2C for ADG2128
    // let i2c = i2c::I2c::new_blocking(
    //     p.I2C1, 
    //     p.PB6, 
    //     p.PB7, 
    //     Hertz(400_000), 
    //     i2c::Config::default()
    // );

    // // Initialize electrodes
    // let mut electrodes = Electrodes::new(i2c);
    // electrodes.reset_all();
    // electrodes.set(Electrode::E1, AD5940Pin::CE0, State::ENABLED);
    // electrodes.set(Electrode::E3, AD5940Pin::AIN1, State::ENABLED);
    // // electrodes.set(Electrode::E12, AD5940Pin::RE0, State::ENABLED);

    // Turn on the green LED
    // ad5940.write_reg_raw(0x0000_0004, 1 << 1).await.unwrap(); 
    // ad5940.write_reg_raw(0x0000_001C, 1 << 1).await.unwrap(); 

    // Create USB driver and start postcard-rpc server
    let driver = Driver::new(p.USB, Irqs, p.PA12, p.PA11);
    init_communication(driver, impedance_setup, spawner);

    // Green led task
    // spawner.must_spawn(green_led(led));

    // Trigger the sequencer
    // impedance_setup.start_measurement().await;

    // Set up interrupt at GPIO for AD5940
    let ad5940_gpio_0 = ExtiInput::new(p.PC8, p.EXTI8, embassy_stm32::gpio::Pull::Up);
    spawner.must_spawn(impedance_setup_readout_task(ad5940_gpio_0, impedance_setup));

    loop {
        // Read chip id
        // let chip_id = ad5940.get_chipid().await;
        // info!("Chip ID: 0x{:04X}", chip_id);
        
        // Read temperature
        // let temp = ad5940.get_temperature().await.unwrap();
        // info!("Temperature: {}°C", temp);

        // let result = electrodes.get_all();
        // info!("Electrodes states: {:?}", result);

        Timer::after_secs(1).await;
    }
}

#[embassy_executor::task]
async fn green_led(mut led: Output<'static>) {
    let mut delay = Duration::from_millis(500);

    loop {
        // Wait for either a frequency change or a timer timeout
        match select(LED_FREQUENCY_SIGNAL.wait(), Timer::after(delay)).await {
            Either::First(frequency) => {
                if frequency > 0.0 {
                    // Avoid divide-by-zero or negative delay
                    let millis = (1000.0 / frequency) as u64 / 2;
                    delay = Duration::from_millis(millis.max(1)); // enforce minimum delay of 1ms
                }
            }
            Either::Second(_) => {
                // Timer expired, proceed to toggle
            }
        }

        led.toggle();
    }
}

#[embassy_executor::task]
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 {
            impedance_setup.start_measurement().await;
        }

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

        if count >= 4 {
            let mut data: [u32; 4] = [0; 4];

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

            let result = calculate_impedance(data);

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

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

            IMPEDANCE_CHANNEL.try_send(data).ok();
        }
    }
}


extern crate libm;

#[derive(Debug, Clone, Copy)]
pub struct Complex {
    real: i32,
    imag: i32,
}

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

// Example Rcal value (Ohms)
const RCAL_VAL: f32 = 1000.0;

/// 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(data: [u32; 4]) -> ImpedanceResult {
    let mut signed_data = [0i32; 4];
    for (i, &val) in data.iter().enumerate() {
        signed_data[i] = sign_extend_18bit(val);
    }

    let dft_rcal = Complex {
        real: signed_data[0],
        imag: signed_data[1],
    };

    let dft_rz = Complex {
        real: signed_data[2],
        imag: signed_data[3],
    };

    let rcal_mag = libm::sqrtf((dft_rcal.real as f32) * (dft_rcal.real as f32)
                             + (dft_rcal.imag as f32) * (dft_rcal.imag as f32));

    let rz_mag = libm::sqrtf((dft_rz.real as f32) * (dft_rz.real as f32)
                           + (dft_rz.imag as f32) * (dft_rz.imag as f32));

    let rcal_phase = libm::atan2f(-(dft_rcal.imag as f32), dft_rcal.real as f32);
    let rz_phase = libm::atan2f(-(dft_rz.imag as f32), dft_rz.real as f32);

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

    ImpedanceResult { magnitude, phase }
}