// Released under MIT License.
// Copyright (c) 2023-2025 Ladislav Bartos

//! Structures and functions for efficient trajectory concatenation.

use std::collections::VecDeque;
use std::marker::PhantomData;
use std::path::{Path, PathBuf};

use crate::prelude::{TrajFullReadOpen, TrajGroupReadOpen};
use crate::{errors::ReadTrajError, io::traj_read::TrajRead, system::System};

use crate::io::traj_read::{
    FrameData, FrameDataTime, TrajRangeRead, TrajStepRead, TrajStepTimeRead,
};

#[cfg(feature = "parallel")]
use crate::prelude::TrajReadOpen;

use super::traj_read::{TrajFile, TrajReader};

pub struct TrajCat<'a, R: TrajRead<'a>> {
    remaining_files: VecDeque<PathBuf>,
    curr_reader: Option<R>,
    opener: Box<dyn FnMut(&Path) -> Result<R, ReadTrajError> + 'a>,
    last_time: Option<f32>,
    last_skipped_time: Option<f32>,
    _phantom: &'a PhantomData<R>,
}

impl<'a, R> TrajFile for TrajCat<'a, R> where R: TrajRead<'a> {}

impl<'a, R: TrajRead<'a>> TrajCat<'a, R> {
    /// Drop the current reader (closing the file) and open the next one.
    /// Returns Ok(true) if a new file was opened, Ok(false) if no files remain.
    fn open_next(&mut self) -> Result<bool, ReadTrajError> {
        self.curr_reader = None;
        match self.remaining_files.pop_front() {
            None => Ok(false),
            Some(path) => {
                self.curr_reader = Some((self.opener)(&path)?);
                Ok(true)
            }
        }
    }
}

pub struct TrajCatFrame<'a, R>
where
    R: TrajRead<'a>,
{
    frame: R::FrameData,
    _phantom: &'a PhantomData<R>,
}

impl<'a, R> FrameData for TrajCatFrame<'a, R>
where
    R: TrajRead<'a>,
    R::FrameData: FrameDataTime,
{
    type TrajFile = TrajCat<'a, R>;

    fn from_frame(
        trajcat: &mut Self::TrajFile,
        system: &System,
    ) -> Option<Result<Self, ReadTrajError>> {
        let frame = match trajcat.curr_reader.as_mut() {
            None => None,
            Some(trajectory) => {
                match R::FrameData::from_frame(trajectory.get_file_handle(), system) {
                    // no frame left -> move to next trajectory
                    None => {
                        trajcat.last_time = Some(system.get_simulation_time());
                        match trajcat.open_next() {
                            Err(e) => return Some(Err(e)),
                            Ok(false) => None,
                            Ok(true) => return Self::from_frame(trajcat, system),
                        }
                    }
                    // error while reading
                    Some(Err(e)) => Some(Err(e)),
                    // normal reading
                    Some(Ok(x)) => Some(Ok(TrajCatFrame {
                        frame: x,
                        _phantom: &PhantomData,
                    })),
                }
            }
        };

        // if we are at the boundary between two trajectory files,
        // check for matching times between the first and last frame
        match (frame, trajcat.last_time, trajcat.last_skipped_time) {
            (Some(Err(e)), _, _) => Some(Err(e)),
            (None, _, _) => None,
            (x, None, None) => x,
            (Some(Ok(curr)), Some(last), None) => {
                // skip this frame, if the time matches the time of the last read frame
                if last == curr.get_time() {
                    Self::from_frame(trajcat, system)
                } else {
                    trajcat.last_time = None;
                    Some(Ok(curr))
                }
            }
            (Some(Ok(curr)), None, Some(last)) => {
                // skip this frame, if the time matches the time of the last skipped frame
                if last == curr.get_time() {
                    Self::from_frame(trajcat, system)
                } else {
                    trajcat.last_time = None;
                    Some(Ok(curr))
                }
            }
            (Some(Ok(curr)), Some(last_read), Some(last_skipped)) => {
                // skip this frame, if the time matches the time of the last read frame or the last skipped frame
                if last_read == curr.get_time() || last_skipped == curr.get_time() {
                    Self::from_frame(trajcat, system)
                } else {
                    trajcat.last_time = None;
                    Some(Ok(curr))
                }
            }
        }
    }

    fn update_system(self, system: &mut System) {
        R::FrameData::update_system(self.frame, system)
    }
}

impl<'a, R> FrameDataTime for TrajCatFrame<'a, R>
where
    R: TrajRead<'a>,
    R::FrameData: FrameDataTime,
{
    fn get_time(&self) -> f32 {
        self.frame.get_time()
    }
}

pub struct TrajConcatenator<'a, R: TrajRead<'a>> {
    // safety: only this system pointer can be used!
    system: *mut System,
    traj_cat: TrajCat<'a, R>,
    _phantom: &'a PhantomData<R>,
}

