//! Artifact type definitions and labels.
//!
//! Defines the types of artifacts that can be detected in MEG/EEG recordings.

use serde::{Deserialize, Serialize};

/// Types of artifacts that can be detected
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ArtifactType {
    /// Eye blink artifact (EOG)
    /// Characterized by large frontal deflections lasting 100-400ms
    EyeBlink,

    /// Eye movement artifact (saccades, smooth pursuit)
    /// Horizontal or vertical eye movements
    EyeMovement,

    /// Muscle artifact (EMG)
    /// High-frequency (>20Hz) contamination from facial/neck muscles
    Muscle,

    /// Heartbeat artifact (ECG/BCG)
    /// Cardiac-related artifact, especially prominent in MEG
    Heartbeat,

    /// Power line noise (50/60 Hz)
    /// Narrow-band interference from electrical mains
    LineNoise,

    /// Movement artifact
    /// Head or body movement causing signal distortion
    Movement,

    /// Electrode pop/jump
    /// Sudden amplitude changes from electrode issues
    ElectrodePop,

    /// Channel noise/bad channel
    /// Persistently noisy or disconnected channel
    ChannelNoise,

    /// Environmental interference
    /// External electromagnetic interference
    Environmental,

    /// Unknown artifact
    /// Detected anomaly of unknown origin
    Unknown,
}

impl ArtifactType {
    /// Get the index for this artifact type (for model output mapping)
    pub fn index(&self) -> usize {
        match self {
            Self::EyeBlink => 0,
            Self::EyeMovement => 1,
            Self::Muscle => 2,
            Self::Heartbeat => 3,
            Self::LineNoise => 4,
            Self::Movement => 5,
            Self::ElectrodePop => 6,
            Self::ChannelNoise => 7,
            Self::Environmental => 8,
            Self::Unknown => 9,
        }
    }

    /// Create from index
    pub fn from_index(index: usize) -> Option<Self> {
        match index {
            0 => Some(Self::EyeBlink),
            1 => Some(Self::EyeMovement),
            2 => Some(Self::Muscle),
            3 => Some(Self::Heartbeat),
            4 => Some(Self::LineNoise),
            5 => Some(Self::Movement),
            6 => Some(Self::ElectrodePop),
            7 => Some(Self::ChannelNoise),
            8 => Some(Self::Environmental),
            9 => Some(Self::Unknown),
            _ => None,
        }
    }

    /// Number of artifact types
    pub const fn count() -> usize {
        10
    }

    /// Get all artifact types
    pub fn all() -> Vec<Self> {
        vec![
            Self::EyeBlink,
            Self::EyeMovement,
            Self::Muscle,
            Self::Heartbeat,
            Self::LineNoise,
            Self::Movement,
            Self::ElectrodePop,
            Self::ChannelNoise,
            Self::Environmental,
            Self::Unknown,
        ]
    }

    /// Get human-readable name
    pub fn name(&self) -> &'static str {
        match self {
            Self::EyeBlink => "Eye Blink",
            Self::EyeMovement => "Eye Movement",
            Self::Muscle => "Muscle",
            Self::Heartbeat => "Heartbeat",
            Self::LineNoise => "Line Noise",
            Self::Movement => "Movement",
            Self::ElectrodePop => "Electrode Pop",
            Self::ChannelNoise => "Channel Noise",
            Self::Environmental => "Environmental",
            Self::Unknown => "Unknown",
        }
    }

    /// Get short code
    pub fn code(&self) -> &'static str {
        match self {
            Self::EyeBlink => "EOG_B",
            Self::EyeMovement => "EOG_M",
            Self::Muscle => "EMG",
            Self::Heartbeat => "ECG",
            Self::LineNoise => "LINE",
            Self::Movement => "MOV",
            Self::ElectrodePop => "POP",
            Self::ChannelNoise => "BAD",
            Self::Environmental => "ENV",
            Self::Unknown => "UNK",
        }
    }

    /// Get typical frequency range for this artifact type
    pub fn frequency_range(&self) -> (f64, f64) {
        match self {
            Self::EyeBlink => (0.5, 4.0),
            Self::EyeMovement => (0.5, 10.0),
            Self::Muscle => (20.0, 200.0),
            Self::Heartbeat => (0.5, 40.0),
            Self::LineNoise => (49.0, 61.0), // Covers both 50 and 60 Hz
            Self::Movement => (0.1, 5.0),
            Self::ElectrodePop => (0.1, 100.0),  // Broadband
            Self::ChannelNoise => (0.1, 200.0),  // Broadband
            Self::Environmental => (0.1, 200.0), // Variable
            Self::Unknown => (0.1, 200.0),
        }
    }

    /// Get typical duration range in milliseconds
    pub fn duration_range_ms(&self) -> (f64, f64) {
        match self {
            Self::EyeBlink => (100.0, 400.0),
            Self::EyeMovement => (50.0, 500.0),
            Self::Muscle => (10.0, 1000.0),    // Can be sustained
            Self::Heartbeat => (200.0, 400.0), // QRS complex
            Self::LineNoise => (0.0, f64::INFINITY), // Continuous
            Self::Movement => (100.0, 5000.0),
            Self::ElectrodePop => (1.0, 50.0), // Very brief
            Self::ChannelNoise => (0.0, f64::INFINITY), // Continuous
            Self::Environmental => (0.0, f64::INFINITY), // Variable
            Self::Unknown => (0.0, f64::INFINITY),
        }
    }

    /// Get color for visualization (RGB)
    pub fn color(&self) -> (u8, u8, u8) {
        match self {
            Self::EyeBlink => (255, 165, 0),     // Orange
            Self::EyeMovement => (255, 200, 0),  // Gold
            Self::Muscle => (255, 0, 0),         // Red
            Self::Heartbeat => (255, 0, 128),    // Pink
            Self::LineNoise => (128, 128, 128),  // Gray
            Self::Movement => (0, 128, 255),     // Blue
            Self::ElectrodePop => (255, 255, 0), // Yellow
            Self::ChannelNoise => (128, 0, 128), // Purple
            Self::Environmental => (0, 128, 0),  // Green
            Self::Unknown => (64, 64, 64),       // Dark gray
        }
    }
}

