imzml 0.1.3

use std::{
    collections::VecDeque,
    fmt::Display,
    fs::File,
    io::{BufRead, BufReader, Read, Seek, Write},
    path::{Path, PathBuf},
    sync::{Arc, Mutex},
};

use hashbrown::HashMap;
use quick_xml::{
    events::{BytesStart, Event},
    Decoder, Error, Reader,
};
use serde::{Deserialize, Serialize};

use crate::{
    error::{FatalParseError, ParseError},
    obo::{Ontology, MZML_ONTOLOGY},
    IndexedMzML, MzML, Representation, SpectrumAccess, SpectrumAccessIterator, BUFFER_SIZE,
};

use self::{
    attributes::{AttributeDefinition, AttributeType, AttributeValue},
    data_processing::{DataProcessing, DataProcessingRef},
    filedescription::{SourceFile, SourceFileRef},
    instrument::{InstrumentConfiguration, InstrumentConfigurationRef},
    referenceableparamgroup::{ReferenceableParamGroup, ReferenceableParamGroupRef},
    sample::Sample,
    scan_settings::ScanSettingsList,
    software::{Software, SoftwareRef},
    writer::Writer,
};

pub(crate) mod attributes;
pub(crate) mod binarydataarray;
pub(crate) mod cvlist;
/// Metadata in the form of cvParam entries (linked to ontology)
pub mod cvparam;
/// Information on data processing
pub mod data_processing;
/// Describes the file (e.g. style of data, source raw files used to create it, ...)
pub mod filedescription;
/// Provide instrument configuration information
pub mod instrument;
/// Handling mzML tag
pub mod mzml;
//mod parser;
/// Sets of cvParams which can be referenced to reduce overhead
pub mod referenceableparamgroup;
/// Captures information about the experiment (run) and data (e.g. list of spectra)
pub mod run;
pub(crate) mod sample;
pub(crate) mod scan;
/// Describes the settings used to acquire the data.
pub mod scan_settings;
pub(crate) mod software;
pub(crate) mod spectrum;

/// Writing .imzML files
pub mod writer;
//pub use writer::*;

/// Trait describing any tag present in an .mzML or .imzML file
pub trait MzMLTag {
    /// Returns the tag being described
    fn tag() -> Tag;

    /// Parse the start tag (capture any attributes, noting any missing or extra as errors).
    fn parse_start_tag<B: BufRead>(
        parser: &mut MzMLReader<B>,
        start_event: &BytesStart,
    ) -> Result<Option<Self>, FatalParseError>
    where
        Self: std::marker::Sized;

    /// Parse the child tags (if necessary). This is only called if there are expected to be child tags.
    fn parse_xml<B: BufRead>(
        &mut self,
        parser: &mut MzMLReader<B>,
        buffer: &mut Vec<u8>,
    ) -> Result<(), FatalParseError>;

    /// Write the tag in XML format.
    fn write_xml<W: Write>(&self, writer: &mut Writer<W>) -> Result<(), quick_xml::Error>;
}

/// DataReader enables reading the binary data stored in an mzML, as well as providing access to the metadata stored in the XML part.
pub struct DataReader<D: Read + Seek> {
    mzml: MzML,
    data_reader: Arc<Mutex<D>>,

    errors: VecDeque<ParseError>,
}

impl<D: Read + Seek> DataReader<D> {
    /// Create a new parser from the supplied reader, using the default mzML ontology
    pub fn new(mzml: MzML, data: D) -> Self {
        DataReader {
            mzml,
            data_reader: Arc::new(Mutex::new(data)),
            errors: VecDeque::new(),
        }
    }

    /// Create a new parser from the supplied reader and ontology
    pub fn with_ontology<B: BufRead>(
        header: B,
        data: D,
        ontology: Ontology,
    ) -> Result<Self, FatalParseError> {
        let reader = MzMLReader::with_ontology(header, ontology)?;

        let mzml = reader.mzml.ok_or_else(|| {
            FatalParseError::UnexpectedError("No MzML generated when parsing".to_string())
        })?;

        Ok(DataReader {
            mzml,
            data_reader: Arc::new(Mutex::new(data)),
            errors: reader.errors,
        })
    }

