//! Protocol and study organization system.

use crate::{NeuroResult, Recording};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};

/// A protocol (study/experiment) containing multiple subjects
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Protocol {
    /// Protocol name
    pub name: String,
    /// Root directory for this protocol
    pub path: PathBuf,
    /// Protocol description
    pub description: String,
    /// Creation date
    pub created: DateTime<Utc>,
    /// Last modified date
    pub modified: DateTime<Utc>,
    /// Subjects in this protocol
    subjects: HashMap<String, Subject>,
    /// Group-level analyses
    group_analyses: Vec<GroupAnalysis>,
    /// Protocol settings
    pub settings: ProtocolSettings,
}

impl Protocol {
    /// Create a new protocol
    #[must_use]
    pub fn new(name: impl Into<String>, path: impl Into<PathBuf>) -> Self {
        let now = Utc::now();
        Self {
            name: name.into(),
            path: path.into(),
            description: String::new(),
            created: now,
            modified: now,
            subjects: HashMap::new(),
            group_analyses: Vec::new(),
            settings: ProtocolSettings::default(),
        }
    }

    /// Add a subject to the protocol
    pub fn add_subject(&mut self, subject: Subject) {
        self.subjects.insert(subject.id.clone(), subject);
        self.modified = Utc::now();
    }

    /// Get a subject by ID
    #[must_use]
    pub fn get_subject(&self, id: &str) -> Option<&Subject> {
        self.subjects.get(id)
    }

    /// Get mutable reference to a subject
    pub fn get_subject_mut(&mut self, id: &str) -> Option<&mut Subject> {
        self.modified = Utc::now();
        self.subjects.get_mut(id)
    }

    /// Remove a subject
    pub fn remove_subject(&mut self, id: &str) -> Option<Subject> {
        self.modified = Utc::now();
        self.subjects.remove(id)
    }

    /// List all subject IDs
    #[must_use]
    pub fn subject_ids(&self) -> Vec<&str> {
        self.subjects.keys().map(String::as_str).collect()
    }

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

    /// Iterate over subjects
    pub fn subjects(&self) -> impl Iterator<Item = &Subject> {
        self.subjects.values()
    }

    /// Add a group analysis
    pub fn add_group_analysis(&mut self, analysis: GroupAnalysis) {
        self.group_analyses.push(analysis);
        self.modified = Utc::now();
    }

    /// Get group analyses
    #[must_use]
    pub fn group_analyses(&self) -> &[GroupAnalysis] {
        &self.group_analyses
    }

    /// Save protocol to disk
    pub fn save(&self) -> NeuroResult<()> {
        let protocol_file = self.path.join("protocol.json");
        let json = serde_json::to_string_pretty(self).map_err(|e| {
            crate::NeuroError::Io(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                e.to_string(),
            ))
        })?;
        std::fs::write(&protocol_file, json)?;
        Ok(())
    }

    /// Load protocol from disk
    pub fn load(path: impl AsRef<Path>) -> NeuroResult<Self> {
        let path = path.as_ref();
        let protocol_file = path.join("protocol.json");
        let json = std::fs::read_to_string(&protocol_file)?;
        let protocol: Self = serde_json::from_str(&json)
            .map_err(|e| crate::NeuroError::InvalidFormat(e.to_string()))?;
        Ok(protocol)
    }
}

/// Protocol-wide settings
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProtocolSettings {
    /// Default sampling frequency for new recordings
    pub default_sfreq: f64,
    /// Default line frequency (50 or 60 Hz)
    pub line_freq: f64,
    /// Use GPU acceleration if available
    pub use_gpu: bool,
    /// Number of parallel workers
    pub n_jobs: usize,
}

impl Default for ProtocolSettings {
    fn default() -> Self {
        Self {
            default_sfreq: 1000.0,
            line_freq: 60.0,
            use_gpu: true,
            n_jobs: 4,
        }
    }
}

/// A subject (participant) in a study
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Subject {
    /// Subject ID
    pub id: String,
    /// Subject name (optional, for display)
    pub name: Option<String>,
    /// Subject comments/notes
    pub comments: String,
    /// Anatomy data
    pub anatomy: Option<Anatomy>,
    /// Raw recordings
    recordings: Vec<Recording>,
    /// Conditions/sessions
    conditions: HashMap<String, Condition>,
}

impl Subject {
    /// Create a new subject
    #[must_use]
    pub fn new(id: impl Into<String>) -> Self {
        Self {
            id: id.into(),
            name: None,
            comments: String::new(),
            anatomy: None,
            recordings: Vec::new(),
            conditions: HashMap::new(),
        }
    }

    /// Add a recording
    pub fn add_recording(&mut self, recording: Recording) {
        self.recordings.push(recording);
    }

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

    /// Add a condition
    pub fn add_condition(&mut self, condition: Condition) {
        self.conditions.insert(condition.name.clone(), condition);
    }

    /// Get a condition by name
    #[must_use]
    pub fn get_condition(&self, name: &str) -> Option<&Condition> {
        self.conditions.get(name)
    }