impl<'a, R> TrajRead<'a> for TrajConcatenator<'a, R>
where
    R: TrajRead<'a>,
    R::FrameData: FrameDataTime,
{
    type FrameData = TrajCatFrame<'a, R>;

    fn get_system(&mut self) -> *mut System {
        self.system
    }

    fn get_file_handle(
        &mut self,
    ) -> &mut <<Self as TrajRead<'a>>::FrameData as FrameData>::TrajFile {
        &mut self.traj_cat
    }
}

impl<'a, R> TrajRangeRead<'a> for TrajConcatenator<'a, R>
where
    R: TrajRangeRead<'a>,
    R::FrameData: FrameDataTime,
{
    fn jump_to_start(&mut self, start_time: f32) -> Result<(), ReadTrajError> {
        // first try the currently open reader
        if let Some(reader) = self.traj_cat.curr_reader.as_mut() {
            if reader.jump_to_start(start_time).is_ok() {
                return Ok(());
            }
        }

        // then try remaining files one by one
        while let Some(path) = self.traj_cat.remaining_files.pop_front() {
            let mut reader = (self.traj_cat.opener)(&path)?;
            if reader.jump_to_start(start_time).is_ok() {
                self.traj_cat.curr_reader = Some(reader);
                return Ok(());
            }
        }

        self.traj_cat.curr_reader = None;
        // if starting time is not found in any trajectory, return an error
        Err(ReadTrajError::StartNotFound(start_time.to_string()))
    }
}

impl<'a, R> TrajStepRead<'a> for TrajConcatenator<'a, R>
where
    R: TrajStepTimeRead<'a>,
    R::FrameData: FrameDataTime,
{
    fn skip_frame(&mut self) -> Result<bool, ReadTrajError> {
        let traj = match self.traj_cat.curr_reader.as_mut() {
            None => return Ok(false),
            Some(x) => x,
        };

        match (
            traj.skip_frame_time(),
            self.traj_cat.last_skipped_time,
            self.traj_cat.last_time,
        ) {
            // if the frame could not be skipped (EOF), move to another trajectory and attempt to skip a frame
            (Ok(None), _, _) => {
                // we are at the file boundary; store the simulation time from the last frame
                unsafe {
                    self.traj_cat.last_time = Some((*self.get_system()).get_simulation_time());
                }
                match self.traj_cat.open_next() {
                    Err(e) => Err(e),
                    Ok(false) => Ok(false),
                    Ok(true) => self.skip_frame(),
                }
            }

            // if an error is returned, propagate it
            (Err(e), _, _) => Err(e),

            // if the frame was skipped and no `last_skipped_frame` is stored, just store
            // the time of the frame
            (Ok(Some(time)), None, None) => {
                self.traj_cat.last_skipped_time = Some(time);
                Ok(true)
            }

            (Ok(Some(time)), None, Some(last)) => {
                // the frame that was skipped has the same time as the last READ frame
                // skip one more frame
                if time == last {
                    self.skip_frame()
                } else {
                    self.traj_cat.last_skipped_time = Some(time);
                    Ok(true)
                }
            }

            (Ok(Some(time)), Some(last), None) => {
                // the frame was skipped but it had the same time as the last skipped frame
                // skip one more frame
                if time == last {
                    self.skip_frame()
                } else {
                    self.traj_cat.last_skipped_time = Some(time);
                    Ok(true)
                }
            }

            (Ok(Some(time)), Some(last_skipped), Some(last_read)) => {
                // the frame was skipped but it has the same time as the last skipped or read frame
                // skip one more frame
                if time == last_skipped || time == last_read {
                    self.skip_frame()
                } else {
                    self.traj_cat.last_skipped_time = Some(time);
                    Ok(true)
                }
            }
        }
    }
}

impl<'a, R> TrajConcatenator<'a, R>
where
    R: TrajRead<'a>,
    R::FrameData: FrameDataTime,
{
    pub(super) fn new(
        filenames: &[impl AsRef<Path>],
        system: *mut System,
        mut opener: Box<dyn FnMut(&Path) -> Result<R, ReadTrajError> + 'a>,
    ) -> Result<TrajReader<'a, Self>, ReadTrajError> {
        let mut remaining_files: VecDeque<PathBuf> =
            filenames.iter().map(|f| f.as_ref().to_path_buf()).collect();

        let first_path = remaining_files
            .pop_front()
            .expect("FATAL GROAN ERROR | TrajConcatenator::new | No trajectory files provided.");

        let curr_reader = Some(opener(&first_path)?);

        let traj_cat = TrajCat {
            remaining_files: remaining_files,
            curr_reader,
            opener: opener,
            last_time: None,
            last_skipped_time: None,
            _phantom: &PhantomData,
        };

        let concatenator = TrajConcatenator {
            system,
            traj_cat,
            _phantom: &PhantomData,
        };

        Ok(TrajReader::wrap_traj(concatenator))
    }
}

/***************************************/
/*        Generic System methods       */
/***************************************/