    /// Returns reference to the `MzML` (metadata) information
    pub fn mzml(&self) -> &MzML {
        &self.mzml
    }

    /// Creates and returns an iterator over all spectra in the image.
    pub fn spectra(&self) -> SpectrumAccessIterator<D> {
        //impl Iterator<Item = DiskSpectrumAccess<D>> + '_ {
        SpectrumAccessIterator::new(
            self.data_reader.clone(),
            self.mzml().spectrum_list().unwrap(),
        )
    }

    /// Returns `SpectrumAccess` to the spectrum at the specified index (or None) if none present
    pub fn spectrum(&self, index: usize) -> Option<SpectrumAccess<D>> {
        if let Some(spectrum) = self.mzml().spectrum(index) {
            return Some(SpectrumAccess::new(
                self.data_reader.clone(),
                spectrum.clone(),
            ));
        }

        None
    }

    /// Return errors found during parsing
    pub fn errors(&self) -> &VecDeque<ParseError> {
        &self.errors
    }

    /// Returns mutable reference to error list
    pub fn errors_mut(&mut self) -> &mut VecDeque<ParseError> {
        &mut self.errors
    }
}

impl DataReader<BufReader<File>> {
    /// Create a DataReader from a path. This assumes that the ibd file has the exact same name and the extension '.ibd'
    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Self, FatalParseError> {
        Self::from_path_with_ontology(path, MZML_ONTOLOGY.clone())
    }

    /// Create an ImzMLParser from a path, using the supplied ontology to map CV params.
    /// This assumes that the ibd file has the exact same name and the extension '.ibd'
    pub fn from_path_with_ontology<P: AsRef<Path>>(
        path: P,
        ontology: Ontology,
    ) -> Result<Self, FatalParseError> {
        let path = PathBuf::from(path.as_ref());
        let file = File::open(&path)?;
        let data_file = File::open(&path)?;

        Self::with_ontology(BufReader::new(file), BufReader::new(data_file), ontology)
    }
}

/// MzMLParser parses an mzML file (or the imzML part of an imzML file)
pub struct MzMLReader<B: BufRead> {
    reader: Reader<B>,

    decoder: Decoder,
    ontology: Ontology,
    errors: VecDeque<ParseError>,

    with_attribute_checks: bool,
    ignore_uncommon_tags: bool,

    breadcrumbs: VecDeque<(Tag, Option<String>)>,

    source_file_list: Vec<Arc<SourceFile>>,
    referenceable_param_group_list: Vec<Arc<ReferenceableParamGroup>>,
    sample_list: Vec<Arc<Sample>>,
    software_list: Vec<Arc<Software>>,
    scan_settings_list: ScanSettingsList,
    instrument_list: Vec<Arc<InstrumentConfiguration>>,
    data_processing_list: Vec<Arc<DataProcessing>>,

    representation: Option<Representation>,

    mzml: Option<MzML>,
}

impl<B: BufRead> MzMLReader<B> {
    /// Turn the `MzMLReader` into a `DataReader` to provide access to the data, using the reader specified
    pub fn with_data<R: Read + Seek>(self, data: R) -> Option<DataReader<R>> {
        Some(DataReader {
            mzml: self.mzml?,
            data_reader: Arc::new(Mutex::new(data)),
            errors: VecDeque::new(),
        })
    }
}

impl<B: BufRead> From<MzMLReader<B>> for MzML {
    fn from(mut parser: MzMLReader<B>) -> Self {
        parser.mzml.take().unwrap()
    }
}

