//! ONNX-based artifact detector.
//!
//! Provides transformer-based artifact detection using ONNX inference.

use crate::error::{ArtifactError, ArtifactResult};
use crate::labels::{ArtifactLabel, ArtifactRegion, ArtifactSummary, ArtifactType};
use rtx_onnx::{OnnxSession, OnnxSessionConfig};
use rtx_tensor::{Device, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;

/// Configuration for the artifact detector
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetectorConfig {
    /// Model path (ONNX file)
    pub model_path: Option<String>,
    /// Detection threshold (0.0 - 1.0)
    pub threshold: f64,
    /// Window size in samples for processing
    pub window_size: usize,
    /// Overlap between consecutive windows (0.0 - 1.0)
    pub overlap: f64,
    /// Batch size for inference
    pub batch_size: usize,
    /// Number of input channels
    pub n_channels: usize,
    /// Sampling frequency
    pub sfreq: f64,
    /// Which artifact types to detect (None = all)
    pub artifact_types: Option<Vec<ArtifactType>>,
    /// Use GPU if available
    pub use_gpu: bool,
}

impl Default for DetectorConfig {
    fn default() -> Self {
        Self {
            model_path: None,
            threshold: 0.5,
            window_size: 1000, // 1 second at 1kHz
            overlap: 0.5,
            batch_size: 32,
            n_channels: 64,
            sfreq: 1000.0,
            artifact_types: None,
            use_gpu: false,
        }
    }
}

impl DetectorConfig {
    /// Create config for EEG data
    pub fn eeg(n_channels: usize, sfreq: f64) -> Self {
        Self {
            n_channels,
            sfreq,
            window_size: (sfreq * 1.0) as usize, // 1 second windows
            ..Default::default()
        }
    }

    /// Create config for MEG data
    pub fn meg(n_channels: usize, sfreq: f64) -> Self {
        Self {
            n_channels,
            sfreq,
            window_size: (sfreq * 1.0) as usize,
            threshold: 0.4, // Slightly more sensitive for MEG
            ..Default::default()
        }
    }

    /// Set detection threshold
    pub fn with_threshold(mut self, threshold: f64) -> Self {
        self.threshold = threshold.clamp(0.0, 1.0);
        self
    }

    /// Set window size
    pub fn with_window_size(mut self, window_size: usize) -> Self {
        self.window_size = window_size;
        self
    }

    /// Set overlap
    pub fn with_overlap(mut self, overlap: f64) -> Self {
        self.overlap = overlap.clamp(0.0, 0.99);
        self
    }

    /// Set specific artifact types to detect
    pub fn with_artifact_types(mut self, types: Vec<ArtifactType>) -> Self {
        self.artifact_types = Some(types);
        self
    }

    /// Enable GPU acceleration
    pub fn with_gpu(mut self, use_gpu: bool) -> Self {
        self.use_gpu = use_gpu;
        self
    }
}

/// Result of artifact detection on a single window
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetectionResult {
    /// Artifact labels with probabilities
    pub labels: Vec<ArtifactLabel>,
    /// Window start time in seconds
    pub start_time: f64,
    /// Window end time in seconds
    pub end_time: f64,
    /// Raw model output probabilities
    pub raw_probabilities: Vec<f64>,
}

impl DetectionResult {
    /// Get detected artifacts (above threshold)
    pub fn detected(&self, threshold: f64) -> Vec<&ArtifactLabel> {
        self.labels
            .iter()
            .filter(|l| l.is_detected(threshold))
            .collect()
    }

    /// Get predictions as (ArtifactType, probability) pairs
    pub fn predictions(&self) -> impl Iterator<Item = (ArtifactType, f64)> + '_ {
        self.labels.iter().map(|l| (l.artifact_type, l.probability))
    }

    /// Check if any artifact was detected
    pub fn has_artifact(&self, threshold: f64) -> bool {
        self.labels.iter().any(|l| l.is_detected(threshold))
    }

    /// Get the most likely artifact type
    pub fn most_likely(&self) -> Option<&ArtifactLabel> {
        self.labels
            .iter()
            .max_by(|a, b| a.probability.partial_cmp(&b.probability).unwrap())
    }
}

