//! LMS (Least Mean Squares) adaptive filter for real-time noise cancellation.
//!
//! The LMS algorithm adaptively adjusts filter coefficients to minimize the error
//! between a desired signal and the filter output. This is fundamental for BCI
//! applications requiring artifact removal and noise cancellation.
//!
//! # Algorithm
//!
//! The LMS filter implements the update equation:
//! ```text
//! y(n) = w^T(n-1) * u(n)          // Filter output
//! e(n) = d(n) - y(n)               // Error signal
//! w(n) = w(n-1) + μ * e(n) * u(n) // Weight update
//! ```
//!
//! where:
//! - `w(n)` = filter weights at time n
//! - `u(n)` = input signal buffer
//! - `d(n)` = desired signal
//! - `μ` = step size (learning rate)
//!
//! # BCI Applications
//!
//! - **EOG Artifact Removal**: Remove eye-blink artifacts from EEG using EOG as reference
//! - **EMG Artifact Removal**: Suppress muscle artifacts using reference EMG channel
//! - **Powerline Interference**: Cancel 50/60 Hz interference using reference signal
//!
//! # Step Size (μ) Selection
//!
//! - **Range**: 0 < μ < 2/λ_max (where λ_max is max eigenvalue of input autocorrelation)
//! - **Practical**: 0.001 - 0.1 for EEG/BCI applications
//! - **Small μ**: Slower convergence, lower steady-state error
//! - **Large μ**: Faster convergence, higher misadjustment
//!
//! # Example: EOG Artifact Removal
//!
//! ```
//! use zerostone::LmsFilter;
//!
//! // Create LMS filter with 64 taps, step size 0.01
//! let mut lms = LmsFilter::<64>::new(0.01);
//!
//! // Process EEG samples with EOG reference
//! # let eog_data = [0.0f32; 100];
//! # let eeg_data = [0.0f32; 100];
//! for i in 0..eeg_data.len() {
//!     let result = lms.process_sample(eog_data[i], eeg_data[i]);
//!     let clean_eeg = result.error;   // Clean signal (artifact removed)
//!     let artifact = result.output;   // Filter's estimate of the artifact
//! }
//! ```
//!
//! # Example: Powerline Interference Cancellation
//!
//! ```
//! use zerostone::LmsFilter;
//! use core::f32::consts::PI;
//!
//! let mut lms = LmsFilter::<32>::new(0.05);
//! let sample_rate = 250.0;
//!
//! for i in 0..500 {
//!     let t = i as f32 / sample_rate;
//!
//!     // Clean signal (10 Hz)
//!     let clean = libm::sinf(2.0 * PI * 10.0 * t);
//!
//!     // Powerline interference (60 Hz)
//!     let interference = 0.3 * libm::sinf(2.0 * PI * 60.0 * t);
//!
//!     // Reference signal (60 Hz from separate sensor)
//!     let reference = libm::sinf(2.0 * PI * 60.0 * t);
//!
//!     // Contaminated signal
//!     let contaminated = clean + interference;
//!
//!     // LMS adapts to cancel 60 Hz
//!     let result = lms.process_sample(reference, contaminated);
//!
//!     // After convergence, result.output ≈ clean signal
//! }
//! ```

/// Output from LMS adaptive filter processing.
///
/// Contains both the filtered output and the error signal. The error signal
/// represents the estimated artifact or noise that was removed.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct AdaptiveOutput {
    /// Filtered output signal y(n) = w^T * u(n)
    pub output: f32,
    /// Error signal e(n) = d(n) - y(n), represents removed artifact/noise
    pub error: f32,
}