impl<B: BufRead + Seek> From<MzMLReader<B>> for DataReader<B> {
    fn from(reader: MzMLReader<B>) -> Self {
        DataReader {
            mzml: reader.mzml.unwrap(),
            data_reader: Arc::new(Mutex::new(reader.reader.into_inner())),
            errors: reader.errors,
        }
    }
}

impl MzMLReader<BufReader<File>> {
    /// Create an MzMLReader from a path. This assumes that the ibd file has the exact same name and the extension '.ibd'
    pub fn from_path<P: AsRef<Path>>(
        path: P,
    ) -> Result<MzMLReader<BufReader<File>>, FatalParseError> {
        Self::from_path_with_ontology(path, MZML_ONTOLOGY.clone())
    }

    /// Create an MzMLReader from a path, using the supplied ontology to map CV params.
    /// This assumes that the ibd file has the exact same name and the extension '.ibd'
    pub fn from_path_with_ontology<P: AsRef<Path>>(
        path: P,
        ontology: Ontology,
    ) -> Result<MzMLReader<BufReader<File>>, FatalParseError> {
        let path = PathBuf::from(path.as_ref());
        let file = File::open(&path)?;

        Self::with_ontology(BufReader::new(file), ontology)
    }
}

impl<B: BufRead> MzMLReader<B> {
    /// Create a new parser from the supplied reader, using the default mzML ontology
    pub fn new(reader: B) -> Result<Self, FatalParseError> {
        Self::with_ontology(reader, MZML_ONTOLOGY.clone())
    }

    /// Create a new parser from the supplied reader and ontology
    pub fn with_ontology(reader: B, ontology: Ontology) -> Result<Self, FatalParseError> {
        let mut reader = quick_xml::Reader::from_reader(reader);
        reader.trim_text(true);

        let mut parser = MzMLReader {
            decoder: reader.decoder(),
            reader,
            ontology,
            errors: VecDeque::new(),
            with_attribute_checks: false,
            ignore_uncommon_tags: false,
            source_file_list: Vec::new(),
            referenceable_param_group_list: Vec::new(),
            sample_list: Vec::new(),
            software_list: Vec::new(),
            scan_settings_list: ScanSettingsList::new(),
            instrument_list: Vec::new(),
            data_processing_list: Vec::new(),

            breadcrumbs: VecDeque::new(),
            representation: None,

            mzml: None,
        };

        let mut buffer = Vec::with_capacity(BUFFER_SIZE);

        loop {
            // Skip comments and declarations
            let next_event = parser.next(&mut buffer)?;

            match next_event {
                Event::Comment(_) => {}
                Event::Decl(_) => {}
                Event::DocType(_) => {}
                Event::Text(_) => {}
                Event::Start(start_event) => match start_event.name().as_ref() {
                    b"mzML" => {
                        // We can unwrap here as this should never return None
                        let mut mzml = MzML::parse_start_tag(&mut parser, &start_event)?.unwrap();

                        mzml.parse_xml(&mut parser, &mut buffer)?;

                        parser.mzml = Some(mzml);

                        break;
                    }
                    b"indexedmzML" => {
                        // We can unwrap here as this should never return None
                        let mut indexed_mzml =
                            IndexedMzML::parse_start_tag(&mut parser, &start_event)?.unwrap();

                        indexed_mzml.parse_xml(&mut parser, &mut buffer)?;

                        parser.mzml = Some(indexed_mzml.into());

                        break;
                    }
                    _ => {
                        return Err(FatalParseError::UnexpectedTag(format!("{:?}", start_event)));
                    }
                },
                _ => {
                    return Err(FatalParseError::UnexpectedEvent(format!(
                        "{:?}",
                        next_event
                    )));
                    // panic!("Unexpected event: {:?}", next_event);
                }
            }
        }

        Ok(parser)
    }

    /// Returns the ontology used to describe controlled vocabulary when parsing
    pub fn ontology(&self) -> &Ontology {
        &self.ontology
    }

    pub(crate) fn last_tag(&self) -> &Tag {
        &self.breadcrumbs.back().unwrap().0
    }