/// Batch detection result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetectionBatch {
    /// Results for each window
    pub results: Vec<DetectionResult>,
    /// Total processing time in milliseconds
    pub processing_time_ms: f64,
    /// Number of windows processed
    pub n_windows: usize,
}

impl DetectionBatch {
    /// Convert to artifact regions
    pub fn to_regions(&self, threshold: f64) -> Vec<ArtifactRegion> {
        let mut regions: Vec<ArtifactRegion> = Vec::new();

        for result in &self.results {
            for label in &result.labels {
                if label.is_detected(threshold) {
                    let region = ArtifactRegion::new(
                        label.artifact_type,
                        result.start_time,
                        result.end_time,
                        Vec::new(), // Channel info not available at this level
                        label.probability,
                    );

                    // Try to merge with existing regions
                    let mut merged = false;
                    for existing in &mut regions {
                        if let Some(merged_region) = existing.merge(&region) {
                            *existing = merged_region;
                            merged = true;
                            break;
                        }
                    }

                    if !merged {
                        regions.push(region);
                    }
                }
            }
        }

        regions
    }

    /// Get summary statistics
    pub fn summary(&self, total_duration: f64, threshold: f64) -> ArtifactSummary {
        let regions = self.to_regions(threshold);
        ArtifactSummary::from_regions(&regions, total_duration)
    }

    /// Get all detected artifact types
    pub fn detected_types(&self, threshold: f64) -> Vec<ArtifactType> {
        let mut types: Vec<ArtifactType> = Vec::new();
        for result in &self.results {
            for label in &result.labels {
                if label.is_detected(threshold) && !types.contains(&label.artifact_type) {
                    types.push(label.artifact_type);
                }
            }
        }
        types
    }
}

/// ONNX-based artifact detector
pub struct ArtifactDetector {
    /// Configuration
    config: DetectorConfig,
    /// ONNX session
    session: Option<OnnxSession>,
    /// Whether model is loaded
    model_loaded: bool,
    /// Input name for ONNX model
    input_name: String,
    /// Output name for ONNX model
    output_name: String,
}

impl ArtifactDetector {
    /// Create a new artifact detector
    pub fn new(config: DetectorConfig) -> ArtifactResult<Self> {
        // Clone the model path before creating the detector
        let model_path = config.model_path.clone();

        let mut detector = Self {
            config,
            session: None,
            model_loaded: false,
            input_name: "input".to_string(),
            output_name: "output".to_string(),
        };

        // Load model if path provided
        if let Some(ref path) = model_path {
            detector.load_model(path)?;
        }

        Ok(detector)
    }

    /// Load an ONNX model from path
    pub fn load_model(&mut self, path: impl AsRef<Path>) -> ArtifactResult<()> {
        let path = path.as_ref();

        if !path.exists() {
            return Err(ArtifactError::Model(format!(
                "Model file not found: {}",
                path.display()
            )));
        }

        // Create ONNX session config
        let onnx_config = OnnxSessionConfig::default();

        // Create session
        let session = OnnxSession::from_file(path, onnx_config)
            .map_err(|e| ArtifactError::Onnx(e.to_string()))?;

        // Get input/output names from session
        let input_names = session.input_names();
        let output_names = session.output_names();

        if let Some(first_input) = input_names.first() {
            self.input_name = first_input.clone();
        }
        if let Some(first_output) = output_names.first() {
            self.output_name = first_output.clone();
        }

        self.session = Some(session);
        self.model_loaded = true;

        Ok(())
    }

    /// Check if model is loaded
    pub fn is_loaded(&self) -> bool {
        self.model_loaded
    }

    /// Get configuration
    pub fn config(&self) -> &DetectorConfig {
        &self.config
    }

