#![no_std]

use bit_field::BitField;
use embedded_hal::{blocking::delay::DelayMs, digital::v2::OutputPin};

/// A device driver for the AD9959 direct digital synthesis (DDS) chip.
///
/// This chip provides four independently controllable digital-to-analog output sinusoids with
/// configurable phase, amplitude, and frequency. All channels are inherently synchronized as they
/// are derived off a common system clock.
///
/// The chip contains a configurable PLL and supports system clock frequencies up to 500 MHz.
///
/// The chip supports a number of serial interfaces to improve data throughput, including normal,
/// dual, and quad SPI configurations.
pub struct Ad9959<INTERFACE, DELAY, UPDATE> {
    interface: INTERFACE,
    delay: DELAY,
    reference_clock_frequency: f32,
    system_clock_multiplier: u8,
    io_update: UPDATE,
}

/// A trait that allows a HAL to provide a means of communicating with the AD9959.
pub trait Interface {
    type Error;

    fn configure_mode(&mut self, mode: Mode) -> Result<(), Self::Error>;

    fn write(&mut self, addr: u8, data: &[u8]) -> Result<(), Self::Error>;

    fn read(&mut self, addr: u8, dest: &mut [u8]) -> Result<(), Self::Error>;
}

/// Indicates various communication modes of the DDS. The value of this enumeration is equivalent to
/// the configuration bits of the DDS CSR register.
#[derive(Copy, Clone, PartialEq)]
pub enum Mode {
    SingleBitTwoWire = 0b00,
    SingleBitThreeWire = 0b01,
    TwoBitSerial = 0b10,
    FourBitSerial = 0b11,
}

/// The configuration registers within the AD9959 DDS device. The values of each register are
/// equivalent to the address.
pub enum Register {
    CSR = 0x00,
    FR1 = 0x01,
    FR2 = 0x02,
    CFR = 0x03,
    CFTW0 = 0x04,
    CPOW0 = 0x05,
    ACR = 0x06,
    LSRR = 0x07,
    RDW = 0x08,
    FDW = 0x09,
    CW1 = 0x0a,
    CW2 = 0x0b,
    CW3 = 0x0c,
    CW4 = 0x0d,
    CW5 = 0x0e,
    CW6 = 0x0f,
    CW7 = 0x10,
    CW8 = 0x11,
    CW9 = 0x12,
    CW10 = 0x13,
    CW11 = 0x14,
    CW12 = 0x15,
    CW13 = 0x16,
    CW14 = 0x17,
    CW15 = 0x18,
}

/// Specifies an output channel of the AD9959 DDS chip.
pub enum Channel {
    One = 0,
    Two = 1,
    Three = 2,
    Four = 3,
}

/// Possible errors generated by the AD9959 driver.
#[derive(Debug)]
pub enum Error {
    Interface,
    Check,
    Bounds,
    Pin,
    Frequency,
}