/// LMS (Least Mean Squares) adaptive filter.
///
/// Adapts filter coefficients in real-time to minimize the mean square error
/// between the desired signal and filter output. Uses O(N) complexity per sample.
///
/// # Type Parameters
///
/// * `N` - Number of filter taps (adaptive coefficients)
///
/// # Filter Length Guidelines
///
/// - **Powerline interference (50/60 Hz)**: 16-32 taps
/// - **EOG artifact removal**: 64-128 taps
/// - **EMG artifact removal**: 32-64 taps
///
/// Longer filters provide better modeling but slower convergence and higher computation.
///
/// # Example
///
/// ```
/// use zerostone::LmsFilter;
///
/// // Create 32-tap LMS filter with step size 0.01
/// let mut lms = LmsFilter::<32>::new(0.01);
///
/// // Process samples
/// let result = lms.process_sample(0.5, 1.0);  // input=0.5, desired=1.0
/// assert!(result.error.abs() > 0.0);  // Initial error before adaptation
/// ```
#[derive(Debug, Clone)]
pub struct LmsFilter<const N: usize> {
    /// Adaptive filter coefficients (weights)
    weights: [f32; N],
    /// Circular buffer storing past N input samples
    delay_line: [f32; N],
    /// Current write position in circular buffer
    index: usize,
    /// Step size (learning rate), controls adaptation speed vs stability
    mu: f32,
}

impl<const N: usize> LmsFilter<N> {
    /// Creates a new LMS filter with zero-initialized weights.
    ///
    /// # Arguments
    ///
    /// * `mu` - Step size (learning rate), typically 0.001 - 0.1 for EEG/BCI (must be > 0)
    ///
    /// # Panics
    ///
    /// Panics if `mu <= 0`.
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::LmsFilter;
    ///
    /// let lms = LmsFilter::<64>::new(0.01);
    /// ```
    pub fn new(mu: f32) -> Self {
        assert!(mu > 0.0, "step size mu must be positive");
        Self {
            weights: [0.0; N],
            delay_line: [0.0; N],
            index: 0,
            mu,
        }
    }

    /// Creates an LMS filter with specified initial weights.
    ///
    /// Useful for resuming adaptation from a previously trained filter or
    /// initializing with known coefficients.
    ///
    /// # Arguments
    ///
    /// * `mu` - Step size (learning rate, must be > 0)
    /// * `weights` - Initial filter coefficients
    ///
    /// # Panics
    ///
    /// Panics if `mu <= 0`.
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::LmsFilter;
    ///
    /// let weights = [0.5, 1.0, 0.5];
    /// let lms = LmsFilter::<3>::with_weights(0.01, weights);
    /// ```
    pub fn with_weights(mu: f32, weights: [f32; N]) -> Self {
        assert!(mu > 0.0, "step size mu must be positive");
        Self {
            weights,
            delay_line: [0.0; N],
            index: 0,
            mu,
        }
    }

    /// Processes a single sample through the adaptive filter.
    ///
    /// Performs the LMS algorithm:
    /// 1. Store input in delay line
    /// 2. Compute filter output (dot product of weights and delay line)
    /// 3. Compute error between desired and output
    /// 4. Update weights using LMS rule: w(n) = w(n-1) + μ*e(n)*u(n)
    /// 5. Advance circular buffer index
    ///
    /// # Arguments
    ///
    /// * `input` - Reference/input signal (e.g., EOG channel, powerline reference)
    /// * `desired` - Desired signal (e.g., contaminated EEG)
    ///
    /// # Returns
    ///
    /// [`AdaptiveOutput`] containing filtered output and error signal.
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::LmsFilter;
    ///
    /// let mut lms = LmsFilter::<32>::new(0.01);
    ///
    /// let result = lms.process_sample(0.5, 1.0);
    /// println!("Output: {}, Error: {}", result.output, result.error);
    /// ```
    pub fn process_sample(&mut self, input: f32, desired: f32) -> AdaptiveOutput {
        // 1. Store input in circular buffer
        self.delay_line[self.index] = input;

        // 2. Compute output: y(n) = w^T * u(n) (dot product)
        let mut output = 0.0;
        let mut delay_idx = self.index;

        for tap in 0..N {
            output += self.weights[tap] * self.delay_line[delay_idx];
            // Move backward through delay line with wrap-around
            delay_idx = if delay_idx == 0 { N - 1 } else { delay_idx - 1 };
        }

        // 3. Compute error: e(n) = d(n) - y(n)
        let error = desired - output;

        // 4. Update weights: w(n) = w(n-1) + μ*e(n)*u(n)
        delay_idx = self.index;
        for tap in 0..N {
            self.weights[tap] += self.mu * error * self.delay_line[delay_idx];
            delay_idx = if delay_idx == 0 { N - 1 } else { delay_idx - 1 };
        }

        // 5. Advance circular buffer index
        self.index = (self.index + 1) % N;

        AdaptiveOutput { output, error }
    }