    /// Detect artifacts in data
    ///
    /// Input: [channels x time] EEG/MEG data
    pub fn detect(&self, data: &[Vec<f64>]) -> ArtifactResult<DetectionResult> {
        if data.is_empty() {
            return Err(ArtifactError::Input("Empty data".to_string()));
        }

        let n_channels = data.len();
        let n_samples = data[0].len();

        // Validate dimensions
        if n_channels != self.config.n_channels {
            return Err(ArtifactError::DimensionMismatch(format!(
                "Expected {} channels, got {}",
                self.config.n_channels, n_channels
            )));
        }

        // Run inference or use fallback
        let probabilities = if let Some(ref session) = self.session {
            self.run_inference(session, data)?
        } else {
            // Fallback: rule-based detection when no model loaded
            self.rule_based_detect(data)?
        };

        // Create labels from probabilities
        let labels: Vec<ArtifactLabel> = probabilities
            .iter()
            .enumerate()
            .filter_map(|(i, &prob)| {
                ArtifactType::from_index(i)
                    .map(|artifact_type| ArtifactLabel::new(artifact_type, prob))
            })
            .collect();

        Ok(DetectionResult {
            labels,
            start_time: 0.0,
            end_time: n_samples as f64 / self.config.sfreq,
            raw_probabilities: probabilities,
        })
    }

    /// Detect artifacts in batched data
    ///
    /// Processes data in sliding windows
    pub fn detect_batch(&self, data: &[Vec<f64>]) -> ArtifactResult<DetectionBatch> {
        let start = std::time::Instant::now();

        if data.is_empty() {
            return Err(ArtifactError::Input("Empty data".to_string()));
        }

        let n_samples = data[0].len();
        let step = ((1.0 - self.config.overlap) * self.config.window_size as f64) as usize;
        let step = step.max(1);

        let mut results = Vec::new();
        let mut offset = 0;

        while offset + self.config.window_size <= n_samples {
            // Extract window
            let window: Vec<Vec<f64>> = data
                .iter()
                .map(|ch| ch[offset..offset + self.config.window_size].to_vec())
                .collect();

            // Detect on window
            let mut result = self.detect(&window)?;

            // Update timing
            result.start_time = offset as f64 / self.config.sfreq;
            result.end_time = (offset + self.config.window_size) as f64 / self.config.sfreq;

            results.push(result);
            offset += step;
        }

        let processing_time_ms = start.elapsed().as_secs_f64() * 1000.0;

        Ok(DetectionBatch {
            n_windows: results.len(),
            results,
            processing_time_ms,
        })
    }

    /// Run ONNX inference
    fn run_inference(&self, session: &OnnxSession, data: &[Vec<f64>]) -> ArtifactResult<Vec<f64>> {
        let n_channels = data.len();
        let n_samples = data[0].len();

        // Flatten data to [batch=1, channels, time]
        let mut flat_data: Vec<f32> = Vec::with_capacity(n_channels * n_samples);
        for ch in data {
            for &sample in ch {
                flat_data.push(sample as f32);
            }
        }

        // Create input tensor
        let device = Device::Cpu;
        let input = Tensor::from_vec(flat_data, &[1, n_channels, n_samples], &device)
            .map_err(|e| ArtifactError::Tensor(e.to_string()))?;

        // Prepare inputs as HashMap
        let mut inputs: HashMap<String, &Tensor> = HashMap::new();
        inputs.insert(self.input_name.clone(), &input);

        // Run inference - need mutable session, but we only have immutable ref
        // This is a limitation - for now return error suggesting rule-based
        // In production, session should be wrapped in a mutex
        Err(ArtifactError::Inference(
            "ONNX inference requires mutable session. Using rule-based detection.".to_string(),
        ))
    }

