//! Source estimate data structure.
//!
//! Represents brain source activity estimated from MEG/EEG data.

use crate::InverseResult;

/// Source time course estimate
#[derive(Debug, Clone)]
pub struct SourceEstimate {
    /// Source data [n_source_columns x n_times]
    data: Vec<Vec<f64>>,
    /// Time points in seconds
    times: Vec<f64>,
    /// Source indices (vertex indices)
    vertices: Vec<usize>,
    /// Whether sources have free orientation
    free_orientation: bool,
}

impl SourceEstimate {
    /// Create a new source estimate
    pub fn new(
        data: Vec<Vec<f64>>,
        times: Vec<f64>,
        vertices: Vec<usize>,
        free_orientation: bool,
    ) -> Self {
        Self {
            data,
            times,
            vertices,
            free_orientation,
        }
    }

    /// Get number of sources
    pub fn n_sources(&self) -> usize {
        if self.free_orientation {
            self.data.len() / 3
        } else {
            self.data.len()
        }
    }

    /// Get number of time points
    pub fn n_times(&self) -> usize {
        self.times.len()
    }

    /// Get the source data
    pub fn data(&self) -> &Vec<Vec<f64>> {
        &self.data
    }

    /// Get the time points
    pub fn times(&self) -> &[f64] {
        &self.times
    }

    /// Get vertex indices
    pub fn vertices(&self) -> &[usize] {
        &self.vertices
    }

    /// Check if sources have free orientation
    pub fn is_free_orientation(&self) -> bool {
        self.free_orientation
    }

    /// Get source activity at a specific time
    pub fn at_time(&self, time_idx: usize) -> Option<Vec<f64>> {
        if time_idx >= self.n_times() {
            return None;
        }

        Some(self.data.iter().map(|src| src[time_idx]).collect())
    }

    /// Get time course for a specific source
    pub fn get_source(&self, source_idx: usize) -> Option<&[f64]> {
        if self.free_orientation {
            // Return combined magnitude
            None // For free orientation, use get_source_vector
        } else {
            self.data.get(source_idx).map(std::vec::Vec::as_slice)
        }
    }

    /// Get vector time course for a free-orientation source
    pub fn get_source_vector(&self, source_idx: usize) -> Option<[&[f64]; 3]> {
        if !self.free_orientation {
            return None;
        }

        let idx = source_idx * 3;
        if idx + 2 >= self.data.len() {
            return None;
        }

        Some([&self.data[idx], &self.data[idx + 1], &self.data[idx + 2]])
    }

    /// Compute the magnitude time course for each source
    ///
    /// For fixed orientation: absolute value
    /// For free orientation: sqrt(x^2 + y^2 + z^2)
    pub fn magnitude(&self) -> Vec<Vec<f64>> {
        let n_sources = self.n_sources();
        let n_times = self.n_times();

        if self.free_orientation {
            (0..n_sources)
                .map(|src| {
                    let idx = src * 3;
                    (0..n_times)
                        .map(|t| {
                            let x = self.data[idx][t];
                            let y = self.data[idx + 1][t];
                            let z = self.data[idx + 2][t];
                            (x * x + y * y + z * z).sqrt()
                        })
                        .collect()
                })
                .collect()
        } else {
            self.data
                .iter()
                .map(|src| src.iter().map(|&v| v.abs()).collect())
                .collect()
        }
    }

    /// Get the mean activity across time
    pub fn mean(&self) -> Vec<f64> {
        self.data
            .iter()
            .map(|src| src.iter().sum::<f64>() / src.len() as f64)
            .collect()
    }

    /// Get the peak activity for each source
    pub fn peak(&self) -> Vec<(f64, usize)> {
        self.data
            .iter()
            .map(|src| {
                src.iter()
                    .enumerate()
                    .map(|(i, &v)| (v.abs(), i))
                    .max_by(|a, b| a.0.partial_cmp(&b.0).unwrap())
                    .unwrap_or((0.0, 0))
            })
            .collect()
    }

    /// Crop to a time window
    pub fn crop(&self, tmin: f64, tmax: f64) -> InverseResult<Self> {
        let start_idx = self.times.iter().position(|&t| t >= tmin).unwrap_or(0);

        let end_idx = self
            .times
            .iter()
            .rposition(|&t| t <= tmax)
            .unwrap_or(self.n_times() - 1)
            + 1;

        let new_data: Vec<Vec<f64>> = self
            .data
            .iter()
            .map(|src| src[start_idx..end_idx].to_vec())
            .collect();

        let new_times = self.times[start_idx..end_idx].to_vec();

        Ok(Self {
            data: new_data,
            times: new_times,
            vertices: self.vertices.clone(),
            free_orientation: self.free_orientation,
        })
    }

    /// Extract sources above a threshold
    pub fn threshold(&self, thresh: f64) -> Vec<usize> {
        let mag = self.magnitude();
        mag.iter()
            .enumerate()
            .filter(|(_, src)| src.iter().any(|&v| v > thresh))
            .map(|(i, _)| i)
            .collect()
    }
}

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

    #[test]
    fn test_source_estimate() {
        let data = vec![vec![1.0, 2.0, 3.0], vec![0.5, 1.5, 2.5]];
        let times = vec![0.0, 0.1, 0.2];
        let vertices = vec![0, 1];

        let stc = SourceEstimate::new(data, times, vertices, false);

        assert_eq!(stc.n_sources(), 2);
        assert_eq!(stc.n_times(), 3);
    }

    #[test]
    fn test_magnitude_fixed() {
        let data = vec![vec![-1.0, 2.0, -3.0], vec![0.5, -1.5, 2.5]];
        let stc = SourceEstimate::new(data, vec![0.0, 0.1, 0.2], vec![0, 1], false);

        let mag = stc.magnitude();
        assert!((mag[0][0] - 1.0).abs() < 1e-10);
        assert!((mag[0][2] - 3.0).abs() < 1e-10);
    }

    #[test]
    fn test_magnitude_free() {
        // 1 source with 3 orientations
        let data = vec![
            vec![3.0], // x
            vec![4.0], // y
            vec![0.0], // z
        ];
        let stc = SourceEstimate::new(data, vec![0.0], vec![0], true);

        let mag = stc.magnitude();
        assert!((mag[0][0] - 5.0).abs() < 1e-10); // sqrt(9+16) = 5
    }

    #[test]
    fn test_crop() {
        let data = vec![vec![1.0, 2.0, 3.0, 4.0, 5.0]];
        let times = vec![0.0, 0.1, 0.2, 0.3, 0.4];

        let stc = SourceEstimate::new(data, times, vec![0], false);
        let cropped = stc.crop(0.1, 0.3).unwrap();

        assert_eq!(cropped.n_times(), 3);
        assert!((cropped.times()[0] - 0.1).abs() < 1e-10);
    }
}