    /// List condition names
    #[must_use]
    pub fn condition_names(&self) -> Vec<&str> {
        self.conditions.keys().map(String::as_str).collect()
    }

    /// Set anatomy
    pub fn set_anatomy(&mut self, anatomy: Anatomy) {
        self.anatomy = Some(anatomy);
    }
}

/// Anatomical data for a subject
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Anatomy {
    /// Path to MRI volume
    pub mri_path: Option<PathBuf>,
    /// Path to cortical surface (left hemisphere)
    pub surface_lh: Option<PathBuf>,
    /// Path to cortical surface (right hemisphere)
    pub surface_rh: Option<PathBuf>,
    /// Path to head surface (for BEM/visualization)
    pub surface_head: Option<PathBuf>,
    /// Fiducial points (NAS, LPA, RPA) in MRI coordinates
    pub fiducials: HashMap<String, [f64; 3]>,
    /// Coordinate transform from MRI to head coordinates
    pub mri_to_head: Option<[[f64; 4]; 4]>,
    /// Whether anatomy is from template (vs individual MRI)
    pub is_template: bool,
    /// Template name if using template
    pub template_name: Option<String>,
}

impl Anatomy {
    /// Create a new anatomy
    #[must_use]
    pub fn new() -> Self {
        Self {
            mri_path: None,
            surface_lh: None,
            surface_rh: None,
            surface_head: None,
            fiducials: HashMap::new(),
            mri_to_head: None,
            is_template: false,
            template_name: None,
        }
    }

    /// Create anatomy from template
    #[must_use]
    pub fn from_template(name: impl Into<String>) -> Self {
        Self {
            is_template: true,
            template_name: Some(name.into()),
            ..Self::new()
        }
    }

    /// Set fiducial point
    pub fn set_fiducial(&mut self, name: impl Into<String>, pos: [f64; 3]) {
        self.fiducials.insert(name.into(), pos);
    }
}

impl Default for Anatomy {
    fn default() -> Self {
        Self::new()
    }
}

/// A condition/task within a subject's data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Condition {
    /// Condition name
    pub name: String,
    /// Description
    pub description: String,
    /// Associated data files
    pub data_files: Vec<PathBuf>,
    /// Epochs file (if computed)
    pub epochs_file: Option<PathBuf>,
    /// Evoked/average file (if computed)
    pub evoked_file: Option<PathBuf>,
    /// Source estimate files
    pub source_files: Vec<PathBuf>,
}

impl Condition {
    /// Create a new condition
    #[must_use]
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            description: String::new(),
            data_files: Vec::new(),
            epochs_file: None,
            evoked_file: None,
            source_files: Vec::new(),
        }
    }

    /// Add a data file
    pub fn add_data_file(&mut self, path: impl Into<PathBuf>) {
        self.data_files.push(path.into());
    }
}

/// Group-level analysis
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GroupAnalysis {
    /// Analysis name
    pub name: String,
    /// Description
    pub description: String,
    /// Subjects included
    pub subjects: Vec<String>,
    /// Conditions compared
    pub conditions: Vec<String>,
    /// Output files
    pub outputs: Vec<PathBuf>,
    /// Creation date
    pub created: DateTime<Utc>,
}

impl GroupAnalysis {
    /// Create a new group analysis
    #[must_use]
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            description: String::new(),
            subjects: Vec::new(),
            conditions: Vec::new(),
            outputs: Vec::new(),
            created: Utc::now(),
        }
    }

    /// Add subjects
    pub fn add_subjects<S: Into<String>>(&mut self, subjects: impl IntoIterator<Item = S>) {
        self.subjects.extend(subjects.into_iter().map(Into::into));
    }

    /// Add conditions
    pub fn add_conditions<S: Into<String>>(&mut self, conditions: impl IntoIterator<Item = S>) {
        self.conditions
            .extend(conditions.into_iter().map(Into::into));
    }
}

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

    #[test]
    fn test_protocol_creation() {
        let protocol = Protocol::new("test_study", "/tmp/test");
        assert_eq!(protocol.name, "test_study");
        assert_eq!(protocol.n_subjects(), 0);
    }

    #[test]
    fn test_add_subject() {
        let mut protocol = Protocol::new("test", "/tmp/test");
        protocol.add_subject(Subject::new("sub-01"));
        protocol.add_subject(Subject::new("sub-02"));

        assert_eq!(protocol.n_subjects(), 2);
        assert!(protocol.get_subject("sub-01").is_some());
        assert!(protocol.get_subject("sub-03").is_none());
    }

    #[test]
    fn test_protocol_save_load() {
        let temp_dir = TempDir::new().unwrap();
        let path = temp_dir.path();

        let mut protocol = Protocol::new("test", path);
        protocol.add_subject(Subject::new("sub-01"));
        protocol.save().unwrap();

        let loaded = Protocol::load(path).unwrap();
        assert_eq!(loaded.name, "test");
        assert_eq!(loaded.n_subjects(), 1);
    }
}