    /// Returns the
    pub fn mzml(&self) -> Option<&MzML> {
        self.mzml.as_ref()
    }

    /// Clone the MzML instance that has been parsed
    pub fn clone_mzml(&self) -> Option<MzML> {
        self.mzml.clone()
    }

    // pub fn coordinates(&self) -> Vec<Coordinate> {

    //     if let Some(mzml) = &self.mzml {
    //     let coordinates = Vec::with_capacity(mzml.num_spectra());
    //         for i in 0..mzml.num_spectra() {
    //             let spectrum = mzml.spectrum(i).unwrap();
    //             spectrum.x_position()
    //         }

    //         coordinates
    //     } else {
    //         Vec::new()
    //     }
    // }

    pub(crate) fn source_file_ref(&self, id: &[u8]) -> Option<SourceFileRef> {
        let id = self.decoder.decode(id).unwrap();

        for source_file in &self.source_file_list {
            if source_file.id() == id {
                return Some(SourceFileRef::Ref(source_file.clone()));
            }
        }

        None
    }

    pub(crate) fn referenceable_param_group_ref(
        &self,
        id: &[u8],
    ) -> Option<ReferenceableParamGroupRef> {
        let id = self.decoder.decode(id).unwrap();

        for referenceable_param_group in &self.referenceable_param_group_list {
            if referenceable_param_group.id() == id {
                return Some(ReferenceableParamGroupRef::Ref(
                    referenceable_param_group.clone(),
                ));
            }
        }

        None
    }

    pub(crate) fn software_ref(&self, id: &[u8]) -> Option<SoftwareRef> {
        let id = self.decoder.decode(id).unwrap();

        for software_ref in &self.software_list {
            if software_ref.id() == id {
                return Some(SoftwareRef::Ref(software_ref.clone()));
            }
        }

        None
    }

    pub(crate) fn instrument_configuration_ref(
        &self,
        id: &[u8],
    ) -> Option<InstrumentConfigurationRef> {
        let id = self.decoder.decode(id).unwrap();

        for instrument in &self.instrument_list {
            if instrument.id() == id {
                return Some(InstrumentConfigurationRef::Ref(instrument.clone()));
            }
        }

        None
    }

    pub(crate) fn data_processing_ref(&self, id: &[u8]) -> Option<DataProcessingRef> {
        let id = self.decoder.decode(id).unwrap();

        for data_processing in &self.data_processing_list {
            if data_processing.id() == id {
                return Some(DataProcessingRef::Ref(data_processing.clone()));
            }
        }

        None
    }

    /// Return errors found during parsing
    pub fn errors(&self) -> &VecDeque<ParseError> {
        &self.errors
    }

    /// Returns mutable reference to error list
    pub fn errors_mut(&mut self) -> &mut VecDeque<ParseError> {
        &mut self.errors
    }

    /*pub fn peek_next(&mut self, buffer: &mut Vec<u8>) -> Result<&Event, Error> {
        let event = self.reader.read_event(buffer)?;

        self.current_event = Some(event.into_owned());

        Ok(self.current_event.as_ref().unwrap())
    }*/