impl System {
    /// Iterate through multiple trajectory files.
    /// Any duplicate frames at the boundaries of the trajectories are skipped.
    ///
    /// ## Examples
    /// Iterate across several xtc trajectories.
    /// ```no_run
    /// # use groan_rs::prelude::*;
    /// #
    /// let mut system = System::from_file("system.gro").unwrap();
    /// let trajectories = vec!["md0001.xtc", "md0002.xtc", "md0003.xtc"];
    ///
    /// for frame in system.traj_cat_iter::<XtcReader>(&trajectories).unwrap() {
    ///     let frame = frame.unwrap();
    ///    
    ///     // perform some operation with the frame
    /// }
    /// ```
    /// Iterate across frames in a specified time interval (100 to 300 ns) that may cross several trajectories.
    /// In the specified time interval, read only every 5th frame.
    /// Also print progress of the trajectory reading.
    /// ```no_run
    /// # use groan_rs::prelude::*;
    /// #
    /// let mut system = System::from_file("system.gro").unwrap();
    /// let trajectories = vec!["md0001.xtc", "md0002.xtc", "md0003.xtc"];
    ///
    /// for frame in system
    ///     .traj_cat_iter::<XtcReader>(&trajectories)
    ///     .unwrap()
    ///     .print_progress(ProgressPrinter::default())
    ///     .with_range(100_000.0, 300_000.0)
    ///     .unwrap()
    ///     .with_step(5)
    ///     .unwrap()
    /// {
    ///     let frame = frame.unwrap();
    ///     // perform some operation with the frame
    /// }
    /// ```
    ///
    /// ## Notes
    /// - The trajectory files are not reordered in any way and are read in the order in which they are provided.
    /// - It is quite common that the last frame of a simulation trajectory is the same frame as the first frame
    ///   of the following simulation trajectory. When concatenating trajectories using this function,
    ///   any duplicate frames, i.e., frames with matching simulation time, at the trajectory boundaries are removed.
    /// - ONLY duplicate frames at the trajectory boundaries are removed, all other frames are kept and iterated through.
    /// - For instance, the following three trajectory files (each composed of 4 frames with specified time)
    ///   will be concatenated into one file like this:
    /// ```text
    ///       trajectory 1            trajectory 2              trajectory 3
    /// [ 0 , 100 , 200 , 300 ] [ 300 , 400 , 500 , 600 ] [ 100 , 200 , 300 , 400 ]
    /// ----->
    ///                   concatenated trajectory
    /// [ 0 , 100 , 200 , 300 , 400 , 500 , 600 , 100 , 200 , 300 , 400 ]
    /// ```
    /// - All the concatenated trajectory files must be of the same type, i.e. it is not possible to concatenate
    ///   e.g. `xtc` with `trr` files.
    /// - Only trajectories which contain information about the simulation time can be concatenated using this function.
    pub fn traj_cat_iter<'a, Read>(
        &'a mut self,
        filenames: &[impl AsRef<Path>],
    ) -> Result<TrajReader<'a, TrajConcatenator<'a, Read>>, ReadTrajError>
    where
        Read: TrajFullReadOpen<'a>,
        Read::FrameData: FrameDataTime,
    {
        if filenames.is_empty() {
            return Err(ReadTrajError::CatNoTrajectories);
        }

        let system = self as *mut System;

        let opener = move |path: &Path| -> Result<Read, ReadTrajError> {
            unsafe { Read::new(&mut *system, path) }
        };

        TrajConcatenator::new(filenames, system, Box::new(opener))
    }

    /// Iterate through multiple trajectory files while reading only the coordinates
    /// of the specified group.
    ///
    /// In all other aspects, works the same as [`System::traj_cat_iter`].
    ///
    /// See also [`System::group_xtc_iter`] for more information about partial frame reading.
    ///
    pub fn group_traj_cat_iter<'a, Read>(
        &mut self,
        filenames: &[impl AsRef<Path>],
        group: &str,
    ) -> Result<TrajReader<'a, TrajConcatenator<'a, Read>>, ReadTrajError>
    where
        Read: TrajGroupReadOpen<'a>,
        Read::FrameData: FrameDataTime,
    {
        if filenames.is_empty() {
            return Err(ReadTrajError::CatNoTrajectories);
        }

        let system = self as *mut System;

        let group = group.to_owned();
        let opener = move |path: &Path| -> Result<Read, ReadTrajError> {
            unsafe { Read::new(&mut *system, path, &group) }
        };

        TrajConcatenator::new(filenames, system, Box::new(opener))
    }
}