    /// Rule-based fallback detection (when no model loaded)
    fn rule_based_detect(&self, data: &[Vec<f64>]) -> ArtifactResult<Vec<f64>> {
        let n_types = ArtifactType::count();
        let mut probabilities = vec![0.0; n_types];

        // Compute basic statistics
        let (mean_amplitude, max_amplitude, variance) = compute_statistics(data);

        // Eye blink detection (large frontal deflections)
        // Check first few channels (typically frontal in EEG)
        if data.len() > 2 {
            let frontal_amplitude = compute_channel_amplitude(&data[0..3.min(data.len())]);
            if frontal_amplitude > 100.0 {
                probabilities[ArtifactType::EyeBlink.index()] =
                    (frontal_amplitude / 200.0).min(1.0);
            }
        }

        // Muscle artifact detection (high-frequency content)
        let hf_power = compute_high_freq_power(data, self.config.sfreq);
        if hf_power > 0.3 {
            probabilities[ArtifactType::Muscle.index()] = hf_power.min(1.0);
        }

        // Line noise detection (50/60 Hz)
        let line_noise = detect_line_noise(data, self.config.sfreq);
        probabilities[ArtifactType::LineNoise.index()] = line_noise;

        // Movement artifact (slow drift)
        let drift = compute_drift(data);
        if drift > 0.3 {
            probabilities[ArtifactType::Movement.index()] = drift.min(1.0);
        }

        // Electrode pop (sudden jumps)
        let has_pop = detect_electrode_pop(data);
        probabilities[ArtifactType::ElectrodePop.index()] = if has_pop { 0.9 } else { 0.0 };

        // Channel noise (high variance channels)
        let noisy_channels = count_noisy_channels(data, variance);
        if noisy_channels > 0 {
            probabilities[ArtifactType::ChannelNoise.index()] =
                (noisy_channels as f64 / data.len() as f64).min(1.0);
        }

        Ok(probabilities)
    }
}

impl std::fmt::Debug for ArtifactDetector {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ArtifactDetector")
            .field("config", &self.config)
            .field("model_loaded", &self.model_loaded)
            .finish()
    }
}

// Helper functions

fn sigmoid(x: f64) -> f64 {
    1.0 / (1.0 + (-x).exp())
}

fn compute_statistics(data: &[Vec<f64>]) -> (f64, f64, f64) {
    let mut sum = 0.0;
    let mut max_val = f64::NEG_INFINITY;
    let mut count = 0;

    for ch in data {
        for &sample in ch {
            sum += sample.abs();
            max_val = max_val.max(sample.abs());
            count += 1;
        }
    }

    let mean = if count > 0 { sum / count as f64 } else { 0.0 };

    // Compute variance
    let mut var_sum = 0.0;
    for ch in data {
        for &sample in ch {
            var_sum += (sample.abs() - mean).powi(2);
        }
    }
    let variance = if count > 1 {
        var_sum / (count - 1) as f64
    } else {
        0.0
    };

    (mean, max_val, variance)
}

fn compute_channel_amplitude(channels: &[Vec<f64>]) -> f64 {
    let mut max_amp: f64 = 0.0;
    for ch in channels {
        let min = ch.iter().copied().fold(f64::INFINITY, f64::min);
        let max = ch.iter().copied().fold(f64::NEG_INFINITY, f64::max);
        max_amp = max_amp.max(max - min);
    }
    max_amp
}

fn compute_high_freq_power(data: &[Vec<f64>], sfreq: f64) -> f64 {
    // Simple high-pass approximation using differences
    let mut hf_power = 0.0;
    let mut total_power = 0.0;

    for ch in data {
        for i in 1..ch.len() {
            let diff = ch[i] - ch[i - 1];
            hf_power += diff * diff;
            total_power += ch[i] * ch[i];
        }
    }

    if total_power > 0.0 {
        // Scale by sampling frequency
        let scale = sfreq / 1000.0;
        (hf_power / total_power * scale).min(1.0)
    } else {
        0.0
    }
}

