//! Epoched data for event-related analysis.

use crate::{ChannelInfo, Event, Events, NeuroData, NeuroResult, SampleRate};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Configuration for creating epochs
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EpochsConfig {
    /// Start time relative to event (negative = before event)
    pub tmin: f64,
    /// End time relative to event
    pub tmax: f64,
    /// Baseline correction window (None = no baseline correction)
    pub baseline: Option<(f64, f64)>,
    /// Reject epochs where any channel exceeds this threshold
    pub reject: Option<HashMap<String, f64>>,
    /// Decimate by this factor
    pub decim: usize,
}

impl Default for EpochsConfig {
    fn default() -> Self {
        Self {
            tmin: -0.2,
            tmax: 0.5,
            baseline: Some((-0.2, 0.0)),
            reject: None,
            decim: 1,
        }
    }
}

impl EpochsConfig {
    /// Create a new epochs configuration
    #[must_use]
    pub fn new(tmin: f64, tmax: f64) -> Self {
        Self {
            tmin,
            tmax,
            ..Default::default()
        }
    }

    /// Set baseline correction window
    #[must_use]
    pub fn with_baseline(mut self, start: f64, end: f64) -> Self {
        self.baseline = Some((start, end));
        self
    }

    /// Disable baseline correction
    #[must_use]
    pub fn no_baseline(mut self) -> Self {
        self.baseline = None;
        self
    }

    /// Set rejection threshold for a channel type
    pub fn reject_by_type(&mut self, ch_type: &str, threshold: f64) {
        self.reject
            .get_or_insert_with(HashMap::new)
            .insert(ch_type.to_string(), threshold);
    }
}

/// Collection of epochs (segmented data around events)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Epochs {
    /// Data: [n_epochs x n_channels x n_times]
    data: Vec<f64>,
    /// Number of epochs
    n_epochs: usize,
    /// Number of channels
    n_channels: usize,
    /// Number of time points per epoch
    n_times: usize,
    /// Sampling frequency
    sfreq: SampleRate,
    /// Start time relative to event
    tmin: f64,
    /// End time relative to event
    tmax: f64,
    /// Channel information
    channels: ChannelInfo,
    /// Events that define each epoch
    events: Vec<Event>,
    /// Indices of dropped epochs
    dropped: Vec<usize>,
    /// Selection (which epochs are currently active)
    selection: Vec<usize>,
}

impl Epochs {
    /// Create epochs from continuous data
    ///
    /// # Arguments
    /// * `raw` - Continuous data
    /// * `events` - Events to epoch around
    /// * `event_ids` - Which event values to include (None = all)
    /// * `config` - Epoching configuration
    pub fn from_data(
        raw: &NeuroData,
        events: &Events,
        event_ids: Option<&[i32]>,
        config: &EpochsConfig,
    ) -> NeuroResult<Self> {
        let sfreq = raw.sfreq();
        let n_channels = raw.n_channels();

        // Calculate sample indices
        let tmin_samples = (config.tmin * sfreq).round() as i64;
        let tmax_samples = (config.tmax * sfreq).round() as i64;
        let n_times = (tmax_samples - tmin_samples) as usize;

        // Filter events by event_ids if specified
        let selected_events: Vec<&Event> = events
            .iter()
            .filter(|e| event_ids.is_none_or(|ids| ids.contains(&e.value)))
            .collect();

        let n_epochs = selected_events.len();
        if n_epochs == 0 {
            return Err(crate::NeuroError::Event(
                "No events found matching the criteria".to_string(),
            ));
        }

        // Allocate data
        let mut data = vec![0.0; n_epochs * n_channels * n_times];
        let mut epoch_events = Vec::with_capacity(n_epochs);
        let mut dropped = Vec::new();

        // Extract epochs
        for (epoch_idx, event) in selected_events.iter().enumerate() {
            let event_sample = event.sample as i64;
            let start_sample = event_sample + tmin_samples;
            let end_sample = event_sample + tmax_samples;

            // Check if epoch is within data bounds
            if start_sample < 0 || end_sample as usize > raw.n_samples() {
                dropped.push(epoch_idx);
                continue;
            }

            let start_sample = start_sample as usize;

            // Copy data for this epoch
            for ch in 0..n_channels {
                let raw_ch = raw.get_channel(ch).unwrap();
                let epoch_offset = epoch_idx * n_channels * n_times + ch * n_times;

                for t in 0..n_times {
                    data[epoch_offset + t] = raw_ch[start_sample + t];
                }
            }

            epoch_events.push(**event);
        }

        let mut epochs = Self {
            data,
            n_epochs,
            n_channels,
            n_times,
            sfreq,
            tmin: config.tmin,
            tmax: config.tmax,
            channels: raw.channels().clone(),
            events: epoch_events,
            dropped,
            selection: (0..n_epochs).collect(),
        };

        // Apply baseline correction
        if let Some((bmin, bmax)) = config.baseline {
            epochs.apply_baseline(bmin, bmax)?;
        }

        Ok(epochs)
    }