#[cfg(feature = "parallel")]
impl System {
    /// Iterate through multiple trajectory files while either reading the full frames
    /// or only a specific group from each frame.
    ///
    /// Helper method for dispatching into either [`System::traj_cat_iter`] or [`System::group_traj_cat_iter`].
    pub(crate) fn traj_cat_iter_initialize<'a, Read>(
        &mut self,
        filenames: &[impl AsRef<Path>],
        group: Option<&str>,
    ) -> Result<TrajReader<'a, TrajConcatenator<'a, Read>>, ReadTrajError>
    where
        Read: TrajReadOpen<'a>,
        Read::FrameData: FrameDataTime,
    {
        if filenames.is_empty() {
            return Err(ReadTrajError::CatNoTrajectories);
        }

        let system = self as *mut System;

        let group = group.map(|g| g.to_string());
        let opener: Box<dyn FnMut(&Path) -> Result<Read, ReadTrajError> + 'a> =
            Box::new(move |path: &Path| unsafe {
                Read::initialize(&mut *system, path, group.as_deref())
            });

        TrajConcatenator::new(filenames, system, opener)
    }
}

/**************************/
/*       UNIT TESTS       */
/**************************/

#[cfg(test)]
mod tests {
    use std::fs::File;

    use float_cmp::assert_approx_eq;

    use crate::io::traj_read::TrajMasterRead;
    #[cfg(not(feature = "no-xdrfile"))]
    use crate::io::trr_io::TrrReader;
    use crate::io::xtc_io::XtcReader;
    use crate::progress::ProgressPrinter;
    use crate::test_utilities::*;

    use self::utilities::*;

    use super::*;

    #[test]
    fn cat_xtc_simple() {
        let mut system_single = System::from_file("test_files/example.gro").unwrap();
        let mut system_cat = System::from_file("test_files/example.gro").unwrap();

        let traj_single = system_single
            .xtc_iter("test_files/short_trajectory.xtc")
            .unwrap();
        let traj_cat = system_cat
            .traj_cat_iter::<XtcReader>(&[
                "test_files/split/traj1.xtc",
                "test_files/split/traj2.xtc",
                "test_files/split/traj3.xtc",
                "test_files/split/traj4.xtc",
                "test_files/split/traj5.xtc",
                "test_files/split/traj6.xtc",
            ])
            .unwrap();

        for (frame_single, frame_cat) in traj_single.zip(traj_cat) {
            let frame_single = frame_single.unwrap();
            let frame_cat = frame_cat.unwrap();

            assert_approx_eq!(
                f32,
                frame_single.get_simulation_time(),
                frame_cat.get_simulation_time()
            );
            assert_eq!(
                frame_single.get_simulation_step(),
                frame_cat.get_simulation_step()
            );

            compare_box(
                frame_single.get_box().unwrap(),
                frame_cat.get_box().unwrap(),
            );

            for (atom_single, atom_cat) in frame_single.atoms_iter().zip(frame_cat.atoms_iter()) {
                compare_atoms(atom_single, atom_cat);
            }
        }
    }

    #[test]
    fn cat_xtc_with_ranges() {
        let mut system_single = System::from_file("test_files/example.gro").unwrap();
        let mut system_cat = System::from_file("test_files/example.gro").unwrap();

        let ranges = vec![(0.0, 570.0), (320.0, f32::MAX), (320.0, 500.0)];

        for (start, end) in ranges {
            let traj_single = system_single
                .xtc_iter("test_files/short_trajectory.xtc")
                .unwrap()
                .with_range(start, end)
                .unwrap();
            let traj_cat = system_cat
                .traj_cat_iter::<XtcReader>(&[
                    "test_files/split/traj1.xtc",
                    "test_files/split/traj2.xtc",
                    "test_files/split/traj3.xtc",
                    "test_files/split/traj4.xtc",
                    "test_files/split/traj5.xtc",
                    "test_files/split/traj6.xtc",
                ])
                .unwrap()
                .with_range(start, end)
                .unwrap();

            for (frame_single, frame_cat) in traj_single.zip(traj_cat) {
                let frame_single = frame_single.unwrap();
                let frame_cat = frame_cat.unwrap();

                assert_approx_eq!(
                    f32,
                    frame_single.get_simulation_time(),
                    frame_cat.get_simulation_time()
                );
                assert_eq!(
                    frame_single.get_simulation_step(),
                    frame_cat.get_simulation_step()
                );

                compare_box(
                    frame_single.get_box().unwrap(),
                    frame_cat.get_box().unwrap(),
                );

                for (atom_single, atom_cat) in frame_single.atoms_iter().zip(frame_cat.atoms_iter())
                {
                    compare_atoms(atom_single, atom_cat);
                }
            }
        }
    }