impl std::fmt::Display for ArtifactType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.name())
    }
}

/// A detected artifact label with probability
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArtifactLabel {
    /// Type of artifact
    pub artifact_type: ArtifactType,
    /// Detection probability (0.0 - 1.0)
    pub probability: f64,
    /// Confidence score (may differ from probability for some models)
    pub confidence: f64,
}

impl ArtifactLabel {
    /// Create a new artifact label
    pub fn new(artifact_type: ArtifactType, probability: f64) -> Self {
        Self {
            artifact_type,
            probability,
            confidence: probability,
        }
    }

    /// Create with separate confidence score
    pub fn with_confidence(artifact_type: ArtifactType, probability: f64, confidence: f64) -> Self {
        Self {
            artifact_type,
            probability,
            confidence,
        }
    }

    /// Check if this artifact is considered detected (probability > threshold)
    pub fn is_detected(&self, threshold: f64) -> bool {
        self.probability > threshold
    }
}

/// A region in the signal where an artifact was detected
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ArtifactRegion {
    /// Artifact type
    pub artifact_type: ArtifactType,
    /// Start time in seconds
    pub start_time: f64,
    /// End time in seconds
    pub end_time: f64,
    /// Affected channel indices
    pub channels: Vec<usize>,
    /// Detection probability
    pub probability: f64,
    /// Peak location (time of maximum artifact influence)
    pub peak_time: Option<f64>,
    /// Severity score (0.0 - 1.0)
    pub severity: f64,
}

impl ArtifactRegion {
    /// Create a new artifact region
    pub fn new(
        artifact_type: ArtifactType,
        start_time: f64,
        end_time: f64,
        channels: Vec<usize>,
        probability: f64,
    ) -> Self {
        Self {
            artifact_type,
            start_time,
            end_time,
            channels,
            probability,
            peak_time: None,
            severity: probability,
        }
    }

    /// Duration in seconds
    pub fn duration(&self) -> f64 {
        self.end_time - self.start_time
    }

    /// Duration in milliseconds
    pub fn duration_ms(&self) -> f64 {
        self.duration() * 1000.0
    }

    /// Check if this region overlaps with a time range
    pub fn overlaps(&self, start: f64, end: f64) -> bool {
        self.start_time < end && self.end_time > start
    }

    /// Check if this region affects a specific channel
    pub fn affects_channel(&self, channel: usize) -> bool {
        self.channels.contains(&channel)
    }

    /// Merge with another overlapping region of the same type
    pub fn merge(&self, other: &Self) -> Option<Self> {
        if self.artifact_type != other.artifact_type {
            return None;
        }

        if !self.overlaps(other.start_time, other.end_time) {
            return None;
        }

        let mut channels: Vec<usize> = self.channels.clone();
        for ch in &other.channels {
            if !channels.contains(ch) {
                channels.push(*ch);
            }
        }
        channels.sort_unstable();

        Some(Self {
            artifact_type: self.artifact_type,
            start_time: self.start_time.min(other.start_time),
            end_time: self.end_time.max(other.end_time),
            channels,
            probability: self.probability.max(other.probability),
            peak_time: match (self.peak_time, other.peak_time) {
                (Some(a), Some(b)) => Some(if self.probability > other.probability {
                    a
                } else {
                    b
                }),
                (Some(a), None) => Some(a),
                (None, Some(b)) => Some(b),
                (None, None) => None,
            },
            severity: self.severity.max(other.severity),
        })
    }
}

