edf-rs 0.2.0

use std::collections::HashMap;
use std::ops::{Bound, RangeBounds};

use crate::EDFSpecifications;
use crate::error::edf_error::EDFError;
use crate::headers::annotation_list::AnnotationList;
use crate::headers::edf_header::EDFHeader;
use crate::headers::signal_header::SignalHeader;
use crate::save::{SaveInstruction, SaveValue};

#[derive(Debug, Default, Clone, PartialEq)]
struct RecordLayout {
    signal_map: HashMap<usize, SignalType>,
    annotation_samples_count: Vec<usize>,
}

#[derive(Debug, Clone, PartialEq)]
enum SignalType {
    Samples(usize),
    Annotation(usize),
}

#[derive(Debug, Clone, PartialEq)]
pub enum Samples {
    Values16Bit(Vec<i16>),
    Values24Bit(Vec<i32>),
}

impl Samples {
    pub fn len(&self) -> usize {
        match self {
            Self::Values16Bit(samples) => samples.len(),
            Self::Values24Bit(samples) => samples.len(),
        }
    }

    pub fn is_empty(&self) -> bool {
        match self {
            Self::Values16Bit(samples) => samples.is_empty(),
            Self::Values24Bit(samples) => samples.is_empty(),
        }
    }

    pub fn extend(&mut self, samples: &Samples) -> Result<(), EDFError> {
        match (self, samples) {
            (Self::Values16Bit(self_samples), Self::Values16Bit(samples)) => self_samples.extend(samples),
            (Self::Values24Bit(self_samples), Self::Values24Bit(samples)) => self_samples.extend(samples),
            _ => return Err(EDFError::MismatchedSampleBits)
        }

        Ok(())
    }

