//! Baseline correction for neuroimaging signals.

use crate::{SignalError, SignalResult};

/// Baseline correction method
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum BaselineMethod {
    /// Subtract mean of baseline period
    #[default]
    Mean,
    /// Subtract median of baseline period
    Median,
    /// Z-score normalization (subtract mean, divide by std)
    Zscore,
    /// Percent change relative to baseline mean
    Percent,
    /// Decibel scale: 10 * log10(signal / baseline_mean)
    Decibel,
    /// Log ratio: log(signal / baseline_mean)
    LogRatio,
}

/// Apply baseline correction to a signal.
///
/// # Arguments
/// * `signal` - Input signal samples
/// * `baseline_start` - Start index of baseline period
/// * `baseline_end` - End index of baseline period (exclusive)
/// * `method` - Baseline correction method
///
/// # Returns
/// Baseline-corrected signal
pub fn baseline_correct(
    signal: &[f64],
    baseline_start: usize,
    baseline_end: usize,
    method: BaselineMethod,
) -> SignalResult<Vec<f64>> {
    if signal.is_empty() {
        return Ok(Vec::new());
    }

    if baseline_start >= baseline_end || baseline_end > signal.len() {
        return Err(SignalError::InvalidParameters(format!(
            "Invalid baseline range [{baseline_start}, {baseline_end}) for signal length {}",
            signal.len()
        )));
    }

    let baseline = &signal[baseline_start..baseline_end];

    match method {
        BaselineMethod::Mean => {
            let mean = baseline.iter().sum::<f64>() / baseline.len() as f64;
            Ok(signal.iter().map(|&x| x - mean).collect())
        }
        BaselineMethod::Median => {
            let median = compute_median(baseline);
            Ok(signal.iter().map(|&x| x - median).collect())
        }
        BaselineMethod::Zscore => {
            let mean = baseline.iter().sum::<f64>() / baseline.len() as f64;
            let variance = baseline.iter().map(|&x| (x - mean).powi(2)).sum::<f64>()
                / (baseline.len() - 1) as f64;
            let std = variance.sqrt();

            if std < 1e-10 {
                return Err(SignalError::Numerical(
                    "Baseline standard deviation is zero".to_string(),
                ));
            }

            Ok(signal.iter().map(|&x| (x - mean) / std).collect())
        }
        BaselineMethod::Percent => {
            let mean = baseline.iter().sum::<f64>() / baseline.len() as f64;

            if mean.abs() < 1e-10 {
                return Err(SignalError::Numerical(
                    "Baseline mean is zero, cannot compute percent change".to_string(),
                ));
            }

            Ok(signal.iter().map(|&x| 100.0 * (x - mean) / mean).collect())
        }
        BaselineMethod::Decibel => {
            let mean = baseline.iter().sum::<f64>() / baseline.len() as f64;

            if mean <= 0.0 {
                return Err(SignalError::Numerical(
                    "Baseline mean must be positive for decibel conversion".to_string(),
                ));
            }

            Ok(signal.iter().map(|&x| 10.0 * (x / mean).log10()).collect())
        }
        BaselineMethod::LogRatio => {
            let mean = baseline.iter().sum::<f64>() / baseline.len() as f64;

            if mean <= 0.0 {
                return Err(SignalError::Numerical(
                    "Baseline mean must be positive for log ratio".to_string(),
                ));
            }

            Ok(signal.iter().map(|&x| (x / mean).ln()).collect())
        }
    }
}

