//! Rust library for reading, writing, and manipulating SWC files for neuronal morphology.
//! Also includes a CLI for basic validation, sorting, and reindexing (with optional feature `cli`).
//!
//! The format was originally proposed in [Cannon, et al. 1998](http://dx.doi.org/10.1016/S0165-0270(98)00091-0),
//! with an implementation in [dohalloran/SWC_BATCH_CHECK](https://github.com/dohalloran/SWC_BATCH_CHECK).
//!
//! While commonly used, many variants exist; this implementation tries to cover the "standardised" version described
//! [here](http://www.neuronland.org/NLMorphologyConverter/MorphologyFormats/SWC/Spec.html),
//! with some ambiguities resolved by the [SWCplus specification](https://neuroinformatics.nl/swcPlus/).
//!
//! The header is a series of `#`-prefixed lines starting at the beginning of the file.
//! Blank lines (i.e. without a `#` or any other non-whitespace content) do not interrupt the header,
//! but are also not included in the parsed header.
//! All other `#`-prefixed and all whitespace-only lines in the file after the first sample are ignored.
//! Samples define a sample ID, a structure represented by that sample, its location, radius,
//! and the sample ID of the sample's parent in the neuron's tree structure (the root's parent is `-1`).
//!
//! The [SwcNeuron] is generic over implementors of [StructureIdentifier] and [Header],
//! and can be parsed from [Read]ers and dumped to [Write]rs.
//! [StructureIdentifier] is implemented for some known sub-specifications.
//! [Header] is implemented for [String] (i.e. a free-text header field).
//! [AnySwc] is a [SwcNeuron] with any structure integer and string header.
use std::cmp::Reverse;
use std::collections::{HashMap, HashSet};
use std::fmt::Debug;
use std::io::{BufRead, Write};
use std::iter::FromIterator;
use std::str::FromStr;

pub mod structures;

pub use crate::structures::{AnyStructure, StructureIdentifier};

pub mod header;
pub use crate::header::Header;

type SampleId = usize;
type Radius = f64;

/// Maximally flexible [SwcNeuron] with any structure specification and a free-text header.
pub type AnySwc = SwcNeuron<AnyStructure, String>;