impl<PinE, INTERFACE, DELAY, UPDATE> Ad9959<INTERFACE, DELAY, UPDATE>
where
    INTERFACE: Interface,
    DELAY: DelayMs<u8>,
    UPDATE: OutputPin<Error = PinE>,
{
    /// Construct and initialize the DDS.
    ///
    /// Args:
    /// * `interface` - An interface to the DDS.
    /// * `reset_pin` - A pin connected to the DDS reset input.
    /// * `io_update` - A pin connected to the DDS io_update input.
    /// * `delay` - A delay implementation for blocking operation for specific amounts of time.
    /// * `desired_mode` - The desired communication mode of the interface to the DDS.
    /// * `clock_frequency` - The clock frequency of the reference clock input.
    /// * `multiplier` - The desired clock multiplier for the system clock. This multiplies
    ///   `clock_frequency` to generate the system clock.
    pub fn new<RST>(
        interface: INTERFACE,
        reset_pin: &mut RST,
        io_update: UPDATE,
        delay: DELAY,
        desired_mode: Mode,
        clock_frequency: f32,
        multiplier: u8,
    ) -> Result<Self, Error>
    where
        RST: OutputPin,
    {
        let mut ad9959 = Ad9959 {
            interface,
            io_update,
            delay,
            reference_clock_frequency: clock_frequency,
            system_clock_multiplier: 1,
        };

        ad9959.io_update.set_low().or_else(|_| Err(Error::Pin))?;

        // Reset the AD9959
        reset_pin.set_high().or_else(|_| Err(Error::Pin))?;

        // Delay for a clock cycle to allow the device to reset.
        ad9959
            .delay
            .delay_ms((1000.0 / clock_frequency as f32) as u8);

        reset_pin.set_low().or_else(|_| Err(Error::Pin))?;

        ad9959
            .interface
            .configure_mode(Mode::SingleBitTwoWire)
            .map_err(|_| Error::Interface)?;

        // Program the interface configuration in the AD9959. Default to all channels enabled.
        let mut csr: [u8; 1] = [0xF0];
        csr[0].set_bits(1..3, desired_mode as u8);
        ad9959
            .interface
            .write(Register::CSR as u8, &csr)
            .map_err(|_| Error::Interface)?;

        // Latch the configuration registers to make them active.
        ad9959.latch_configuration()?;

        ad9959
            .interface
            .configure_mode(desired_mode)
            .map_err(|_| Error::Interface)?;

        // Read back the CSR to ensure it specifies the mode correctly.
        let mut updated_csr: [u8; 1] = [0];
        ad9959
            .interface
            .read(Register::CSR as u8, &mut updated_csr)
            .map_err(|_| Error::Interface)?;
        if updated_csr[0] != csr[0] {
            return Err(Error::Check);
        }

        // Set the clock frequency to configure the device as necessary.
        ad9959.configure_system_clock(clock_frequency, multiplier)?;
        Ok(ad9959)
    }

    /// Latch the DDS configuration to ensure it is active on the output channels.
    fn latch_configuration(&mut self) -> Result<(), Error> {
        self.io_update.set_high().or_else(|_| Err(Error::Pin))?;
        // The SYNC_CLK is 1/4 the system clock frequency. The IO_UPDATE pin must be latched for one
        // full SYNC_CLK pulse to register. For safety, we latch for 5 here.
        self.delay
            .delay_ms((5000.0 / self.system_clock_frequency()) as u8);
        self.io_update.set_low().or_else(|_| Err(Error::Pin))?;

        Ok(())
    }

    /// Configure the internal system clock of the chip.
    ///
    /// Arguments:
    /// * `reference_clock_frequency` - The reference clock frequency provided to the AD9959 core.
    /// * `multiplier` - The frequency multiplier of the system clock. Must be 1 or 4-20.
    ///
    /// Returns:
    /// The actual frequency configured for the internal system clock.
    pub fn configure_system_clock(
        &mut self,
        reference_clock_frequency: f32,
        multiplier: u8,
    ) -> Result<f64, Error> {
        self.reference_clock_frequency = reference_clock_frequency;

        if multiplier != 1 && (multiplier > 20 || multiplier < 4) {
            return Err(Error::Bounds);
        }

        let frequency =
            multiplier as f64 * self.reference_clock_frequency as f64;
        if frequency > 500_000_000.0f64 {
            return Err(Error::Frequency);
        }

        // TODO: Update / disable any enabled channels?
        let mut fr1: [u8; 3] = [0, 0, 0];
        self.interface
            .read(Register::FR1 as u8, &mut fr1)
            .map_err(|_| Error::Interface)?;
        fr1[0].set_bits(2..=6, multiplier);

        let vco_range = frequency > 255e6;
        fr1[0].set_bit(7, vco_range);

        self.interface
            .write(Register::FR1 as u8, &fr1)
            .map_err(|_| Error::Interface)?;
        self.system_clock_multiplier = multiplier;

        Ok(self.system_clock_frequency())
    }

    /// Get the current reference clock frequency in Hz.
    pub fn get_reference_clock_frequency(&self) -> f32 {
        self.reference_clock_frequency
    }

    /// Get the current reference clock multiplier.
    pub fn get_reference_clock_multiplier(&mut self) -> Result<u8, Error> {
        let mut fr1: [u8; 3] = [0, 0, 0];
        self.interface
            .read(Register::FR1 as u8, &mut fr1)
            .map_err(|_| Error::Interface)?;

        Ok(fr1[0].get_bits(2..=6) as u8)
    }

    /// Perform a self-test of the communication interface.
    ///
    /// Note:
    /// This modifies the existing channel enables. They are restored upon exit.
    ///
    /// Returns:
    /// True if the self test succeeded. False otherwise.
    pub fn self_test(&mut self) -> Result<bool, Error> {
        let mut csr: [u8; 1] = [0];
        self.interface
            .read(Register::CSR as u8, &mut csr)
            .map_err(|_| Error::Interface)?;
        let old_csr = csr[0];

        // Enable all channels.
        csr[0].set_bits(4..8, 0xF);
        self.interface
            .write(Register::CSR as u8, &csr)
            .map_err(|_| Error::Interface)?;

        // Read back the enable.
        csr[0] = 0;
        self.interface
            .read(Register::CSR as u8, &mut csr)
            .map_err(|_| Error::Interface)?;
        if csr[0].get_bits(4..8) != 0xF {
            return Ok(false);
        }

        // Clear all channel enables.
        csr[0].set_bits(4..8, 0x0);
        self.interface
            .write(Register::CSR as u8, &csr)
            .map_err(|_| Error::Interface)?;

        // Read back the enable.
        csr[0] = 0xFF;
        self.interface
            .read(Register::CSR as u8, &mut csr)
            .map_err(|_| Error::Interface)?;
        if csr[0].get_bits(4..8) != 0 {
            return Ok(false);
        }

        // Restore the CSR.
        csr[0] = old_csr;
        self.interface
            .write(Register::CSR as u8, &csr)
            .map_err(|_| Error::Interface)?;

        Ok(true)
    }

    /// Get the current system clock frequency in Hz.
    fn system_clock_frequency(&self) -> f64 {
        self.system_clock_multiplier as f64
            * self.reference_clock_frequency as f64
    }

    /// Enable an output channel.
    pub fn enable_channel(&mut self, channel: Channel) -> Result<(), Error> {
        let mut csr: [u8; 1] = [0];
        self.interface
            .read(Register::CSR as u8, &mut csr)
            .map_err(|_| Error::Interface)?;
        csr[0].set_bit(channel as usize + 4, true);
        self.interface
            .write(Register::CSR as u8, &csr)
            .map_err(|_| Error::Interface)?;

        Ok(())
    }

    /// Disable an output channel.
    pub fn disable_channel(&mut self, channel: Channel) -> Result<(), Error> {
        let mut csr: [u8; 1] = [0];
        self.interface
            .read(Register::CSR as u8, &mut csr)
            .map_err(|_| Error::Interface)?;
        csr[0].set_bit(channel as usize + 4, false);
        self.interface
            .write(Register::CSR as u8, &csr)
            .map_err(|_| Error::Interface)?;

        Ok(())
    }

    /// Determine if an output channel is enabled.
    pub fn is_enabled(&mut self, channel: Channel) -> Result<bool, Error> {
        let mut csr: [u8; 1] = [0; 1];
        self.interface
            .read(Register::CSR as u8, &mut csr)
            .map_err(|_| Error::Interface)?;

        Ok(csr[0].get_bit(channel as usize + 4))
    }

    /// Update an output channel configuration register.
    ///
    /// Args:
    /// * `channel` - The channel to configure.
    /// * `register` - The register to update.
    /// * `data` - The contents to write to the provided register.
    fn modify_channel(
        &mut self,
        channel: Channel,
        register: Register,
        data: &[u8],
    ) -> Result<(), Error> {
        // Disable all other outputs so that we can update the configuration register of only the
        // specified channel.
        let mut csr: [u8; 1] = [0];
        self.interface
            .read(Register::CSR as u8, &mut csr)
            .map_err(|_| Error::Interface)?;

        let mut new_csr = csr;
        new_csr[0].set_bits(4..8, 0);
        new_csr[0].set_bit(4 + channel as usize, true);

        self.interface
            .write(Register::CSR as u8, &new_csr)
            .map_err(|_| Error::Interface)?;

        self.interface
            .write(register as u8, &data)
            .map_err(|_| Error::Interface)?;

        // Latch the configuration and restore the previous CSR. Note that the re-enable of the
        // channel happens immediately, so the CSR update does not need to be latched.
        self.latch_configuration()?;

        Ok(())
    }

    /// Read a configuration register of a specific channel.
    ///
    /// Args:
    /// * `channel` - The channel to read.
    /// * `register` - The register to read.
    /// * `data` - A location to store the read register contents.
    fn read_channel(
        &mut self,
        channel: Channel,
        register: Register,
        mut data: &mut [u8],
    ) -> Result<(), Error> {
        // Disable all other channels in the CSR so that we can read the configuration register of
        // only the desired channel.
        let mut csr: [u8; 1] = [0];
        self.interface
            .read(Register::CSR as u8, &mut csr)
            .map_err(|_| Error::Interface)?;

        let mut new_csr = csr;
        new_csr[0].set_bits(4..8, 0);
        new_csr[0].set_bit(4 + channel as usize, true);

        self.interface
            .write(Register::CSR as u8, &new_csr)
            .map_err(|_| Error::Interface)?;

        self.interface
            .read(register as u8, &mut data)
            .map_err(|_| Error::Interface)?;

        // Restore the previous CSR. Note that the re-enable of the channel happens immediately, so
        // the CSR update does not need to be latched.
        self.interface
            .write(Register::CSR as u8, &csr)
            .map_err(|_| Error::Interface)?;

        Ok(())
    }

    /// Configure the phase of a specified channel.
    ///
    /// Arguments:
    /// * `channel` - The channel to configure the frequency of.
    /// * `phase_turns` - The desired phase offset in turns.
    ///
    /// Returns:
    /// The actual programmed phase offset of the channel in turns.
    pub fn set_phase(
        &mut self,
        channel: Channel,
        phase_turns: f32,
    ) -> Result<f32, Error> {
        let phase_offset: u16 =
            (phase_turns * (1 << 14) as f32) as u16 & 0x3FFFu16;

        self.modify_channel(
            channel,
            Register::CPOW0,
            &phase_offset.to_be_bytes(),
        )?;

        Ok((phase_offset as f32) / ((1 << 14) as f32))
    }

    /// Get the current phase of a specified channel.
    ///
    /// Args:
    /// * `channel` - The channel to get the phase of.
    ///
    /// Returns:
    /// The phase of the channel in turns.
    pub fn get_phase(&mut self, channel: Channel) -> Result<f32, Error> {
        let mut phase_offset: [u8; 2] = [0; 2];
        self.read_channel(channel, Register::CPOW0, &mut phase_offset)?;

        let phase_offset = u16::from_be_bytes(phase_offset) & 0x3FFFu16;

        Ok((phase_offset as f32) / ((1 << 14) as f32))
    }

    /// Configure the amplitude of a specified channel.
    ///
    /// Arguments:
    /// * `channel` - The channel to configure the frequency of.
    /// * `amplitude` - A normalized amplitude setting [0, 1].
    ///
    /// Returns:
    /// The actual normalized amplitude of the channel relative to full-scale range.
    pub fn set_amplitude(
        &mut self,
        channel: Channel,
        amplitude: f32,
    ) -> Result<f32, Error> {
        if amplitude < 0.0 || amplitude > 1.0 {
            return Err(Error::Bounds);
        }

        let amplitude_control: u16 = (amplitude * (1 << 10) as f32) as u16;

        let mut acr: [u8; 3] = [0; 3];

        // Enable the amplitude multiplier for the channel if required. The amplitude control has
        // full-scale at 0x3FF (amplitude of 1), so the multiplier should be disabled whenever
        // full-scale is used.
        if amplitude_control < (1 << 10) {
            let masked_control = amplitude_control & 0x3FF;
            acr[1] = masked_control.to_be_bytes()[0];
            acr[2] = masked_control.to_be_bytes()[1];

            // Enable the amplitude multiplier
            acr[1].set_bit(4, true);
        }

        self.modify_channel(channel, Register::ACR, &acr)?;

        Ok(amplitude_control as f32 / (1 << 10) as f32)
    }

    /// Get the configured amplitude of a channel.
    ///
    /// Args:
    /// * `channel` - The channel to get the amplitude of.
    ///
    /// Returns:
    /// The normalized amplitude of the channel.
    pub fn get_amplitude(&mut self, channel: Channel) -> Result<f32, Error> {
        let mut acr: [u8; 3] = [0; 3];
        self.read_channel(channel, Register::ACR, &mut acr)?;

        if acr[1].get_bit(4) {
            let amplitude_control: u16 =
                (((acr[1] as u16) << 8) | (acr[2] as u16)) & 0x3FF;
            Ok(amplitude_control as f32 / (1 << 10) as f32)
        } else {
            Ok(1.0)
        }
    }

    /// Configure the frequency of a specified channel.
    ///
    /// Arguments:
    /// * `channel` - The channel to configure the frequency of.
    /// * `frequency` - The desired output frequency in Hz.
    ///
    /// Returns:
    /// The actual programmed frequency of the channel.
    pub fn set_frequency(
        &mut self,
        channel: Channel,
        frequency: f64,
    ) -> Result<f64, Error> {
        if frequency < 0.0 || frequency > self.system_clock_frequency() {
            return Err(Error::Bounds);
        }

        // The function for channel frequency is `f_out = FTW * f_s / 2^32`, where FTW is the
        // frequency tuning word and f_s is the system clock rate.
        let tuning_word: u32 =
            ((frequency as f64 / self.system_clock_frequency())
                * 1u64.wrapping_shl(32) as f64) as u32;

        self.modify_channel(
            channel,
            Register::CFTW0,
            &tuning_word.to_be_bytes(),
        )?;
        Ok((tuning_word as f64 / 1u64.wrapping_shl(32) as f64)
            * self.system_clock_frequency())
    }

    /// Get the frequency of a channel.
    ///
    /// Arguments:
    /// * `channel` - The channel to get the frequency of.
    ///
    /// Returns:
    /// The frequency of the channel in Hz.
    pub fn get_frequency(&mut self, channel: Channel) -> Result<f64, Error> {
        // Read the frequency tuning word for the channel.
        let mut tuning_word: [u8; 4] = [0; 4];
        self.read_channel(channel, Register::CFTW0, &mut tuning_word)?;
        let tuning_word = u32::from_be_bytes(tuning_word);

        // Convert the tuning word into a frequency.
        Ok(tuning_word as f64 * self.system_clock_frequency()
            / (1u64 << 32) as f64)
    }

    pub fn write_profile(&mut self, channel: Channel, freq: f64, turns: f32, amplitude: f32) -> Result<(), Error> {
        // The function for channel frequency is `f_out = FTW * f_s / 2^32`, where FTW is the
        // frequency tuning word and f_s is the system clock rate.
        let tuning_word: u32 =
            ((freq as f64 / self.system_clock_frequency())
                * 1u64.wrapping_shl(32) as f64) as u32;

        let phase_offset: u16 = (turns * (1 << 14) as f32) as u16 & 0x3FFFu16;

        let amplitude_control: u16 = (amplitude * (1 << 10) as f32) as u16;
        let mut acr: [u8; 3] = [0; 3];

        // Enable the amplitude multiplier for the channel if required. The amplitude control has
        // full-scale at 0x3FF (amplitude of 1), so the multiplier should be disabled whenever
        // full-scale is used.
        if amplitude_control < (1 << 10) {
            let masked_control = amplitude_control & 0x3FF;
            acr[1] = masked_control.to_be_bytes()[0];
            acr[2] = masked_control.to_be_bytes()[1];

            // Enable the amplitude multiplier
            acr[1].set_bit(4, true);
        }

        self.modify_channel_closure(channel, |interface| {
            let mut data: [u8; 11] = [0; 11];
            data[0..2].copy_from_slice(&phase_offset.to_be_bytes());
            data[2] = Register::CFTW0 as u8;
            data[3..7].copy_from_slice(&tuning_word.to_be_bytes());
            data[7] = Register::ACR as u8;
            data[8..11].copy_from_slice(&acr);
            interface.write(Register::CPOW0 as u8, &data).map_err(|_| Error::Interface)
        })?;

        Ok(())
    }

    fn modify_channel_closure<F>(&mut self, channel: Channel, f: F) -> Result<(), Error>
    where
        F: FnOnce(&mut INTERFACE) -> Result<(), Error>,
    {
        // Disable all other outputs so that we can update the configuration register of only the
        // specified channel.
        let mut csr: [u8; 1] = [0];
        self.interface
            .read(Register::CSR as u8, &mut csr)
            .map_err(|_| Error::Interface)?;

        let mut new_csr = csr;
        new_csr[0].set_bits(4..8, 0);
        new_csr[0].set_bit(4 + channel as usize, true);

        self.interface
            .write(Register::CSR as u8, &new_csr)
            .map_err(|_| Error::Interface)?;

        let result = f(&mut self.interface);

        // Latch the configuration and restore the previous CSR. Note that the re-enable of the
        // channel happens immediately, so the CSR update does not need to be latched.
        self.latch_configuration()?;

        result
    }
}