    #[test]
    fn cat_xtc_steps() {
        for step in 2..=11 {
            let mut system_single = System::from_file("test_files/example.gro").unwrap();
            let mut system_cat = System::from_file("test_files/example.gro").unwrap();

            let traj_single = system_single
                .xtc_iter("test_files/short_trajectory.xtc")
                .unwrap()
                .with_step(step)
                .unwrap();
            let traj_cat = system_cat
                .traj_cat_iter::<XtcReader>(&[
                    "test_files/split/traj1.xtc",
                    "test_files/split/traj2.xtc",
                    "test_files/split/traj3.xtc",
                    "test_files/split/traj4.xtc",
                    "test_files/split/traj5.xtc",
                    "test_files/split/traj6.xtc",
                ])
                .unwrap()
                .with_step(step)
                .unwrap();

            for (frame_single, frame_cat) in traj_single.zip(traj_cat) {
                let frame_single = frame_single.unwrap();
                let frame_cat = frame_cat.unwrap();

                assert_approx_eq!(
                    f32,
                    frame_single.get_simulation_time(),
                    frame_cat.get_simulation_time()
                );
                assert_eq!(
                    frame_single.get_simulation_step(),
                    frame_cat.get_simulation_step()
                );

                compare_box(
                    frame_single.get_box().unwrap(),
                    frame_cat.get_box().unwrap(),
                );

                for (atom_single, atom_cat) in frame_single.atoms_iter().zip(frame_cat.atoms_iter())
                {
                    compare_atoms(atom_single, atom_cat);
                }
            }
        }
    }

    #[test]
    fn cat_xtc_duplicate_not_at_boundary() {
        let mut system_cat = System::from_file("test_files/example.gro").unwrap();

        let traj_cat = system_cat
            .traj_cat_iter::<XtcReader>(&[
                "test_files/split/traj1.xtc",
                "test_files/split/traj2.xtc",
                "test_files/split/traj3b.xtc",
                "test_files/split/traj4.xtc",
                "test_files/split/traj5.xtc",
                "test_files/split/traj6.xtc",
            ])
            .unwrap();

        for (frame, time) in traj_cat.zip([
            0.0, 100.0, 200.0, 300.0, 300.0, 400.0, 500.0, 600.0, 700.0, 800.0, 900.0, 1000.0,
        ]) {
            let frame = frame.unwrap();
            assert_approx_eq!(f32, frame.get_simulation_time(), time);
        }
    }

    #[test]
    fn cat_xtc_duplicate_not_at_boundary_step3() {
        let mut system_cat = System::from_file("test_files/example.gro").unwrap();

        let traj_cat = system_cat
            .traj_cat_iter::<XtcReader>(&[
                "test_files/split/traj1.xtc",
                "test_files/split/traj2.xtc",
                "test_files/split/traj3b.xtc",
                "test_files/split/traj4.xtc",
                "test_files/split/traj5.xtc",
                "test_files/split/traj6.xtc",
            ])
            .unwrap()
            .with_step(3)
            .unwrap();

        for (frame, time) in traj_cat.zip([0.0, 300.0, 500.0, 800.0]) {
            let frame = frame.unwrap();
            assert_approx_eq!(f32, frame.get_simulation_time(), time);
        }
    }

    #[test]
    fn cat_xtc_steps_with_ranges() {
        let ranges = vec![(0.0, 570.0), (320.0, f32::MAX), (220.0, 800.0)];

        for step in 2..=11 {
            for (start, end) in &ranges {
                let mut system_single = System::from_file("test_files/example.gro").unwrap();
                let mut system_cat = System::from_file("test_files/example.gro").unwrap();

                let traj_single = system_single
                    .xtc_iter("test_files/short_trajectory.xtc")
                    .unwrap()
                    .with_step(step)
                    .unwrap()
                    .with_range(*start, *end)
                    .unwrap();
                let traj_cat = system_cat
                    .traj_cat_iter::<XtcReader>(&[
                        "test_files/split/traj1.xtc",
                        "test_files/split/traj2.xtc",
                        "test_files/split/traj3.xtc",
                        "test_files/split/traj4.xtc",
                        "test_files/split/traj5.xtc",
                        "test_files/split/traj6.xtc",
                    ])
                    .unwrap()
                    .with_step(step)
                    .unwrap()
                    .with_range(*start, *end)
                    .unwrap();

                for (frame_single, frame_cat) in traj_single.zip(traj_cat) {
                    let frame_single = frame_single.unwrap();
                    let frame_cat = frame_cat.unwrap();

                    assert_approx_eq!(
                        f32,
                        frame_single.get_simulation_time(),
                        frame_cat.get_simulation_time()
                    );
                    assert_eq!(
                        frame_single.get_simulation_step(),
                        frame_cat.get_simulation_step()
                    );

                    compare_box(
                        frame_single.get_box().unwrap(),
                        frame_cat.get_box().unwrap(),
                    );

                    for (atom_single, atom_cat) in
                        frame_single.atoms_iter().zip(frame_cat.atoms_iter())
                    {
                        compare_atoms(atom_single, atom_cat);
                    }
                }
            }
        }
    }