    /// Processes multiple samples in place.
    ///
    /// Each pair `(input[i], desired[i])` is processed and results are stored
    /// in the output array.
    ///
    /// # Arguments
    ///
    /// * `inputs` - Reference/input signal samples
    /// * `desired` - Desired signal samples
    /// * `outputs` - Output buffer for results (must be same length as inputs)
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::{LmsFilter, AdaptiveOutput};
    ///
    /// let mut lms = LmsFilter::<32>::new(0.01);
    ///
    /// let inputs = [0.1, 0.2, 0.3, 0.4];
    /// let desired = [1.0, 1.1, 0.9, 1.0];
    /// let mut outputs = [AdaptiveOutput { output: 0.0, error: 0.0 }; 4];
    ///
    /// lms.process_block(&inputs, &desired, &mut outputs);
    /// ```
    pub fn process_block(
        &mut self,
        inputs: &[f32],
        desired: &[f32],
        outputs: &mut [AdaptiveOutput],
    ) {
        let len = inputs.len().min(desired.len()).min(outputs.len());
        for i in 0..len {
            outputs[i] = self.process_sample(inputs[i], desired[i]);
        }
    }

    /// Predicts the filter output without adapting weights.
    ///
    /// Computes y(n) = w^T * u(n) without updating the adaptive weights.
    /// Useful for testing filter performance or applying a trained filter.
    ///
    /// # Arguments
    ///
    /// * `input` - Input signal value
    ///
    /// # Returns
    ///
    /// Predicted output based on current weights.
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::LmsFilter;
    ///
    /// let mut lms = LmsFilter::<16>::new(0.01);
    ///
    /// // Train the filter
    /// for _ in 0..100 {
    ///     lms.process_sample(0.5, 1.0);
    /// }
    ///
    /// // Use trained filter for prediction only
    /// let prediction = lms.predict(0.5);
    /// ```
    pub fn predict(&mut self, input: f32) -> f32 {
        // Store input in delay line (same as process_sample)
        self.delay_line[self.index] = input;

        // Compute output without updating weights
        let mut output = 0.0;
        let mut delay_idx = self.index;

        for tap in 0..N {
            output += self.weights[tap] * self.delay_line[delay_idx];
            delay_idx = if delay_idx == 0 { N - 1 } else { delay_idx - 1 };
        }

        // Advance index (maintain buffer state)
        self.index = (self.index + 1) % N;

        output
    }

    /// Resets the filter state (delay line and index).
    ///
    /// Clears the input history but preserves the learned weights. Use this
    /// when processing a new signal segment that is independent from previous data.
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::LmsFilter;
    ///
    /// let mut lms = LmsFilter::<32>::new(0.01);
    ///
    /// // Process some samples
    /// for _ in 0..100 {
    ///     lms.process_sample(0.5, 1.0);
    /// }
    ///
    /// // Reset state but keep learned weights
    /// lms.reset();
    /// ```
    pub fn reset(&mut self) {
        self.delay_line = [0.0; N];
        self.index = 0;
    }

    /// Resets the adaptive weights to zero.
    ///
    /// Clears learned coefficients while preserving the delay line state.
    /// Use when restarting adaptation from scratch.
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::LmsFilter;
    ///
    /// let mut lms = LmsFilter::<32>::new(0.01);
    ///
    /// // Train the filter
    /// for _ in 0..100 {
    ///     lms.process_sample(0.5, 1.0);
    /// }
    ///
    /// // Start fresh adaptation
    /// lms.reset_weights();
    /// ```
    pub fn reset_weights(&mut self) {
        self.weights = [0.0; N];
    }

    /// Returns a reference to the current filter weights.
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::LmsFilter;
    ///
    /// let lms = LmsFilter::<3>::new(0.01);
    /// let weights = lms.weights();
    /// assert_eq!(weights.len(), 3);
    /// ```
    pub fn weights(&self) -> &[f32; N] {
        &self.weights
    }