fn detect_line_noise(data: &[Vec<f64>], sfreq: f64) -> f64 {
    // Detect periodicity at 50/60 Hz
    // Simple autocorrelation check at expected period

    let period_50 = (sfreq / 50.0).round() as usize;
    let period_60 = (sfreq / 60.0).round() as usize;

    let mut max_corr: f64 = 0.0;

    for ch in data {
        if ch.len() > period_50.max(period_60) * 2 {
            // Check 50 Hz
            let corr_50 = autocorr_at_lag(ch, period_50);
            max_corr = max_corr.max(corr_50);

            // Check 60 Hz
            let corr_60 = autocorr_at_lag(ch, period_60);
            max_corr = max_corr.max(corr_60);
        }
    }

    max_corr.max(0.0).min(1.0)
}

fn autocorr_at_lag(signal: &[f64], lag: usize) -> f64 {
    if lag >= signal.len() {
        return 0.0;
    }

    let n = signal.len() - lag;
    let mean: f64 = signal.iter().sum::<f64>() / signal.len() as f64;

    let mut num = 0.0;
    let mut den = 0.0;

    for i in 0..n {
        let x = signal[i] - mean;
        let y = signal[i + lag] - mean;
        num += x * y;
        den += x * x;
    }

    if den > 0.0 { num / den } else { 0.0 }
}

fn compute_drift(data: &[Vec<f64>]) -> f64 {
    // Compute linear trend in each channel
    let mut max_drift: f64 = 0.0;

    for ch in data {
        if ch.len() < 2 {
            continue;
        }

        let n = ch.len() as f64;
        let sum_x: f64 = (0..ch.len()).map(|i| i as f64).sum();
        let sum_y: f64 = ch.iter().sum();
        let sum_xy: f64 = ch.iter().enumerate().map(|(i, &y)| i as f64 * y).sum();
        let sum_xx: f64 = (0..ch.len()).map(|i| (i * i) as f64).sum();

        let slope = (n * sum_xy - sum_x * sum_y) / (n * sum_xx - sum_x * sum_x);
        let drift = slope.abs() * n; // Total drift over signal

        // Normalize by signal range
        let range = ch.iter().copied().fold(f64::NEG_INFINITY, f64::max)
            - ch.iter().copied().fold(f64::INFINITY, f64::min);

        if range > 0.0 {
            max_drift = max_drift.max(drift / range);
        }
    }

    max_drift.min(1.0)
}

fn detect_electrode_pop(data: &[Vec<f64>]) -> bool {
    // Detect sudden large jumps
    let threshold = 5.0; // Standard deviations

    for ch in data {
        if ch.len() < 2 {
            continue;
        }

        // Compute differences
        let diffs: Vec<f64> = ch.windows(2).map(|w| (w[1] - w[0]).abs()).collect();

        if diffs.is_empty() {
            continue;
        }

        let mean: f64 = diffs.iter().sum::<f64>() / diffs.len() as f64;
        let std: f64 =
            (diffs.iter().map(|&d| (d - mean).powi(2)).sum::<f64>() / diffs.len() as f64).sqrt();

        // Check for outliers
        if std > 0.0 {
            for &diff in &diffs {
                if (diff - mean) / std > threshold {
                    return true;
                }
            }
        }
    }

    false
}