    #[test]
    #[cfg(not(feature = "no-xdrfile"))]
    fn cat_trr_simple() {
        let mut system_single = System::from_file("test_files/example.gro").unwrap();
        let mut system_cat = System::from_file("test_files/example.gro").unwrap();

        let traj_single = system_single
            .trr_iter("test_files/short_trajectory.trr")
            .unwrap();
        let traj_cat = system_cat
            .traj_cat_iter::<TrrReader>(&[
                "test_files/split/traj1.trr",
                "test_files/split/traj2.trr",
                "test_files/split/traj3.trr",
                "test_files/split/traj4.trr",
                "test_files/split/traj5.trr",
                "test_files/split/traj6.trr",
            ])
            .unwrap();

        for (frame_single, frame_cat) in traj_single.zip(traj_cat) {
            let frame_single = frame_single.unwrap();
            let frame_cat = frame_cat.unwrap();

            assert_approx_eq!(
                f32,
                frame_single.get_simulation_time(),
                frame_cat.get_simulation_time()
            );
            assert_eq!(
                frame_single.get_simulation_step(),
                frame_cat.get_simulation_step()
            );

            compare_box(
                frame_single.get_box().unwrap(),
                frame_cat.get_box().unwrap(),
            );

            for (atom_single, atom_cat) in frame_single.atoms_iter().zip(frame_cat.atoms_iter()) {
                compare_atoms(atom_single, atom_cat);
            }
        }
    }

    #[test]
    #[cfg(not(feature = "no-xdrfile"))]
    fn cat_trr_with_ranges() {
        let mut system_single = System::from_file("test_files/example.gro").unwrap();
        let mut system_cat = System::from_file("test_files/example.gro").unwrap();

        let ranges = vec![(0.0, 400.0), (250.0, f32::MAX), (250.0, 400.0)];

        for (start, end) in ranges {
            let traj_single = system_single
                .trr_iter("test_files/short_trajectory.trr")
                .unwrap()
                .with_range(start, end)
                .unwrap();
            let traj_cat = system_cat
                .traj_cat_iter::<TrrReader>(&[
                    "test_files/split/traj1.trr",
                    "test_files/split/traj2.trr",
                    "test_files/split/traj3.trr",
                    "test_files/split/traj4.trr",
                    "test_files/split/traj5.trr",
                    "test_files/split/traj6.trr",
                ])
                .unwrap()
                .with_range(start, end)
                .unwrap();

            for (frame_single, frame_cat) in traj_single.zip(traj_cat) {
                let frame_single = frame_single.unwrap();
                let frame_cat = frame_cat.unwrap();

                assert_approx_eq!(
                    f32,
                    frame_single.get_simulation_time(),
                    frame_cat.get_simulation_time()
                );
                assert_eq!(
                    frame_single.get_simulation_step(),
                    frame_cat.get_simulation_step()
                );

                compare_box(
                    frame_single.get_box().unwrap(),
                    frame_cat.get_box().unwrap(),
                );

                for (atom_single, atom_cat) in frame_single.atoms_iter().zip(frame_cat.atoms_iter())
                {
                    compare_atoms(atom_single, atom_cat);
                }
            }
        }
    }

    #[test]
    #[cfg(not(feature = "no-xdrfile"))]
    fn cat_trr_steps() {
        for step in 2..=11 {
            let mut system_single = System::from_file("test_files/example.gro").unwrap();
            let mut system_cat = System::from_file("test_files/example.gro").unwrap();

            let traj_single = system_single
                .trr_iter("test_files/short_trajectory.trr")
                .unwrap()
                .with_step(step)
                .unwrap();
            let traj_cat = system_cat
                .traj_cat_iter::<TrrReader>(&[
                    "test_files/split/traj1.trr",
                    "test_files/split/traj2.trr",
                    "test_files/split/traj3.trr",
                    "test_files/split/traj4.trr",
                    "test_files/split/traj5.trr",
                    "test_files/split/traj6.trr",
                ])
                .unwrap()
                .with_step(step)
                .unwrap();

            for (frame_single, frame_cat) in traj_single.zip(traj_cat) {
                let frame_single = frame_single.unwrap();
                let frame_cat = frame_cat.unwrap();

                assert_approx_eq!(
                    f32,
                    frame_single.get_simulation_time(),
                    frame_cat.get_simulation_time()
                );
                assert_eq!(
                    frame_single.get_simulation_step(),
                    frame_cat.get_simulation_step()
                );

                compare_box(
                    frame_single.get_box().unwrap(),
                    frame_cat.get_box().unwrap(),
                );

                for (atom_single, atom_cat) in frame_single.atoms_iter().zip(frame_cat.atoms_iter())
                {
                    compare_atoms(atom_single, atom_cat);
                }
            }
        }
    }