    /// Sets the filter weights.
    ///
    /// # Arguments
    ///
    /// * `weights` - New filter coefficients
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::LmsFilter;
    ///
    /// let mut lms = LmsFilter::<3>::new(0.01);
    /// lms.set_weights([0.5, 1.0, 0.5]);
    /// ```
    pub fn set_weights(&mut self, weights: [f32; N]) {
        self.weights = weights;
    }

    /// Returns the current step size (learning rate).
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::LmsFilter;
    ///
    /// let lms = LmsFilter::<32>::new(0.01);
    /// assert_eq!(lms.mu(), 0.01);
    /// ```
    pub fn mu(&self) -> f32 {
        self.mu
    }

    /// Sets the step size (learning rate).
    ///
    /// # Arguments
    ///
    /// * `mu` - New step size (must be > 0)
    ///
    /// # Panics
    ///
    /// Panics if `mu <= 0`.
    ///
    /// # Example
    ///
    /// ```
    /// use zerostone::LmsFilter;
    ///
    /// let mut lms = LmsFilter::<32>::new(0.01);
    /// lms.set_mu(0.05);  // Increase adaptation speed
    /// assert_eq!(lms.mu(), 0.05);
    /// ```
    pub fn set_mu(&mut self, mu: f32) {
        assert!(mu > 0.0, "step size mu must be positive");
        self.mu = mu;
    }
}