    /// Number of epochs
    #[must_use]
    pub fn n_epochs(&self) -> usize {
        self.selection.len()
    }

    /// Total number of epochs (including dropped)
    #[must_use]
    pub fn n_epochs_total(&self) -> usize {
        self.n_epochs
    }

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

    /// Number of time points per epoch
    #[must_use]
    pub fn n_times(&self) -> usize {
        self.n_times
    }

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

    /// Start time relative to event
    #[must_use]
    pub fn tmin(&self) -> f64 {
        self.tmin
    }

    /// End time relative to event
    #[must_use]
    pub fn tmax(&self) -> f64 {
        self.tmax
    }

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

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

    /// Get events
    #[must_use]
    pub fn events(&self) -> &[Event] {
        &self.events
    }

    /// Get indices of dropped epochs
    #[must_use]
    pub fn dropped(&self) -> &[usize] {
        &self.dropped
    }

    /// Get data for a single epoch [n_channels x n_times]
    #[must_use]
    pub fn get_epoch(&self, epoch: usize) -> Option<Vec<Vec<f64>>> {
        if epoch >= self.n_epochs {
            return None;
        }

        let epoch_offset = epoch * self.n_channels * self.n_times;
        let mut result = Vec::with_capacity(self.n_channels);

        for ch in 0..self.n_channels {
            let ch_offset = epoch_offset + ch * self.n_times;
            result.push(self.data[ch_offset..ch_offset + self.n_times].to_vec());
        }

        Some(result)
    }

    /// Get data for a single epoch and channel
    #[must_use]
    pub fn get_epoch_channel(&self, epoch: usize, channel: usize) -> Option<&[f64]> {
        if epoch >= self.n_epochs || channel >= self.n_channels {
            return None;
        }

        let offset = epoch * self.n_channels * self.n_times + channel * self.n_times;
        Some(&self.data[offset..offset + self.n_times])
    }

    /// Apply baseline correction in-place
    pub fn apply_baseline(&mut self, bmin: f64, bmax: f64) -> NeuroResult<()> {
        // Convert time to samples
        let bmin_sample = ((bmin - self.tmin) * self.sfreq).round() as usize;
        let bmax_sample = ((bmax - self.tmin) * self.sfreq).round() as usize;

        if bmin_sample >= self.n_times || bmax_sample > self.n_times {
            return Err(crate::NeuroError::InvalidParameter(format!(
                "Baseline window [{bmin}, {bmax}] outside epoch range [{}, {}]",
                self.tmin, self.tmax
            )));
        }

        let baseline_len = bmax_sample - bmin_sample;

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

            for ch in 0..self.n_channels {
                let ch_offset = epoch_offset + ch * self.n_times;

                // Compute baseline mean
                let baseline_sum: f64 = self.data[ch_offset + bmin_sample..ch_offset + bmax_sample]
                    .iter()
                    .sum();
                let baseline_mean = baseline_sum / baseline_len as f64;

                // Subtract baseline
                for t in 0..self.n_times {
                    self.data[ch_offset + t] -= baseline_mean;
                }
            }
        }

        Ok(())
    }

    /// Compute evoked response (average across epochs)
    #[must_use]
    pub fn average(&self) -> Evoked {
        let n_active = self.selection.len();
        let mut averaged = vec![0.0; self.n_channels * self.n_times];

        for &epoch_idx in &self.selection {
            let epoch_offset = epoch_idx * self.n_channels * self.n_times;

            for ch in 0..self.n_channels {
                let ch_offset = ch * self.n_times;
                let epoch_ch_offset = epoch_offset + ch_offset;

                for t in 0..self.n_times {
                    averaged[ch_offset + t] += self.data[epoch_ch_offset + t];
                }
            }
        }

        // Divide by number of epochs
        for x in &mut averaged {
            *x /= n_active as f64;
        }

        Evoked {
            data: averaged,
            n_channels: self.n_channels,
            n_times: self.n_times,
            sfreq: self.sfreq,
            tmin: self.tmin,
            tmax: self.tmax,
            channels: self.channels.clone(),
            n_averaged: n_active,
            comment: String::new(),
        }
    }

    /// Drop epochs where any channel exceeds threshold
    pub fn drop_bad(&mut self, threshold: f64) {
        let mut new_selection = Vec::new();

        for &epoch_idx in &self.selection {
            let epoch_offset = epoch_idx * self.n_channels * self.n_times;
            let mut is_bad = false;

            'outer: for ch in 0..self.n_channels {
                let ch_offset = epoch_offset + ch * self.n_times;
                for t in 0..self.n_times {
                    if self.data[ch_offset + t].abs() > threshold {
                        is_bad = true;
                        break 'outer;
                    }
                }
            }

            if !is_bad {
                new_selection.push(epoch_idx);
            } else {
                self.dropped.push(epoch_idx);
            }
        }

        self.selection = new_selection;
    }
}