    pub fn range<R: RangeBounds<usize>>(&self, range: R) -> Samples {
        let start = match range.start_bound() {
            Bound::Included(&n) => n,
            Bound::Excluded(&n) => n + 1,
            Bound::Unbounded => 0,
        };
        let end = match range.end_bound() {
            Bound::Included(&n) => Some(n + 1),
            Bound::Excluded(&n) => Some(n),
            Bound::Unbounded => None,
        };

        match self {
            Self::Values16Bit(samples) => Self::Values16Bit(samples[start..end.unwrap_or(samples.len())].to_vec()),
            Self::Values24Bit(samples) => Self::Values24Bit(samples[start..end.unwrap_or(samples.len())].to_vec()),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct Record {
    layout: RecordLayout,
    specification: EDFSpecifications,
    pub(crate) default_offset: f64,
    pub raw_signal_samples: Vec<Samples>,
    pub annotations: Vec<Vec<AnnotationList>>,
}

impl Record {
    pub fn new(signal_headers: &Vec<SignalHeader>, specification: &EDFSpecifications) -> Self {
        let mut raw_signal_samples = Vec::new();
        let mut annotations = Vec::new();
        let mut annotation_samples_count = Vec::new();
        let mut signal_map = HashMap::new();
        for (i, signal) in signal_headers.iter().enumerate() {
            if signal.is_annotation() {
                signal_map.insert(i, SignalType::Annotation(annotations.len()));
                annotation_samples_count.push(signal.samples_count);
                annotations.push(Vec::new());
            } else {
                signal_map.insert(i, SignalType::Samples(raw_signal_samples.len()));
                raw_signal_samples.push(match specification {
                    EDFSpecifications::EDF | EDFSpecifications::EDFPlus => Samples::Values16Bit(vec![0; signal.samples_count]),
                    EDFSpecifications::BDF | EDFSpecifications::BDFPlus => Samples::Values24Bit(vec![0; signal.samples_count])
                });
            }
        }

        Self {
            layout: RecordLayout {
                signal_map,
                annotation_samples_count,
            },
            specification: specification.clone(),
            default_offset: 0.0,
            raw_signal_samples,
            annotations,
        }
    }

    pub fn new_samples(&self, signal_index: usize, capacity: usize) -> Result<Samples, EDFError> {
        let Some(SignalType::Samples(idx)) = self.layout.signal_map.get(&signal_index) else {
            return Err(EDFError::ItemNotFound);
        };

        let Some(samples) = self.raw_signal_samples.get(*idx) else {
            return Err(EDFError::ItemNotFound);
        };

        Ok(match samples {
            Samples::Values16Bit(_) => Samples::Values16Bit(Vec::with_capacity(capacity)),
            Samples::Values24Bit(_) => Samples::Values24Bit(Vec::with_capacity(capacity)),
        })
    }

    pub fn patch_record(&mut self, instructions: &Vec<SaveInstruction>) -> Result<(), EDFError> {
        if instructions.is_empty() {
            return Ok(());
        }

        // Process each instruction
        let mut signal_idx = instructions[0].index();
        let mut instruction_idx = 0;
        loop {
            let Some(tr) = instructions.get(instruction_idx) else {
                break;
            };

            match tr {
                SaveInstruction::Remove(idx) if *idx == signal_idx => {
                    instruction_idx += 1;
                    self.remove_signal(*idx)?;
                }
                SaveInstruction::Insert(idx, SaveValue::Signal(value)) if *idx == signal_idx => {
                    instruction_idx += 1;
                    if value.is_annotation() {
                        self.insert_annotation(*idx, value.samples_count)?;
                    } else {
                        self.insert_signal_samples(*idx, value.samples_count)?;
                    }
                }
                SaveInstruction::Update(idx, SaveValue::Signal(value)) if *idx == signal_idx => {
                    signal_idx += 1;
                    instruction_idx += 1;
                    self.update_samples_count(*idx, value.samples_count)?;
                }
                _ => {
                    signal_idx += 1;
                }
            }
        }

        Ok(())
    }

    pub fn insert_signal_samples(
        &mut self,
        signal_index: usize,
        samples_count: usize,
    ) -> Result<(), EDFError> {
        // Get count of signal indices which are samples and lower than the target index
        let insert_idx = (0..signal_index)
            .filter(|i| {
                self.layout
                    .signal_map
                    .get(&i)
                    .is_some_and(|s| matches!(s, SignalType::Samples(idx) if *idx < signal_index))
            })
            .count();

        // Increase the global signal index pointers in the signal map as well as the sample signal index pointers
        self.apply_index_change_samples(signal_index, insert_idx, 1);
        self.layout
            .signal_map
            .insert(signal_index, SignalType::Samples(insert_idx));

        // Insert the new annotation signal values
        self.raw_signal_samples.insert(insert_idx, match self.specification {
            EDFSpecifications::EDF | EDFSpecifications::EDFPlus => Samples::Values16Bit(vec![0; samples_count]),
            EDFSpecifications::BDF | EDFSpecifications::BDFPlus => Samples::Values24Bit(vec![0; samples_count])
        });

        Ok(())
    }

    pub fn insert_annotation(
        &mut self,
        signal_index: usize,
        samples_count: usize,
    ) -> Result<(), EDFError> {
        // Get count of signal indices which are annotations and lower than the target index
        let insert_idx = (0..signal_index)
            .filter(|i| {
                self.layout.signal_map.get(&i).is_some_and(
                    |s| matches!(s, SignalType::Annotation(idx) if *idx < signal_index),
                )
            })
            .count();

        // Increase the global signal index pointers in the signal map as well as the annotation signal index pointers
        self.apply_index_change_annotation(signal_index, insert_idx, 1);
        self.layout
            .signal_map
            .insert(signal_index, SignalType::Annotation(insert_idx));

        // Insert the new annotation signal values
        self.layout
            .annotation_samples_count
            .insert(insert_idx, samples_count);
        self.annotations.insert(insert_idx, Vec::new());

        Ok(())
    }

    pub fn remove_signal(&mut self, signal_index: usize) -> Result<(), EDFError> {
        match self.layout.signal_map.remove(&signal_index) {
            Some(SignalType::Samples(idx)) => {
                self.raw_signal_samples.remove(idx);
                self.apply_index_change_samples(signal_index, idx, -1);
            }
            Some(SignalType::Annotation(idx)) => {
                self.layout.annotation_samples_count.remove(idx);
                self.annotations.remove(idx);
                self.apply_index_change_annotation(signal_index, idx, -1);
            }
            _ => return Err(EDFError::ItemNotFound),
        }

        Ok(())
    }

    pub fn update_samples_count(
        &mut self,
        signal_index: usize,
        samples_count: usize,
    ) -> Result<(), EDFError> {
        match self.layout.signal_map.get(&signal_index) {
            Some(SignalType::Samples(idx)) => {
                match self.raw_signal_samples.get_mut(*idx) {
                    Some(Samples::Values16Bit(values)) => values.resize(samples_count, 0),
                    Some(Samples::Values24Bit(values)) => values.resize(samples_count, 0),
                    _ => return Err(EDFError::ItemNotFound)
                }
            }
            Some(SignalType::Annotation(idx)) => {
                if let Some(count) = self.layout.annotation_samples_count.get_mut(*idx) {
                    *count = samples_count;
                } else {
                    return Err(EDFError::ItemNotFound);
                }
            }
            _ => return Err(EDFError::ItemNotFound),
        }

        Ok(())
    }

    pub fn set_annotation(
        &mut self,
        signal_index: usize,
        annotations: Vec<AnnotationList>,
    ) -> Result<(), EDFError> {
        let Some(SignalType::Annotation(idx)) = self.layout.signal_map.get(&signal_index) else {
            return Err(EDFError::ItemNotFound);
        };

        let Some(old_annotations) = self.annotations.get_mut(*idx) else {
            return Err(EDFError::ItemNotFound);
        };

        *old_annotations = annotations;

        Ok(())
    }

    pub fn set_samples(&mut self, signal_index: usize, samples: Samples) -> Result<(), EDFError> {
        let Some(SignalType::Samples(idx)) = self.layout.signal_map.get(&signal_index) else {
            return Err(EDFError::ItemNotFound);
        };

        let Some(old_samples) = self.raw_signal_samples.get_mut(*idx) else {
            return Err(EDFError::ItemNotFound);
        };

        if old_samples.len() != samples.len() {
            return Err(EDFError::InvalidSamplesCount);
        }

        *old_samples = samples;

        Ok(())
    }

    pub fn get_digital_samples(&self, signal: &SignalHeader) -> Vec<Vec<i32>> {
        self.raw_signal_samples.iter().map(|signals| {
            match signals {
                Samples::Values16Bit(samples) => signal.to_digital_samples(samples),
                Samples::Values24Bit(samples) => signal.to_digital_samples(samples),
            }
        }).collect()
    }

    pub fn get_physical_samples(&self, signal: &SignalHeader) -> Vec<Vec<f64>> {
        let range = (signal.physical_maximum - signal.physical_minimum) / (signal.digital_maximum - signal.digital_minimum) as f64;
        let offset = signal.physical_maximum / range - signal.digital_maximum as f64;

        self.raw_signal_samples.iter().map(|signals| {
            match signals {
                Samples::Values16Bit(samples) => signal.to_physical_samples(samples, range, offset),
                Samples::Values24Bit(samples) => signal.to_physical_samples(samples, range, offset),
            }
        }).collect()
    }

    fn apply_index_change_annotation(
        &mut self,
        signal_index: usize,
        target_index: usize,
        direction: i8,
    ) {
        let mut new = HashMap::new();
        for (k, v) in self.layout.signal_map.drain() {
            let new_global_index =
                (k as i64 + direction as i64 * (k >= signal_index) as i64) as usize;
            let value = if let SignalType::Annotation(idx) = v
                && idx >= target_index
            {
                SignalType::Annotation((idx as i64 + direction as i64) as usize)
            } else {
                v
            };
            new.insert(new_global_index, value);
        }
        self.layout.signal_map = new;
    }

    fn apply_index_change_samples(
        &mut self,
        signal_index: usize,
        target_index: usize,
        direction: i8,
    ) {
        let mut new = HashMap::new();
        for (k, v) in self.layout.signal_map.drain() {
            let new_global_index =
                (k as i64 + direction as i64 * (k >= signal_index) as i64) as usize;
            let value = if let SignalType::Samples(idx) = v
                && idx >= target_index
            {
                SignalType::Samples((idx as i64 + direction as i64) as usize)
            } else {
                v
            };
            new.insert(new_global_index, value);
        }
        self.layout.signal_map = new;
    }

    /// Returns the onset of the current record relative to the start of the recording of the EDF+ file.
    /// This only returns useful information for EDF+ files. Regular EDF files will always return the
    /// index of the data-record multiplied by the data-record duration as records are missing the time keeping context.
    /// If there were to be multiple signals labeled `EDF Annotations`, the first one will be used to check for the
    /// Time-keeping-list entry
    pub fn get_start_offset(&self) -> f64 {
        self.annotations
            .first()
            .map(|tals| tals.iter().find(|a| a.is_time_keeping()).map(|a| a.onset))
            .flatten()
            .unwrap_or(self.default_offset)
    }

    pub fn serialize(&self) -> Result<Vec<u8>, EDFError> {
        let sample_bytes = match self.specification {
            EDFSpecifications::EDF | EDFSpecifications::EDFPlus => 2,
            EDFSpecifications::BDF | EDFSpecifications::BDFPlus => 3
        };
        let mut result_buffer = vec![];

        for signal_idx in 0..self.layout.signal_map.len() {
            match self.layout.signal_map.get(&signal_idx) {
                Some(SignalType::Annotation(idx)) => {
                    if let Some(annotation) = self.annotations.get(*idx)
                        && let Some(sample_count) = self.layout.annotation_samples_count.get(*idx)
                    {
                        let tals = annotation
                            .iter()
                            .map(|a| a.serialize())
                            .collect::<Vec<_>>()
                            .join("");
                        let mut tal_bytes = tals.as_bytes().to_vec();
                        tal_bytes.extend(vec![0; sample_bytes * sample_count - tal_bytes.len()]);
                        result_buffer.extend(tal_bytes);
                    }
                }
                Some(SignalType::Samples(idx)) => {
                    if let Some(signal) = self.raw_signal_samples.get(*idx) {
                        result_buffer.extend(match signal {
                            Samples::Values16Bit(samples) => samples.into_iter().map(|s| s.to_le_bytes()).flatten().collect::<Vec<_>>(),
                            Samples::Values24Bit(samples) => samples.into_iter().map(|s| i24_to_le_bytes(s)).flatten().collect::<Vec<_>>()
                        });
                    }
                }
                _ => {
                    panic!("Invalid record signal mapping index. This should not be possible")
                }
            }
        }

        Ok(result_buffer)
    }

    pub fn matches_signals(&self, signal_headers: &Vec<SignalHeader>) -> bool {
        // Validate the signal count of the record matches the provided signal header count
        let actual_count = self.annotations.len() + self.raw_signal_samples.len();
        if actual_count != signal_headers.len()
            || actual_count != self.layout.signal_map.len()
            || actual_count
                != self
                    .layout
                    .signal_map
                    .keys()
                    .max()
                    .map(|k| *k + 1)
                    .unwrap_or(0)
        {
            return false;
        }

        // Validate the sample count of every signal in the record matches the provided signal header
        for i in 0..actual_count {
            match self.layout.signal_map.get(&i) {
                Some(SignalType::Samples(idx)) => {
                    if !self
                        .raw_signal_samples
                        .get(*idx)
                        .is_some_and(|s| s.len() == signal_headers[i].samples_count)
                    {
                        return false;
                    }
                }
                Some(SignalType::Annotation(idx)) => {
                    if !self
                        .layout
                        .annotation_samples_count
                        .get(*idx)
                        .is_some_and(|c| *c == signal_headers[i].samples_count)
                    {
                        return false;
                    }
                }
                _ => return false,
            }
        }

        true
    }
}

fn i24_to_le_bytes(value: &i32) -> [u8; 3] {
    let bytes = value.to_le_bytes();
    [bytes[0], bytes[1], bytes[2]]
}

#[derive(Debug, Clone, PartialEq)]
pub struct RelativeRecordData {
    pub offset: f64,
    pub raw_signal_samples: Samples,
}

impl RelativeRecordData {
    pub fn new(offset: f64, specification: &EDFSpecifications) -> Self {
        Self {
            offset,
            raw_signal_samples: match specification {
                EDFSpecifications::EDF | EDFSpecifications::EDFPlus => Samples::Values16Bit(Vec::new()),
                EDFSpecifications::BDF | EDFSpecifications::BDFPlus => Samples::Values24Bit(Vec::new()),
            },
        }
    }

    pub fn get_digital_samples(&self, signal: &SignalHeader) -> Vec<i32> {
        match &self.raw_signal_samples {
            Samples::Values16Bit(samples) => signal.to_digital_samples(samples),
            Samples::Values24Bit(samples) => signal.to_digital_samples(samples),
        }
    }

    pub fn get_physical_samples(&self, signal: &SignalHeader) -> Vec<f64> {
        let range = (signal.physical_maximum - signal.physical_minimum) / (signal.digital_maximum - signal.digital_minimum) as f64;
        let offset = signal.physical_maximum / range - signal.digital_maximum as f64;

        match &self.raw_signal_samples {
            Samples::Values16Bit(samples) => signal.to_physical_samples(samples, range, offset),
            Samples::Values24Bit(samples) => signal.to_physical_samples(samples, range, offset),
        }
    }
}

#[derive(Debug, Default, Clone, PartialEq)]
pub struct SpanningRecord {
    pub raw_signal_samples: Vec<Vec<RelativeRecordData>>,
    pub annotations: Vec<Vec<AnnotationList>>,
}

impl SpanningRecord {
    pub fn new(header: &EDFHeader) -> Self {
        let signal_count = header.signals.iter().filter(|s| !s.is_annotation()).count();
        Self {
            raw_signal_samples: vec![Vec::new(); signal_count],
            annotations: Vec::new(),
        }
    }

    pub fn is_spanning_wait(&self) -> bool {
        self.raw_signal_samples
            .iter()
            .all(|sp| sp.last().is_some_and(|data| data.raw_signal_samples.is_empty()))
    }

    pub fn remove_last_spanning_wait(&mut self) -> bool {
        if self.is_spanning_wait() {
            for signal in &mut self.raw_signal_samples {
                signal.remove(signal.len() - 1);
            }

            return true;
        }
        return false;
    }

    pub fn insert_spanning_wait(&mut self, offset: f64, specification: &EDFSpecifications) {
        self.remove_last_spanning_wait();

        // Check if the last spanning entry is the same offset
        if self
            .raw_signal_samples
            .first()
            .map(|s| s.last())
            .flatten()
            .is_some_and(|s| s.offset == offset)
        {
            return;
        }

        // In case it is a new offset, insert the new spanning values
        for signal in &mut self.raw_signal_samples {
            signal.push(RelativeRecordData::new(offset, specification));
        }
    }

    pub fn finish(&mut self) {
        self.remove_last_spanning_wait();
        // TODO: This should probably also go through all annotations and remove all
        // Time-keeping entries.
    }

    pub fn extend_samples(&mut self, signal_index: usize, samples: &Samples) -> Result<(), EDFError> {
        if let Some(signal) = self.raw_signal_samples.get_mut(signal_index) {
            if let Some(data) = signal.last_mut() {
                data.raw_signal_samples.extend(samples)?;
                return Ok(())
            }
        }

        Err(EDFError::IndexOutOfBounds)
    }
}