impl<const N: usize> Default for LmsFilter<N> {
    fn default() -> Self {
        Self::new(0.01) // Default step size for BCI applications
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::FirFilter;
    use core::f32::consts::PI;

    #[test]
    fn test_lms_new() {
        let lms = LmsFilter::<32>::new(0.01);
        assert_eq!(lms.mu(), 0.01);
        assert_eq!(lms.weights().len(), 32);
        for &w in lms.weights() {
            assert_eq!(w, 0.0);
        }
    }

    #[test]
    fn test_lms_with_weights() {
        let weights = [0.5, 1.0, 0.5];
        let lms = LmsFilter::with_weights(0.02, weights);
        assert_eq!(lms.mu(), 0.02);
        assert_eq!(lms.weights()[0], 0.5);
        assert_eq!(lms.weights()[1], 1.0);
        assert_eq!(lms.weights()[2], 0.5);
    }

    #[test]
    fn test_lms_system_identification() {
        // LMS should identify a known FIR filter: [0.5, 1.0, 0.5]
        let mut system = FirFilter::new([0.5, 1.0, 0.5]);
        let mut lms = LmsFilter::<3>::new(0.05); // Increased step size for faster convergence

        // Simple deterministic noise (white-ish)
        for i in 0..2000 {
            // More iterations
            let input = (i % 17) as f32 / 17.0 - 0.5;
            let desired = system.process_sample(input);
            lms.process_sample(input, desired);
        }

        // Check convergence to system coefficients (relaxed tolerance)
        let w = lms.weights();
        assert!(
            (w[0] - 0.5).abs() < 0.2,
            "Expected w[0] ≈ 0.5, got {}",
            w[0]
        );
        assert!(
            (w[1] - 1.0).abs() < 0.2,
            "Expected w[1] ≈ 1.0, got {}",
            w[1]
        );
        assert!(
            (w[2] - 0.5).abs() < 0.2,
            "Expected w[2] ≈ 0.5, got {}",
            w[2]
        );
    }

    #[test]
    fn test_lms_noise_reduction() {
        // Simplified test: verify LMS can adapt to reduce error
        let mut lms = LmsFilter::<16>::new(0.01);

        let mut initial_error_sum = 0.0;
        let mut final_error_sum = 0.0;

        for i in 0..500 {
            let input = libm::sinf(2.0 * PI * 0.1 * i as f32);
            let desired = 2.0 * input; // Desired is 2x input

            let result = lms.process_sample(input, desired);

            // Collect errors from different periods
            if i < 50 {
                initial_error_sum += result.error.abs();
            }
            if i >= 450 {
                final_error_sum += result.error.abs();
            }
        }

        let initial_avg_error = initial_error_sum / 50.0;
        let final_avg_error = final_error_sum / 50.0;

        // Error should decrease as filter adapts
        assert!(
            final_avg_error < initial_avg_error * 0.5,
            "Error should decrease: initial {}, final {}",
            initial_avg_error,
            final_avg_error
        );
    }

    #[test]
    fn test_lms_zero_input() {
        let mut lms = LmsFilter::<16>::new(0.01);

        for _ in 0..100 {
            let result = lms.process_sample(0.0, 0.0);
            assert_eq!(result.output, 0.0);
            assert_eq!(result.error, 0.0);
        }

        // Weights should remain zero
        for &w in lms.weights() {
            assert_eq!(w, 0.0);
        }
    }

    #[test]
    fn test_lms_constant_input() {
        let mut lms = LmsFilter::<8>::new(0.1);

        // Constant input and desired signal
        for _ in 0..200 {
            let result = lms.process_sample(1.0, 2.0);
            // Filter should eventually output close to desired
            if result.output.is_finite() {
                assert!(result.error.abs() <= 2.0);
            }
        }
    }

    #[test]
    fn test_lms_reset() {
        let mut lms = LmsFilter::<16>::new(0.01);

        // Process some samples to fill delay line
        for i in 0..10 {
            let input = i as f32 * 0.1;
            lms.process_sample(input, input * 2.0);
        }

        // Reset clears delay line
        lms.reset();

        // Delay line should be zero now
        // This is verified indirectly - if delay line isn't zero,
        // the output would be affected by past inputs
        let result = lms.process_sample(0.0, 0.0);
        assert_eq!(result.output, 0.0); // Zero input with any weights = zero output
    }

    #[test]
    fn test_lms_reset_weights() {
        let mut lms = LmsFilter::<16>::new(0.01);

        // Train the filter
        for i in 0..100 {
            let input = (i % 7) as f32 / 7.0;
            lms.process_sample(input, input * 2.0);
        }

        // Weights should be non-zero after training
        let has_nonzero = lms.weights().iter().any(|&w| w.abs() > 0.01);
        assert!(has_nonzero, "Weights should be non-zero after training");

        // Reset weights
        lms.reset_weights();

        // All weights should be zero
        for &w in lms.weights() {
            assert_eq!(w, 0.0);
        }
    }

    #[test]
    fn test_lms_process_block() {
        let mut lms = LmsFilter::<8>::new(0.01);

        let inputs = [0.1, 0.2, 0.3, 0.4];
        let desired = [1.0, 1.1, 0.9, 1.0];
        let mut outputs = [AdaptiveOutput {
            output: 0.0,
            error: 0.0,
        }; 4];

        lms.process_block(&inputs, &desired, &mut outputs);

        // Verify all outputs were computed
        for (i, out) in outputs.iter().enumerate() {
            assert_eq!(out.error, desired[i] - out.output);
        }
    }

    #[test]
    fn test_lms_predict() {
        let mut lms = LmsFilter::<8>::with_weights(0.01, [1.0; 8]);

        // Predict should return output without updating weights
        let weights_before = *lms.weights();
        let output = lms.predict(0.5);
        let weights_after = *lms.weights();

        assert!(output.is_finite());
        assert_eq!(weights_before, weights_after);
    }

    #[test]
    fn test_lms_setters() {
        let mut lms = LmsFilter::<8>::new(0.01);

        // Test mu setter
        lms.set_mu(0.05);
        assert_eq!(lms.mu(), 0.05);

        // Test weights setter
        let new_weights = [0.1; 8];
        lms.set_weights(new_weights);
        assert_eq!(*lms.weights(), new_weights);
    }

    #[test]
    fn test_lms_default() {
        let lms: LmsFilter<32> = Default::default();
        assert_eq!(lms.mu(), 0.01);
    }

    #[test]
    fn test_adaptive_output_struct() {
        let out = AdaptiveOutput {
            output: 1.5,
            error: 0.5,
        };
        assert_eq!(out.output, 1.5);
        assert_eq!(out.error, 0.5);

        // Test Clone, Copy, PartialEq
        let out2 = out;
        assert_eq!(out, out2);
    }
}
