//! Core data structures for neuroimaging data.

use crate::{Annotations, ChannelInfo, Events, NeuroResult, SampleRate};
use serde::{Deserialize, Serialize};

/// Main container for continuous neuroimaging data.
///
/// This structure holds the raw or processed signal data along with
/// all associated metadata (channels, events, sampling rate, etc.).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NeuroData {
    /// Signal data as a 2D array [n_channels x n_samples]
    /// Stored as row-major (each row is a channel's time series)
    data: Vec<f64>,
    /// Number of channels
    n_channels: usize,
    /// Number of samples per channel
    n_samples: usize,
    /// Sampling frequency in Hz
    sfreq: SampleRate,
    /// Channel information
    channels: ChannelInfo,
    /// Events/triggers
    events: Events,
    /// Annotations
    annotations: Annotations,
    /// First sample number (for time alignment with original recording)
    first_sample: usize,
}

impl NeuroData {
    /// Create new `NeuroData` from raw data.
    ///
    /// # Arguments
    /// * `data` - Flattened data array [n_channels x n_samples] in row-major order
    /// * `n_channels` - Number of channels
    /// * `sfreq` - Sampling frequency in Hz
    /// * `channels` - Channel information
    ///
    /// # Errors
    /// Returns error if data length doesn't match n_channels * n_samples
    pub fn new(
        data: Vec<f64>,
        n_channels: usize,
        sfreq: SampleRate,
        channels: ChannelInfo,
    ) -> NeuroResult<Self> {
        if data.len() % n_channels != 0 {
            return Err(crate::NeuroError::DimensionMismatch {
                expected: format!("multiple of {n_channels}"),
                got: format!("{}", data.len()),
            });
        }

        let n_samples = data.len() / n_channels;

        if channels.len() != n_channels {
            return Err(crate::NeuroError::DimensionMismatch {
                expected: format!("{n_channels} channels"),
                got: format!("{} channel info entries", channels.len()),
            });
        }

        Ok(Self {
            data,
            n_channels,
            n_samples,
            sfreq,
            channels,
            events: Events::new(),
            annotations: Annotations::new(),
            first_sample: 0,
        })
    }

    /// Create empty `NeuroData` with specified dimensions
    #[must_use]
    pub fn zeros(n_channels: usize, n_samples: usize, sfreq: SampleRate) -> Self {
        let mut channels = ChannelInfo::new();
        for i in 0..n_channels {
            channels.add_channel(crate::Channel::new(
                format!("CH{i:03}"),
                crate::ChannelType::EegScalp,
            ));
        }

        Self {
            data: vec![0.0; n_channels * n_samples],
            n_channels,
            n_samples,
            sfreq,
            channels,
            events: Events::new(),
            annotations: Annotations::new(),
            first_sample: 0,
        }
    }

    /// Number of channels
    #[must_use]
    pub fn n_channels(&self) -> usize {
        self.n_channels
    }

    /// Number of samples per channel
    #[must_use]
    pub fn n_samples(&self) -> usize {
        self.n_samples
    }

    /// Sampling frequency in Hz
    #[must_use]
    pub fn sfreq(&self) -> SampleRate {
        self.sfreq
    }

    /// Duration in seconds
    #[must_use]
    pub fn duration(&self) -> f64 {
        self.n_samples as f64 / self.sfreq
    }

    /// First sample number
    #[must_use]
    pub fn first_sample(&self) -> usize {
        self.first_sample
    }

    /// Last sample number (exclusive)
    #[must_use]
    pub fn last_sample(&self) -> usize {
        self.first_sample + self.n_samples
    }

    /// Get time vector in seconds
    #[must_use]
    pub fn times(&self) -> Vec<f64> {
        (0..self.n_samples).map(|i| i as f64 / self.sfreq).collect()
    }

    /// Reference to channel information
    #[must_use]
    pub fn channels(&self) -> &ChannelInfo {
        &self.channels
    }