/// Evoked response (averaged epochs)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Evoked {
    /// Data: [n_channels x n_times]
    data: Vec<f64>,
    /// Number of channels
    n_channels: usize,
    /// Number of time points
    n_times: usize,
    /// Sampling frequency
    sfreq: SampleRate,
    /// Start time relative to event
    tmin: f64,
    /// End time relative to event
    tmax: f64,
    /// Channel information
    channels: ChannelInfo,
    /// Number of epochs averaged
    n_averaged: usize,
    /// Comment/description
    comment: String,
}

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

    /// Number of time points
    #[must_use]
    pub fn n_times(&self) -> usize {
        self.n_times
    }

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

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

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

    /// Number of epochs averaged
    #[must_use]
    pub fn n_averaged(&self) -> usize {
        self.n_averaged
    }

    /// Get raw data as slice
    #[must_use]
    pub fn data(&self) -> &[f64] {
        &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_times;
        Some(&self.data[start..start + self.n_times])
    }

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

    /// Find peak amplitude and latency for a channel
    #[must_use]
    pub fn peak(&self, ch: usize, tmin: Option<f64>, tmax: Option<f64>) -> Option<(f64, f64)> {
        let ch_data = self.get_channel(ch)?;
        let times = self.times();

        let tmin = tmin.unwrap_or(self.tmin);
        let tmax = tmax.unwrap_or(self.tmax);

        let mut max_val = f64::NEG_INFINITY;
        let mut max_time = 0.0;

        for (i, (&val, &t)) in ch_data.iter().zip(times.iter()).enumerate() {
            if t >= tmin && t <= tmax {
                let abs_val = val.abs();
                if abs_val > max_val {
                    max_val = abs_val;
                    max_time = t;
                    if val < 0.0 {
                        max_val = -max_val;
                    }
                }
            }
        }

        if max_val.is_finite() {
            Some((max_val, max_time))
        } else {
            None
        }
    }
}

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

    fn create_test_epochs() -> Epochs {
        // Create simple test data: 2 epochs, 2 channels, 100 samples each
        let n_epochs = 2;
        let n_channels = 2;
        let n_times = 100;
        let sfreq = 1000.0;

        let mut data = Vec::with_capacity(n_epochs * n_channels * n_times);
        for epoch in 0..n_epochs {
            for ch in 0..n_channels {
                for t in 0..n_times {
                    data.push((epoch * 100 + ch * 10 + t) as f64);
                }
            }
        }

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

        Epochs {
            data,
            n_epochs,
            n_channels,
            n_times,
            sfreq,
            tmin: -0.05,
            tmax: 0.05,
            channels,
            events: vec![Event::new(500, 1), Event::new(1500, 1)],
            dropped: Vec::new(),
            selection: vec![0, 1],
        }
    }

    #[test]
    fn test_epochs_average() {
        let epochs = create_test_epochs();
        let evoked = epochs.average();

        assert_eq!(evoked.n_channels(), 2);
        assert_eq!(evoked.n_times(), 100);
        assert_eq!(evoked.n_averaged(), 2);

        // Average of ch0 at t0: (0 + 100) / 2 = 50
        // Epoch 0: 0*100 + 0*10 + 0 = 0
        // Epoch 1: 1*100 + 0*10 + 0 = 100
        let ch0 = evoked.get_channel(0).unwrap();
        assert!((ch0[0] - 50.0).abs() < 1e-10);
    }

    #[test]
    fn test_epochs_config() {
        let config = EpochsConfig::new(-0.1, 0.4).with_baseline(-0.1, 0.0);

        assert_eq!(config.tmin, -0.1);
        assert_eq!(config.tmax, 0.4);
        assert_eq!(config.baseline, Some((-0.1, 0.0)));
    }
}