    #[test]
    #[cfg(not(feature = "no-xdrfile"))]
    fn cat_trr_steps_with_ranges() {
        let ranges = vec![(0.0, 400.0), (250.0, f32::MAX), (250.0, 400.0)];

        for step in 2..=11 {
            for (start, end) in &ranges {
                let mut system_single = System::from_file("test_files/example.gro").unwrap();
                let mut system_cat = System::from_file("test_files/example.gro").unwrap();

                let traj_single = system_single
                    .trr_iter("test_files/short_trajectory.trr")
                    .unwrap()
                    .with_step(step)
                    .unwrap()
                    .with_range(*start, *end)
                    .unwrap();
                let traj_cat = system_cat
                    .traj_cat_iter::<TrrReader>(&[
                        "test_files/split/traj1.trr",
                        "test_files/split/traj2.trr",
                        "test_files/split/traj3.trr",
                        "test_files/split/traj4.trr",
                        "test_files/split/traj5.trr",
                        "test_files/split/traj6.trr",
                    ])
                    .unwrap()
                    .with_step(step)
                    .unwrap()
                    .with_range(*start, *end)
                    .unwrap();

                for (frame_single, frame_cat) in traj_single.zip(traj_cat) {
                    let frame_single = frame_single.unwrap();
                    let frame_cat = frame_cat.unwrap();

                    assert_approx_eq!(
                        f32,
                        frame_single.get_simulation_time(),
                        frame_cat.get_simulation_time()
                    );
                    assert_eq!(
                        frame_single.get_simulation_step(),
                        frame_cat.get_simulation_step()
                    );

                    compare_box(
                        frame_single.get_box().unwrap(),
                        frame_cat.get_box().unwrap(),
                    );

                    for (atom_single, atom_cat) in
                        frame_single.atoms_iter().zip(frame_cat.atoms_iter())
                    {
                        compare_atoms(atom_single, atom_cat);
                    }
                }
            }
        }
    }

    #[test]
    fn cat_traj_print_progress() {
        let mut system_single = System::from_file("test_files/example.gro").unwrap();
        let mut system_cat = System::from_file("test_files/example.gro").unwrap();

        let output_single = File::create("tmp_cat_traj_print_progress_single.txt").unwrap();
        let output_cat = File::create("tmp_cat_traj_print_progress_cat.txt").unwrap();

        let single_printer = ProgressPrinter::new()
            .with_print_freq(3)
            .with_output(Box::from(output_single))
            .with_colored(false);
        let cat_printer = ProgressPrinter::new()
            .with_print_freq(3)
            .with_output(Box::from(output_cat))
            .with_colored(false);

        let traj_single = system_single
            .xtc_iter("test_files/short_trajectory.xtc")
            .unwrap()
            .print_progress(single_printer);
        let traj_cat = system_cat
            .traj_cat_iter::<XtcReader>(&[
                "test_files/split/traj1.xtc",
                "test_files/split/traj2.xtc",
                "test_files/split/traj3.xtc",
                "test_files/split/traj4.xtc",
                "test_files/split/traj5.xtc",
                "test_files/split/traj6.xtc",
            ])
            .unwrap()
            .print_progress(cat_printer);

        for frame in traj_single {
            frame.unwrap();
        }

        for frame in traj_cat {
            frame.unwrap();
        }

        let mut result = File::open("tmp_cat_traj_print_progress_cat.txt").unwrap();
        let mut expected = File::open("tmp_cat_traj_print_progress_single.txt").unwrap();
        assert!(file_diff::diff_files(&mut result, &mut expected));

        std::fs::remove_file("tmp_cat_traj_print_progress_cat.txt").unwrap();
        std::fs::remove_file("tmp_cat_traj_print_progress_single.txt").unwrap();
    }

    #[test]
    fn cat_traj_duplicates() {
        let mut system = System::from_file("test_files/example.gro").unwrap();

        let times = [
            0.0, 100.0, 200.0, 300.0, 400.0, 0.0, 100.0, 200.0, 300.0, 400.0, 900.0, 1000.0, 0.0,
            100.0, 200.0,
        ];

        for (i, frame) in system
            .traj_cat_iter::<XtcReader>(&[
                "test_files/split/traj1.xtc",
                "test_files/split/traj2.xtc",
                "test_files/split/traj3.xtc",
                "test_files/split/traj1.xtc",
                "test_files/split/traj3.xtc",
                "test_files/split/traj6.xtc",
                "test_files/split/traj1.xtc",
                "test_files/split/traj2.xtc",
            ])
            .unwrap()
            .enumerate()
        {
            let frame = frame.unwrap();
            assert_approx_eq!(f32, frame.get_simulation_time(), times[i]);
        }
    }

    #[test]
    fn cat_traj_empty() {
        let mut system = System::from_file("test_files/example.gro").unwrap();
        let empty: Vec<&str> = vec![];

        match system.traj_cat_iter::<XtcReader>(&empty) {
            Ok(_) => panic!("Function should have failed."),
            Err(e) => assert_eq!(ReadTrajError::CatNoTrajectories, e),
        };
    }

    #[test]
    #[cfg(not(feature = "no-xdrfile"))]
    fn cat_traj_second_file_is_nonexistent() {
        let mut system = System::from_file("test_files/example.gro").unwrap();
        let files: Vec<&str> = vec![
            "test_files/split/traj1.trr",
            "test_files/split/traj_nonexistent.trr",
            "test_files/split/traj3.trr",
        ];

        assert!(system
            .traj_cat_iter::<TrrReader>(&files)
            .unwrap()
            .any(|frame| matches!(frame, Err(ReadTrajError::FileNotFound(_)))));
    }