/// Apply baseline correction to multiple epochs.
///
/// # Arguments
/// * `data` - Epoch data [n_epochs x n_channels x n_times] flattened
/// * `n_epochs` - Number of epochs
/// * `n_channels` - Number of channels
/// * `n_times` - Number of time points per epoch
/// * `baseline_start` - Start time index within each epoch
/// * `baseline_end` - End time index within each epoch
/// * `method` - Baseline correction method
pub fn baseline_correct_epochs(
    data: &mut [f64],
    n_epochs: usize,
    n_channels: usize,
    n_times: usize,
    baseline_start: usize,
    baseline_end: usize,
    method: BaselineMethod,
) -> SignalResult<()> {
    if baseline_start >= baseline_end || baseline_end > n_times {
        return Err(SignalError::InvalidParameters(format!(
            "Invalid baseline range [{baseline_start}, {baseline_end}) for n_times {n_times}"
        )));
    }

    for epoch in 0..n_epochs {
        let epoch_offset = epoch * n_channels * n_times;

        for ch in 0..n_channels {
            let ch_offset = epoch_offset + ch * n_times;
            let ch_data = &mut data[ch_offset..ch_offset + n_times];

            // Compute baseline statistics
            let baseline = &ch_data[baseline_start..baseline_end];
            let baseline_len = baseline.len() as f64;

            match method {
                BaselineMethod::Mean => {
                    let mean: f64 = baseline.iter().sum::<f64>() / baseline_len;
                    for x in ch_data.iter_mut() {
                        *x -= mean;
                    }
                }
                BaselineMethod::Median => {
                    let median = compute_median(baseline);
                    for x in ch_data.iter_mut() {
                        *x -= median;
                    }
                }
                BaselineMethod::Zscore => {
                    let mean: f64 = baseline.iter().sum::<f64>() / baseline_len;
                    let variance: f64 = baseline.iter().map(|&x| (x - mean).powi(2)).sum::<f64>()
                        / (baseline_len - 1.0);
                    let std = variance.sqrt();

                    if std > 1e-10 {
                        for x in ch_data.iter_mut() {
                            *x = (*x - mean) / std;
                        }
                    }
                }
                BaselineMethod::Percent => {
                    let mean: f64 = baseline.iter().sum::<f64>() / baseline_len;
                    if mean.abs() > 1e-10 {
                        for x in ch_data.iter_mut() {
                            *x = 100.0 * (*x - mean) / mean;
                        }
                    }
                }
                BaselineMethod::Decibel => {
                    let mean: f64 = baseline.iter().sum::<f64>() / baseline_len;
                    if mean > 0.0 {
                        for x in ch_data.iter_mut() {
                            *x = 10.0 * (*x / mean).log10();
                        }
                    }
                }
                BaselineMethod::LogRatio => {
                    let mean: f64 = baseline.iter().sum::<f64>() / baseline_len;
                    if mean > 0.0 {
                        for x in ch_data.iter_mut() {
                            *x = (*x / mean).ln();
                        }
                    }
                }
            }
        }
    }

    Ok(())
}

/// Compute the median of a slice
fn compute_median(data: &[f64]) -> f64 {
    if data.is_empty() {
        return 0.0;
    }

    let mut sorted: Vec<f64> = data.to_vec();
    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

    let mid = sorted.len() / 2;
    if sorted.len().is_multiple_of(2) {
        f64::midpoint(sorted[mid - 1], sorted[mid])
    } else {
        sorted[mid]
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_baseline_mean() {
        let signal = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0];
        // Baseline is first 5 samples: mean = 3.0
        let corrected = baseline_correct(&signal, 0, 5, BaselineMethod::Mean).unwrap();

        assert!((corrected[0] - (-2.0)).abs() < 1e-10);
        assert!((corrected[4] - 2.0).abs() < 1e-10);
        assert!((corrected[9] - 7.0).abs() < 1e-10);
    }

    #[test]
    fn test_baseline_zscore() {
        let signal = vec![0.0, 1.0, 2.0, 3.0, 4.0]; // mean=1, std=~1.58
        let corrected = baseline_correct(&signal, 0, 3, BaselineMethod::Zscore).unwrap();

        // After z-scoring, baseline should have mean 0, std ~1
        let baseline_mean: f64 = corrected[..3].iter().sum::<f64>() / 3.0;
        assert!(baseline_mean.abs() < 1e-10);
    }

    #[test]
    fn test_median() {
        assert!((compute_median(&[1.0, 2.0, 3.0]) - 2.0).abs() < 1e-10);
        assert!((compute_median(&[1.0, 2.0, 3.0, 4.0]) - 2.5).abs() < 1e-10);
    }
}