/// Error in parsing a SWC sample (row).
#[derive(thiserror::Error, Debug)]
pub enum SampleParseError {
    /// A field which should be an integer could not be parsed.
    #[error("Integer field parse error")]
    IntField(#[from] std::num::ParseIntError),
    /// A field which should be a float/ decimal could not be parsed.
    #[error("Float field parse error")]
    FloatField(#[from] std::num::ParseFloatError),
    /// The value of the `structure` field was not recognised.
    /// Not checked for structure enums with a catch-all variant.
    #[error("Reference to unknown structure {0}")]
    UnknownStructure(isize),
    /// Line terminated early; more fields expected.
    #[error("Incorrect number of fields: found {0}")]
    IncorrectNumFields(usize),
}

/// Struct representing a row of a SWC file (a single sample from the neuron),
/// which gives information about a single node in the tree.
///
/// If the `parent_id` is None, this is a root node.
#[derive(Debug, Clone, Copy)]
pub struct SwcSample<S: StructureIdentifier> {
    pub sample_id: SampleId,
    pub structure: S,
    pub x: f64,
    pub y: f64,
    pub z: f64,
    pub radius: Radius,
    pub parent_id: Option<SampleId>,
}

impl<S: StructureIdentifier> FromStr for SwcSample<S> {
    type Err = SampleParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let trimmed = s.trim();
        let mut items = trimmed.split_whitespace();

        let sample_id = items
            .next()
            .ok_or(SampleParseError::IncorrectNumFields(0))?
            .parse()?;

        let struct_int = items
            .next()
            .ok_or(SampleParseError::IncorrectNumFields(1))?
            .parse::<isize>()?;
        let structure =
            S::try_from(struct_int).map_err(|_e| SampleParseError::UnknownStructure(struct_int))?;
        let x = items
            .next()
            .ok_or(SampleParseError::IncorrectNumFields(2))?
            .parse::<f64>()?;
        let y = items
            .next()
            .ok_or(SampleParseError::IncorrectNumFields(3))?
            .parse::<f64>()?;
        let z = items
            .next()
            .ok_or(SampleParseError::IncorrectNumFields(4))?
            .parse::<f64>()?;
        let radius = items
            .next()
            .ok_or(SampleParseError::IncorrectNumFields(5))?
            .parse::<f64>()?;
        let parent_id = match items.next() {
            Some("-1") => None,
            Some(p_str) => Some(p_str.parse()?),
            None => None,
        };

        let count: usize = items.fold(0, |x, _| x + 1);
        if count > 0 {
            return Err(SampleParseError::IncorrectNumFields(7 + count));
        }

        Ok(Self {
            sample_id,
            structure,
            x,
            y,
            z,
            radius,
            parent_id,
        })
    }
}

impl<S: StructureIdentifier> ToString for SwcSample<S> {
    fn to_string(&self) -> String {
        let structure: isize = self.structure.into();
        format!(
            "{} {} {} {} {} {} {}",
            self.sample_id,
            structure,
            self.x,
            self.y,
            self.z,
            self.radius,
            self.parent_id.map_or("-1".to_string(), |p| p.to_string()),
        )
    }
}

impl<S: StructureIdentifier> SwcSample<S> {
    /// Create a new [SwcSample] with  different sample and parent IDs.
    ///
    /// The other features are the same.
    fn with_ids(self, sample: SampleId, parent: Option<SampleId>) -> Self {
        SwcSample {
            sample_id: sample,
            structure: self.structure,
            x: self.x,
            y: self.y,
            z: self.z,
            radius: self.radius,
            parent_id: parent,
        }
    }
}

/// Struct representing a neuron skeleton as a tree of [SwcSample]s.
#[derive(Debug, Clone)]
pub struct SwcNeuron<S: StructureIdentifier, H: Header> {
    pub samples: Vec<SwcSample<S>>,
    pub header: Option<H>,
}

impl<S: StructureIdentifier, H: Header> FromIterator<SwcSample<S>> for SwcNeuron<S, H> {
    fn from_iter<I: IntoIterator<Item = SwcSample<S>>>(iter: I) -> Self {
        SwcNeuron {
            samples: iter.into_iter().collect(),
            header: None,
        }
    }
}

/// Error for a SWC file which cannot be parsed
#[derive(thiserror::Error, Debug)]
pub enum SwcParseError {
    /// A line could not be read from the file.
    #[error("Line read error")]
    Read(#[from] std::io::Error),
    /// The line could not be parsed as a SwcSample.
    #[error("Sample parse error")]
    SampleParse(#[from] SampleParseError),
    /// The SWC header could not be parsed.
    #[error("Header parse error")]
    HeaderParse(String),
}

/// Error where a sample refers to an unknown parent sample.
#[derive(thiserror::Error, Debug)]
#[error("Parent sample {parent_sample} does not exist")]
pub struct MissingSampleError {
    parent_sample: usize,
}

/// Error where a neuron represented by a SWC file is inconsistent.
#[derive(thiserror::Error, Debug)]
pub enum InconsistentNeuronError {
    /// Neuron is not a tree as it has multiple roots (parentless nodes).
    #[error("Neuron has >1 root")]
    MultipleRoots,
    /// Neuron is not a tree as it has several disconnected parts.
    #[error("Neuron has >1 connected component")]
    Disconnected,
    /// Neuron has no root: it could be empty or be cyclic.
    #[error("Neuron has no root")]
    NoRootError,
    /// Samples refer to unknown parent samples.
    #[error("Neuron has missing sample(s)")]
    MissingSample(#[from] MissingSampleError),
    /// Samples have clashing IDs.
    #[error("Neuron has multiple samples numbered {0}")]
    DuplicateSample(SampleId),
}

impl<S: StructureIdentifier, H: Header> SwcNeuron<S, H> {
    /// Sort the neuron's samples by their index.
    pub fn sort_index(mut self) -> Self {
        self.samples
            .sort_unstable_by(|a, b| a.sample_id.cmp(&b.sample_id));
        self
    }

    /// Re-index the neuron's samples in order of occurrence, starting at 1.
    /// Returns error if the SWC is missing a parent sample
    pub fn reindex(self) -> Result<Self, MissingSampleError> {
        let old_to_new: HashMap<SampleId, SampleId> = self
            .samples
            .iter()
            .enumerate()
            .map(|(idx, row)| (row.sample_id, idx + 1))
            .collect();

        let mut samples = Vec::with_capacity(self.samples.len());
        for (idx, row) in self.samples.into_iter().enumerate() {
            let parent;
            if let Some(old) = row.parent_id {
                parent = Some(
                    *old_to_new
                        .get(&old)
                        .ok_or(MissingSampleError { parent_sample: old })?,
                );
            } else {
                parent = None;
            }
            samples.push(row.with_ids(idx + 1, parent));
        }

        Ok(Self {
            samples,
            header: self.header,
        })
    }

    /// Re-order its samples in pre-order depth first search.
    /// Children are visited in the order of their original sample number.
    ///
    /// Returns an error if the neuron is not a self-consistent tree.
    pub fn sort_topo(self, reindex: bool) -> Result<Self, InconsistentNeuronError> {
        let mut parent_to_children: HashMap<SampleId, Vec<SampleId>> =
            HashMap::with_capacity(self.samples.len());
        let mut id_to_sample: HashMap<SampleId, SwcSample<S>> =
            HashMap::with_capacity(self.samples.len());
        let mut root = None;

        for row in self.samples {
            if let Some(p) = row.parent_id {
                let entry = parent_to_children.entry(p).or_default();
                entry.push(row.sample_id);
                if id_to_sample.insert(row.sample_id, row).is_some() {
                    return Err(InconsistentNeuronError::DuplicateSample(row.sample_id));
                }
            } else if root.is_some() {
                return Err(InconsistentNeuronError::MultipleRoots);
            } else {
                root = Some(row);
            }
        }

        let mut samples = Vec::with_capacity(id_to_sample.len() + 1);

        let mut to_visit = vec![];
        let mut next_id: SampleId = 1;

        if let Some(r) = root {
            let children = parent_to_children.remove(&r.sample_id).map_or_else(
                || Vec::with_capacity(0),
                |mut v| {
                    v.sort_unstable_by_key(|&x| Reverse(x));
                    v
                },
            );
            if reindex {
                samples.push(r.with_ids(next_id, None));
                to_visit.extend(children.into_iter().map(|c| (next_id, c)));
                next_id += 1;
            } else {
                samples.push(r);
                to_visit.extend(children.into_iter().map(|c| (r.sample_id, c)));
            }
        } else {
            return Err(InconsistentNeuronError::NoRootError);
        }

        while let Some((parent_id, row_id)) = to_visit.pop() {
            let row = id_to_sample
                .remove(&row_id)
                .expect("Just constructed this vec");
            let children = parent_to_children.remove(&row.sample_id).map_or_else(
                || Vec::with_capacity(0),
                |mut v| {
                    v.sort_unstable_by_key(|&id| Reverse(id));
                    v
                },
            );
            if reindex {
                samples.push(row.with_ids(next_id, Some(parent_id)));
                to_visit.extend(children.into_iter().map(|c| (next_id, c)));
                next_id += 1;
            } else {
                samples.push(row);
                to_visit.extend(children.into_iter().map(|c| (row.sample_id, c)));
            }
        }

        if id_to_sample.is_empty() {
            Ok(SwcNeuron {
                samples,
                header: self.header,
            })
        } else {
            Err(InconsistentNeuronError::Disconnected)
        }
        // todo: possible MissingSampleError
    }

    /// Ensure that [SwcNeuron] is a self-consistent tree.
    ///
    /// If `validate_order` is `false`, samples may be defined earlier in the file than their parents.
    pub fn validate(&self, validate_order: bool) -> Result<(), InconsistentNeuronError> {
        use InconsistentNeuronError::*;

        // use sample IDs
        let mut parent_to_children: HashMap<SampleId, Vec<SampleId>> =
            HashMap::with_capacity(self.samples.len());
        let mut sample_ids: HashSet<usize> = HashSet::with_capacity(self.samples.len());
        let mut root = None;

        for row in self.samples.iter() {
            if !sample_ids.insert(row.sample_id) {
                return Err(DuplicateSample(row.sample_id));
            }
            if let Some(p) = row.parent_id {
                if validate_order && !sample_ids.contains(&p) {
                    return Err(MissingSample(MissingSampleError { parent_sample: p }));
                }
                let entry = parent_to_children.entry(p).or_default();
                entry.push(row.sample_id);
            } else if root.is_some() {
                return Err(MultipleRoots);
            } else {
                root = Some(row.sample_id);
            }
        }

        let Some(rt) = root else {
            return Err(NoRootError);
        };

        let mut to_visit = Vec::default();
        to_visit.push(rt);

        while let Some(sample_id) = to_visit.pop() {
            let children = parent_to_children
                .remove(&sample_id)
                .unwrap_or_else(|| Vec::with_capacity(0));
            to_visit.extend(children.into_iter());
        }

        if parent_to_children.is_empty() {
            Ok(())
        } else if let Some((unknown_p, _)) = parent_to_children
            .iter()
            .take_while(|(p, _)| !sample_ids.contains(p))
            .next()
        {
            Err(MissingSample(MissingSampleError {
                parent_sample: *unknown_p,
            }))
        } else {
            Err(Disconnected)
        }
    }

    /// Parse a [SwcNeuron] from a [Read]er.
    ///
    /// Does not check the neuron for consistency, but does check for valid structures.
    pub fn from_reader<R: BufRead>(reader: R) -> Result<Self, SwcParseError> {
        let mut is_header = true;
        let mut header_lines = Vec::default();
        let mut samples = Vec::default();
        for line_res in parse_swc_lines::<R, S>(reader) {
            match line_res? {
                SwcLine::Comment(cmt) => {
                    if is_header {
                        header_lines.push(cmt);
                    }
                }
                SwcLine::Sample(s) => {
                    is_header = false;
                    samples.push(s);
                }
            }
        }

        let header: Option<H> = if header_lines.is_empty() {
            None
        } else {
            let header_str = header_lines.join("\n");
            Some(
                header_str
                    .parse()
                    .map_err(|_e| SwcParseError::HeaderParse(header_str))?,
            )
        };

        Ok(Self { samples, header })
    }

    /// Dump this [SwcNeuron] to a [Write]r.
    ///
    /// These writes are quite small: consider wrapping it in a [BufWriter].
    pub fn to_writer<W: Write>(&self, writer: &mut W) -> Result<(), std::io::Error> {
        if let Some(h) = self.header.clone() {
            for line in h.to_string().lines() {
                writeln!(writer, "#{}", line)?;
            }
        }
        for s in self.samples.iter() {
            writeln!(writer, "{}", s.to_string())?;
        }
        Ok(())
    }

    /// Replace the existing header with a new one, returning the existing.
    pub fn replace_header(&mut self, header: Option<H>) -> Option<H> {
        std::mem::replace(&mut self.header, header)
    }
}

pub enum SwcLine<S: StructureIdentifier> {
    Comment(String),
    Sample(SwcSample<S>),
}

/// Parse lines from a SWC file.
///
/// Skips blank lines and treats all `#`-prefixed lines as comments.
/// Does not check the neuron for consistency, but does check for valid structures.
pub fn parse_swc_lines<R: BufRead, S: StructureIdentifier>(
    reader: R,
) -> impl Iterator<Item = Result<SwcLine<S>, SwcParseError>> {
    reader.lines().filter_map(|ln_res| {
        let raw_line = match ln_res {
            Ok(ln) => ln,
            Err(e) => return Some(Err(SwcParseError::Read(e))),
        };
        let line = raw_line.trim_end();

        if line.is_empty() {
            None
        } else if let Some(remainder) = line.strip_prefix('#') {
            Some(Ok(SwcLine::Comment(remainder.to_string())))
        } else {
            match SwcSample::from_str(line) {
                Ok(sample) => Some(Ok(SwcLine::Sample(sample))),
                Err(err) => Some(Err(SwcParseError::SampleParse(err))),
            }
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs::{read_dir, File};
    use std::io::BufReader;
    use std::path::PathBuf;

    fn data_dir() -> PathBuf {
        let mut p = PathBuf::from_str(env!("CARGO_MANIFEST_DIR")).unwrap();
        p.push("data");
        p
    }

    fn data_files() -> impl IntoIterator<Item = PathBuf> {
        let root = data_dir();
        read_dir(&root).unwrap().into_iter().filter_map(|er| {
            let e = er.unwrap();
            let p = e.path();
            if p.is_file() && p.ends_with(".swc") {
                Some(p)
            } else {
                None
            }
        })
    }

    fn all_swcs() -> impl IntoIterator<Item = AnySwc> {
        data_files().into_iter().map(|p| {
            let f = File::open(&p).unwrap();
            AnySwc::from_reader(BufReader::new(f))
                .expect(format!("Could not read {:?}", p.as_os_str()).as_str())
        })
    }

    #[test]
    fn can_read() {
        for _ in all_swcs() {
            ();
        }
    }

    #[test]
    fn can_validate() {
        for swc in all_swcs() {
            swc.validate(true).expect("Invalid SWC");
        }
    }
}