    #[test]
    #[cfg(not(feature = "no-xdrfile"))]
    fn cat_traj_first_file_is_nonexistent() {
        let mut system = System::from_file("test_files/example.gro").unwrap();
        let files: Vec<&str> = vec![
            "test_files/split/traj_nonexistent.trr",
            "test_files/split/traj1.trr",
            "test_files/split/traj3.trr",
        ];

        match system.traj_cat_iter::<TrrReader>(&files) {
            Ok(_) => panic!("Function should have failed."),
            Err(ReadTrajError::FileNotFound(file)) => assert_eq!(
                file.to_str().unwrap(),
                "test_files/split/traj_nonexistent.trr"
            ),
            Err(e) => panic!("Incorrect error type returned `{}`", e),
        };
    }

    #[test]
    #[cfg(feature = "molly")]
    fn cat_xtc_group_simple() {
        use crate::prelude::GroupXtcReader;

        let mut system_single = System::from_file("test_files/example.gro").unwrap();
        system_single.group_create("Protein", "@protein").unwrap();
        let mut system_cat = System::from_file("test_files/example.gro").unwrap();
        system_cat.group_create("Protein", "@protein").unwrap();

        let traj_single = system_single
            .group_xtc_iter("test_files/short_trajectory.xtc", "Protein")
            .unwrap();
        let traj_cat = system_cat
            .group_traj_cat_iter::<GroupXtcReader>(
                &[
                    "test_files/split/traj1.xtc",
                    "test_files/split/traj2.xtc",
                    "test_files/split/traj3.xtc",
                    "test_files/split/traj4.xtc",
                    "test_files/split/traj5.xtc",
                    "test_files/split/traj6.xtc",
                ],
                "Protein",
            )
            .unwrap();

        for (frame_single, frame_cat) in traj_single.zip(traj_cat) {
            let frame_single = frame_single.unwrap();
            let frame_cat = frame_cat.unwrap();

            assert_approx_eq!(
                f32,
                frame_single.get_simulation_time(),
                frame_cat.get_simulation_time()
            );
            assert_eq!(
                frame_single.get_simulation_step(),
                frame_cat.get_simulation_step()
            );

            compare_box(
                frame_single.get_box().unwrap(),
                frame_cat.get_box().unwrap(),
            );

            for (atom_single, atom_cat) in frame_single.atoms_iter().zip(frame_cat.atoms_iter()) {
                compare_atoms(atom_single, atom_cat);
            }
        }
    }

    #[test]
    #[cfg(feature = "molly")]
    fn cat_xtc_group_steps_with_ranges() {
        use crate::prelude::GroupXtcReader;

        let ranges = vec![(0.0, 570.0), (320.0, f32::MAX), (220.0, 800.0)];

        for step in 2..=11 {
            for (start, end) in &ranges {
                let mut system_single = System::from_file("test_files/example.gro").unwrap();
                system_single.group_create("Protein", "@protein").unwrap();
                let mut system_cat = System::from_file("test_files/example.gro").unwrap();
                system_cat.group_create("Protein", "@protein").unwrap();

                let traj_single = system_single
                    .group_xtc_iter("test_files/short_trajectory.xtc", "Protein")
                    .unwrap()
                    .with_step(step)
                    .unwrap()
                    .with_range(*start, *end)
                    .unwrap();
                let traj_cat = system_cat
                    .group_traj_cat_iter::<GroupXtcReader>(
                        &[
                            "test_files/split/traj1.xtc",
                            "test_files/split/traj2.xtc",
                            "test_files/split/traj3.xtc",
                            "test_files/split/traj4.xtc",
                            "test_files/split/traj5.xtc",
                            "test_files/split/traj6.xtc",
                        ],
                        "Protein",
                    )
                    .unwrap()
                    .with_step(step)
                    .unwrap()
                    .with_range(*start, *end)
                    .unwrap();

                for (frame_single, frame_cat) in traj_single.zip(traj_cat) {
                    let frame_single = frame_single.unwrap();
                    let frame_cat = frame_cat.unwrap();

                    assert_approx_eq!(
                        f32,
                        frame_single.get_simulation_time(),
                        frame_cat.get_simulation_time()
                    );
                    assert_eq!(
                        frame_single.get_simulation_step(),
                        frame_cat.get_simulation_step()
                    );

                    compare_box(
                        frame_single.get_box().unwrap(),
                        frame_cat.get_box().unwrap(),
                    );

                    for (atom_single, atom_cat) in
                        frame_single.atoms_iter().zip(frame_cat.atoms_iter())
                    {
                        compare_atoms(atom_single, atom_cat);
                    }
                }
            }
        }
    }
}