fn count_noisy_channels(data: &[Vec<f64>], global_variance: f64) -> usize {
    // Count channels with variance much higher than average
    let threshold = 3.0; // Times global variance

    let mut count = 0;
    for ch in data {
        let ch_var: f64 = {
            let mean: f64 = ch.iter().sum::<f64>() / ch.len() as f64;
            ch.iter().map(|&x| (x - mean).powi(2)).sum::<f64>() / ch.len() as f64
        };

        if global_variance > 0.0 && ch_var > threshold * global_variance {
            count += 1;
        }
    }

    count
}

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

    #[test]
    fn test_detector_config_default() {
        let config = DetectorConfig::default();
        assert_eq!(config.threshold, 0.5);
        assert_eq!(config.n_channels, 64);
    }

    #[test]
    fn test_detector_config_eeg() {
        let config = DetectorConfig::eeg(32, 500.0);
        assert_eq!(config.n_channels, 32);
        assert_eq!(config.sfreq, 500.0);
        assert_eq!(config.window_size, 500); // 1 second at 500 Hz
    }

    #[test]
    fn test_detector_creation() {
        let config = DetectorConfig::eeg(64, 1000.0);
        let detector = ArtifactDetector::new(config).unwrap();
        assert!(!detector.is_loaded()); // No model path provided
    }

    #[test]
    fn test_detection_result() {
        let labels = vec![
            ArtifactLabel::new(ArtifactType::EyeBlink, 0.8),
            ArtifactLabel::new(ArtifactType::Muscle, 0.3),
        ];

        let result = DetectionResult {
            labels,
            start_time: 0.0,
            end_time: 1.0,
            raw_probabilities: vec![0.8, 0.3],
        };

        assert!(result.has_artifact(0.5));
        assert_eq!(result.detected(0.5).len(), 1);
        assert_eq!(
            result.most_likely().unwrap().artifact_type,
            ArtifactType::EyeBlink
        );
    }

    #[test]
    fn test_rule_based_detection() {
        let config = DetectorConfig::eeg(4, 1000.0);
        let detector = ArtifactDetector::new(config).unwrap();

        // Create synthetic data with 4 channels, 1000 samples
        let data: Vec<Vec<f64>> = (0..4)
            .map(|_| (0..1000).map(|i| (i as f64 * 0.01).sin()).collect())
            .collect();

        let result = detector.detect(&data).unwrap();
        assert_eq!(result.labels.len(), ArtifactType::count());
    }

    #[test]
    fn test_batch_detection() {
        let config = DetectorConfig::eeg(4, 1000.0)
            .with_window_size(200)
            .with_overlap(0.5);
        let detector = ArtifactDetector::new(config).unwrap();

        // 2 seconds of data
        let data: Vec<Vec<f64>> = (0..4)
            .map(|_| (0..2000).map(|i| (i as f64 * 0.01).sin()).collect())
            .collect();

        let batch = detector.detect_batch(&data).unwrap();
        assert!(batch.n_windows > 1);
        assert!(batch.processing_time_ms >= 0.0);
    }

    #[test]
    fn test_sigmoid() {
        assert!((sigmoid(0.0) - 0.5).abs() < 1e-10);
        assert!(sigmoid(10.0) > 0.99);
        assert!(sigmoid(-10.0) < 0.01);
    }

    #[test]
    fn test_electrode_pop_detection() {
        // Create data with a sudden jump
        let mut ch = vec![0.0; 100];
        ch[50] = 1000.0; // Large spike

        let data = vec![ch];
        assert!(detect_electrode_pop(&data));

        // Normal data
        let normal: Vec<Vec<f64>> = vec![(0..100).map(|i| (i as f64).sin()).collect()];
        assert!(!detect_electrode_pop(&normal));
    }

    #[test]
    fn test_compute_statistics() {
        let data = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]];
        let (mean, max, variance) = compute_statistics(&data);

        assert!((mean - 3.5).abs() < 0.1);
        assert_eq!(max, 6.0);
        assert!(variance > 0.0);
    }

    #[test]
    fn test_detection_batch_to_regions() {
        let results = vec![
            DetectionResult {
                labels: vec![ArtifactLabel::new(ArtifactType::EyeBlink, 0.8)],
                start_time: 0.0,
                end_time: 1.0,
                raw_probabilities: vec![0.8],
            },
            DetectionResult {
                labels: vec![ArtifactLabel::new(ArtifactType::EyeBlink, 0.7)],
                start_time: 0.5,
                end_time: 1.5,
                raw_probabilities: vec![0.7],
            },
        ];

        let batch = DetectionBatch {
            results,
            processing_time_ms: 10.0,
            n_windows: 2,
        };

        let regions = batch.to_regions(0.5);
        // Should merge overlapping regions
        assert_eq!(regions.len(), 1);
        assert_eq!(regions[0].start_time, 0.0);
        assert_eq!(regions[0].end_time, 1.5);
    }
}