    /// Mutable reference to channel information
    pub fn channels_mut(&mut self) -> &mut ChannelInfo {
        &mut self.channels
    }

    /// Reference to events
    #[must_use]
    pub fn events(&self) -> &Events {
        &self.events
    }

    /// Mutable reference to events
    pub fn events_mut(&mut self) -> &mut Events {
        &mut self.events
    }

    /// Set events
    pub fn set_events(&mut self, events: Events) {
        self.events = events;
    }

    /// Reference to annotations
    #[must_use]
    pub fn annotations(&self) -> &Annotations {
        &self.annotations
    }

    /// Mutable reference to annotations
    pub fn annotations_mut(&mut self) -> &mut Annotations {
        &mut self.annotations
    }

    /// Get raw data as slice (row-major: [ch0_t0, ch0_t1, ..., ch1_t0, ch1_t1, ...])
    #[must_use]
    pub fn data(&self) -> &[f64] {
        &self.data
    }

    /// Get mutable raw data
    pub fn data_mut(&mut self) -> &mut [f64] {
        &mut self.data
    }

    /// Get data for a single channel
    #[must_use]
    pub fn get_channel(&self, ch: usize) -> Option<&[f64]> {
        if ch >= self.n_channels {
            return None;
        }
        let start = ch * self.n_samples;
        let end = start + self.n_samples;
        Some(&self.data[start..end])
    }

    /// Get mutable data for a single channel
    pub fn get_channel_mut(&mut self, ch: usize) -> Option<&mut [f64]> {
        if ch >= self.n_channels {
            return None;
        }
        let start = ch * self.n_samples;
        let end = start + self.n_samples;
        Some(&mut self.data[start..end])
    }

    /// Get a single sample value
    #[must_use]
    pub fn get(&self, ch: usize, sample: usize) -> Option<f64> {
        if ch >= self.n_channels || sample >= self.n_samples {
            return None;
        }
        Some(self.data[ch * self.n_samples + sample])
    }

    /// Set a single sample value
    pub fn set(&mut self, ch: usize, sample: usize, value: f64) -> bool {
        if ch >= self.n_channels || sample >= self.n_samples {
            return false;
        }
        self.data[ch * self.n_samples + sample] = value;
        true
    }

    /// Get data as 2D vector [n_channels][n_samples]
    #[must_use]
    pub fn to_2d(&self) -> Vec<Vec<f64>> {
        (0..self.n_channels)
            .map(|ch| self.get_channel(ch).unwrap().to_vec())
            .collect()
    }

    /// Crop data to a time range
    ///
    /// # Arguments
    /// * `tmin` - Start time in seconds
    /// * `tmax` - End time in seconds
    #[must_use]
    pub fn crop(&self, tmin: f64, tmax: f64) -> Self {
        let start_sample = (tmin * self.sfreq).round() as usize;
        let end_sample = (tmax * self.sfreq).round() as usize;

        let start_sample = start_sample.min(self.n_samples);
        let end_sample = end_sample.min(self.n_samples).max(start_sample);

        let new_n_samples = end_sample - start_sample;
        let mut new_data = Vec::with_capacity(self.n_channels * new_n_samples);

        for ch in 0..self.n_channels {
            let ch_data = self.get_channel(ch).unwrap();
            new_data.extend_from_slice(&ch_data[start_sample..end_sample]);
        }

        Self {
            data: new_data,
            n_channels: self.n_channels,
            n_samples: new_n_samples,
            sfreq: self.sfreq,
            channels: self.channels.clone(),
            events: Events::new(), // Events would need to be adjusted
            annotations: Annotations::new(),
            first_sample: self.first_sample + start_sample,
        }
    }