/// Summary statistics for artifact detection
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ArtifactSummary {
    /// Total number of artifacts detected
    pub total_artifacts: usize,
    /// Count per artifact type
    pub counts: Vec<(ArtifactType, usize)>,
    /// Total contaminated duration (seconds)
    pub total_duration: f64,
    /// Percentage of signal contaminated
    pub contamination_percent: f64,
    /// Most common artifact type
    pub most_common: Option<ArtifactType>,
    /// Channels most affected
    pub most_affected_channels: Vec<usize>,
}

impl ArtifactSummary {
    /// Create from a list of artifact regions
    pub fn from_regions(regions: &[ArtifactRegion], total_signal_duration: f64) -> Self {
        let mut counts: std::collections::HashMap<ArtifactType, usize> =
            std::collections::HashMap::new();
        let mut total_duration = 0.0;
        let mut channel_counts: std::collections::HashMap<usize, usize> =
            std::collections::HashMap::new();

        for region in regions {
            *counts.entry(region.artifact_type).or_insert(0) += 1;
            total_duration += region.duration();

            for &ch in &region.channels {
                *channel_counts.entry(ch).or_insert(0) += 1;
            }
        }

        let counts_vec: Vec<(ArtifactType, usize)> = counts.into_iter().collect();
        let most_common = counts_vec.iter().max_by_key(|(_, c)| *c).map(|(t, _)| *t);

        let mut channel_vec: Vec<(usize, usize)> = channel_counts.into_iter().collect();
        channel_vec.sort_by(|a, b| b.1.cmp(&a.1));
        let most_affected_channels: Vec<usize> =
            channel_vec.into_iter().take(5).map(|(ch, _)| ch).collect();

        let contamination_percent = if total_signal_duration > 0.0 {
            (total_duration / total_signal_duration) * 100.0
        } else {
            0.0
        };

        Self {
            total_artifacts: regions.len(),
            counts: counts_vec,
            total_duration,
            contamination_percent,
            most_common,
            most_affected_channels,
        }
    }
}

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

    #[test]
    fn test_artifact_type_index() {
        assert_eq!(ArtifactType::EyeBlink.index(), 0);
        assert_eq!(ArtifactType::Muscle.index(), 2);
        assert_eq!(ArtifactType::from_index(2), Some(ArtifactType::Muscle));
    }

    #[test]
    fn test_artifact_type_all() {
        let all = ArtifactType::all();
        assert_eq!(all.len(), ArtifactType::count());
    }

    #[test]
    fn test_artifact_label() {
        let label = ArtifactLabel::new(ArtifactType::EyeBlink, 0.85);
        assert!(label.is_detected(0.5));
        assert!(!label.is_detected(0.9));
    }

    #[test]
    fn test_artifact_region() {
        let region = ArtifactRegion::new(ArtifactType::Muscle, 1.0, 1.5, vec![10, 11, 12], 0.9);

        assert_eq!(region.duration(), 0.5);
        assert_eq!(region.duration_ms(), 500.0);
        assert!(region.affects_channel(10));
        assert!(!region.affects_channel(5));
    }

    #[test]
    fn test_region_overlap() {
        let region = ArtifactRegion::new(ArtifactType::Muscle, 1.0, 2.0, vec![0], 0.9);

        assert!(region.overlaps(0.5, 1.5));
        assert!(region.overlaps(1.5, 2.5));
        assert!(!region.overlaps(2.5, 3.0));
        assert!(!region.overlaps(0.0, 0.5));
    }

    #[test]
    fn test_region_merge() {
        let r1 = ArtifactRegion::new(ArtifactType::Muscle, 1.0, 2.0, vec![0, 1], 0.8);
        let r2 = ArtifactRegion::new(ArtifactType::Muscle, 1.5, 2.5, vec![1, 2], 0.9);

        let merged = r1.merge(&r2).unwrap();
        assert_eq!(merged.start_time, 1.0);
        assert_eq!(merged.end_time, 2.5);
        assert_eq!(merged.channels, vec![0, 1, 2]);
        assert_eq!(merged.probability, 0.9);
    }

    #[test]
    fn test_artifact_summary() {
        let regions = vec![
            ArtifactRegion::new(ArtifactType::EyeBlink, 0.0, 0.3, vec![0, 1], 0.9),
            ArtifactRegion::new(ArtifactType::EyeBlink, 1.0, 1.3, vec![0, 1], 0.85),
            ArtifactRegion::new(ArtifactType::Muscle, 2.0, 2.5, vec![10], 0.7),
        ];

        let summary = ArtifactSummary::from_regions(&regions, 10.0);

        assert_eq!(summary.total_artifacts, 3);
        assert!((summary.total_duration - 1.1).abs() < 0.01);
        assert_eq!(summary.most_common, Some(ArtifactType::EyeBlink));
    }
}