    #[inline]
    pub(crate) fn next<'b>(&mut self, buffer: &'b mut Vec<u8>) -> Result<Event<'b>, Error> {
        self.reader.read_event_into(buffer)
    }

    pub(crate) fn process_attributes<'b>(
        &mut self,
        //parser: &mut Parser<'a, B>,
        tag: Tag,
        allowed: &'static HashMap<&'static [u8], AttributeDefinition>,
        e: &'b BytesStart<'b>,
    ) -> Result<HashMap<&'static str, AttributeValue<'b>>, FatalParseError> {
        let mut attributes = HashMap::with_capacity(allowed.capacity());

        for att in e.attributes().with_checks(self.with_attribute_checks) {
            match att {
                Ok(att) => {
                    let key = att.key.as_ref();

                    if let Some(definition) = allowed.get(key) {
                        match self.parse_string(tag, &att.value) {
                            Some(value) => {
                                let attribute_value = match definition.attribute_type {
                                    AttributeType::Integer => {
                                        let value = value.parse();
                                        match value {
                                            Ok(value) => AttributeValue::Integer(value),
                                            Err(error) => {
                                                self.errors
                                                    .push_back(ParseError::IntError((tag, error)));
                                                AttributeValue::String(att.value)
                                            }
                                        }
                                    }
                                    AttributeType::String => AttributeValue::String(att.value),
                                };

                                attributes.insert(definition.name.as_str(), attribute_value);
                            }
                            None => {
                                self.errors.push_back(ParseError::MissingAttributeValue((
                                    tag,
                                    definition.name.clone(),
                                    definition.attribute_type,
                                )));
                            }
                        }
                    } else {
                        self.errors.push_back(ParseError::UnexpectedAttribute((
                            tag,
                            std::str::from_utf8(att.key.as_ref())?.to_string(),
                        )));
                    }
                }
                Err(error) => {
                    self.errors
                        .push_back(ParseError::XMLError((tag, error.into())));
                }
            }
        }

        for (_name, definition) in allowed {
            if definition.required && !attributes.contains_key(definition.name.as_str()) {
                self.errors.push_back(ParseError::MissingAttribute((
                    tag,
                    definition.name.to_string(),
                )));
            }
        }

        Ok(attributes)
    }

    #[inline]
    fn parse_string<'b>(&mut self, tag: Tag, value: &'b [u8]) -> Option<&'b str> {
        match std::str::from_utf8(value) {
            Ok(value) => Some(value),
            Err(error) => {
                self.errors.push_back(ParseError::Utf8Error((tag, error)));

                None
            }
        }
    }
}

/// Enum describing the imzML tag.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub enum Tag {
    /// <indexedMzML>
    IndexedMzML,
    /// <indexList>
    IndexList,
    /// <mzML>
    MzML,
    /// <cvList>
    CVList,
    /// <cv>
    CV,
    /// <cvParam>
    CVParam,
    /// <userParam>
    UserParam,
    /// <fileDescription>
    FileDescription,
    /// <fileContent>
    FileContent,
    /// <sourceFileList>
    SourceFileList,
    /// <sourceFile>
    SourceFile,
    /// <sourceFileRefList>
    SourceFileRefList,
    /// <sourceFileRef>
    SourceFileRef,
    /// <contact>
    Contact,
    /// <referanceableParamGroupList>
    RefParamGroupList,
    /// <referenceableParamGroup>
    RefParamGroup,
    /// <referenceableParamGroupRef>
    RefParamGroupRef,
    /// <sampleList>
    SampleList,
    /// <sample>
    Sample,
    /// <softwareList>
    SoftwareList,
    /// <software>
    Software,
    /// <softwareRef>
    SoftwareRef,
    /// <scanSettingsList>
    ScanSettingsList,
    /// <scanSettings>
    ScanSettings,
    /// <instrumentConfigurationList>
    InstrumentConfigurationList,
    /// <instrumentConfiguration>
    InstrumentConfiguration,
    /// <componentList>
    ComponentList,
    /// General tag for either <source>, <analyser> or <detector>
    Component,
    /// <source>
    Source,
    /// <analyser>
    Analyser,
    /// <detector>
    Detector,
    /// <dataProcessingList>
    DataProcessingList,
    /// <dataProcessing>
    DataProcessing,
    /// <processingMethod>
    ProcessingMethod,
    /// <run>
    Run,
    /// <spectrumList>
    SpectrumList,
    /// <spectrum>
    Spectrum,
    /// <chromatogramList>
    ChromatogramList,
    /// <chromatogram>
    Chromatogram,
    /// <scanList>
    ScanList,
    /// <precursorList>
    PrecursorList,
    /// <precursor>
    Precursor,
    /// <product>
    Product,
    /// <isolationWindow>
    IsolationWindow,
    /// <selectedIonList>
    SelectedIonList,
    /// <selectedIon>
    SelectedIon,
    /// <activation>
    Activation,
    /// <scan>
    Scan,
    /// <scanWindowList>
    ScanWindowList,
    /// <scanWindow>
    ScanWindow,
    /// <binaryDataArrayList>
    BinaryDataArrayList,
    /// <binaryDataArray>
    BinaryDataArray,
    /// <binary>
    Binary,
}