    /// Pick subset of channels
    pub fn pick_channels(&self, indices: &[usize]) -> NeuroResult<Self> {
        // Validate indices
        for &idx in indices {
            if idx >= self.n_channels {
                return Err(crate::NeuroError::Channel(format!(
                    "Channel index {idx} out of range (max {})",
                    self.n_channels - 1
                )));
            }
        }

        let new_n_channels = indices.len();
        let mut new_data = Vec::with_capacity(new_n_channels * self.n_samples);
        let mut new_channels = ChannelInfo::new();

        for &idx in indices {
            new_data.extend_from_slice(self.get_channel(idx).unwrap());
            new_channels.add_channel(self.channels[idx].clone());
        }

        Ok(Self {
            data: new_data,
            n_channels: new_n_channels,
            n_samples: self.n_samples,
            sfreq: self.sfreq,
            channels: new_channels,
            events: self.events.clone(),
            annotations: self.annotations.clone(),
            first_sample: self.first_sample,
        })
    }

    /// Apply a function to each sample in-place
    pub fn apply<F>(&mut self, f: F)
    where
        F: Fn(f64) -> f64,
    {
        for x in &mut self.data {
            *x = f(*x);
        }
    }

    /// Apply a function to each channel in-place
    pub fn apply_per_channel<F>(&mut self, f: F)
    where
        F: Fn(&mut [f64]),
    {
        for ch in 0..self.n_channels {
            let start = ch * self.n_samples;
            let end = start + self.n_samples;
            f(&mut self.data[start..end]);
        }
    }

    /// Compute mean across all samples for each channel
    #[must_use]
    pub fn mean_per_channel(&self) -> Vec<f64> {
        (0..self.n_channels)
            .map(|ch| {
                let data = self.get_channel(ch).unwrap();
                data.iter().sum::<f64>() / data.len() as f64
            })
            .collect()
    }

    /// Compute standard deviation for each channel
    #[must_use]
    pub fn std_per_channel(&self) -> Vec<f64> {
        let means = self.mean_per_channel();
        (0..self.n_channels)
            .map(|ch| {
                let data = self.get_channel(ch).unwrap();
                let mean = means[ch];
                let variance =
                    data.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (data.len() - 1) as f64;
                variance.sqrt()
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Channel, ChannelType};

    fn create_test_data() -> NeuroData {
        let n_channels = 3;
        let n_samples = 100;
        let sfreq = 1000.0;

        // Create test data: each channel has values 0, 1, 2, ..., 99
        let mut data = Vec::with_capacity(n_channels * n_samples);
        for _ch in 0..n_channels {
            for s in 0..n_samples {
                data.push(s as f64);
            }
        }

        let mut channels = ChannelInfo::new();
        channels.add_channel(Channel::new("Ch1", ChannelType::EegScalp));
        channels.add_channel(Channel::new("Ch2", ChannelType::EegScalp));
        channels.add_channel(Channel::new("Ch3", ChannelType::EegScalp));

        NeuroData::new(data, n_channels, sfreq, channels).unwrap()
    }

    #[test]
    fn test_neuro_data_creation() {
        let data = create_test_data();
        assert_eq!(data.n_channels(), 3);
        assert_eq!(data.n_samples(), 100);
        assert_eq!(data.sfreq(), 1000.0);
        assert!((data.duration() - 0.1).abs() < 1e-10);
    }

    #[test]
    fn test_get_channel() {
        let data = create_test_data();
        let ch0 = data.get_channel(0).unwrap();
        assert_eq!(ch0.len(), 100);
        assert_eq!(ch0[0], 0.0);
        assert_eq!(ch0[99], 99.0);
    }

    #[test]
    fn test_crop() {
        let data = create_test_data();
        let cropped = data.crop(0.01, 0.05); // 10-50 samples at 1000 Hz
        assert_eq!(cropped.n_samples(), 40);
        assert_eq!(cropped.first_sample(), 10);
    }

    #[test]
    fn test_pick_channels() {
        let data = create_test_data();
        let picked = data.pick_channels(&[0, 2]).unwrap();
        assert_eq!(picked.n_channels(), 2);
        assert_eq!(picked.channels().names(), vec!["Ch1", "Ch3"]);
    }
}