impl Display for Tag {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Tag::IndexedMzML => write!(f, "indexedMzML"),
            Tag::IndexList => write!(f, "indexList"),
            Tag::MzML => write!(f, "mzML"),
            Tag::CVList => write!(f, "cvList"),
            Tag::CV => write!(f, "cv"),
            Tag::CVParam => write!(f, "cvParam"),
            Tag::UserParam => write!(f, "userParam"),
            Tag::FileDescription => write!(f, "fileDescription"),
            Tag::FileContent => write!(f, "fileContent"),
            Tag::SourceFileList => write!(f, "sourceFileList"),
            Tag::SourceFile => write!(f, "sourceFile"),
            Tag::SourceFileRefList => write!(f, "sourceFileRefList"),
            Tag::SourceFileRef => write!(f, "sourceFileRef"),
            Tag::Contact => write!(f, "contact"),
            Tag::RefParamGroupList => write!(f, "referenceableParamGroupList"),
            Tag::RefParamGroup => write!(f, "referenceableParamGroup"),
            Tag::RefParamGroupRef => write!(f, "referenceableParamGroupRef"),
            Tag::SampleList => write!(f, "sampleList"),
            Tag::Sample => write!(f, "sample"),
            Tag::SoftwareList => write!(f, "softwareList"),
            Tag::Software => write!(f, "software"),
            Tag::SoftwareRef => write!(f, "softwareRef"),
            Tag::ScanSettingsList => write!(f, "scanSettingsList"),
            Tag::ScanSettings => write!(f, "scanSettings"),
            Tag::InstrumentConfigurationList => write!(f, "instrumentConfigurationList"),
            Tag::InstrumentConfiguration => write!(f, "instrumentConfiguration"),
            Tag::ComponentList => write!(f, "componentList"),
            Tag::Component => write!(f, "source/analyzer/detector"),
            Tag::Source => write!(f, "source"),
            Tag::Analyser => write!(f, "analyzer"),
            Tag::Detector => write!(f, "detector"),
            Tag::DataProcessingList => write!(f, "dataProcessingList"),
            Tag::DataProcessing => write!(f, "dataProcessing"),
            Tag::ProcessingMethod => write!(f, "processingMethod"),
            Tag::Run => write!(f, "run"),
            Tag::SpectrumList => write!(f, "spectrumList"),
            Tag::Spectrum => write!(f, "spectrum"),
            Tag::ChromatogramList => write!(f, "chromatogramList"),
            Tag::Chromatogram => write!(f, "chromatogram"),
            Tag::ScanList => write!(f, "scanList"),
            Tag::PrecursorList => write!(f, "precursorList"),
            Tag::Precursor => write!(f, "precursor"),
            Tag::Product => write!(f, "product"),
            Tag::IsolationWindow => write!(f, "isolationWindow"),
            Tag::SelectedIonList => write!(f, "selectedIonList"),
            Tag::SelectedIon => write!(f, "selectedIon"),
            Tag::Activation => write!(f, "activation"),
            Tag::Scan => write!(f, "scan"),
            Tag::ScanWindowList => write!(f, "scanWindowList"),
            Tag::ScanWindow => write!(f, "scanWindow"),
            Tag::BinaryDataArrayList => write!(f, "binaryDataArrayList"),
            Tag::BinaryDataArray => write!(f, "binaryDataArray"),
            Tag::Binary => write!(f, "binary"),
        }
    }
}