molar 1.3.1

use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator};

use super::utils::check_topology_state_sizes;
use crate::prelude::*;
use std::path::Path;

//================================================
/// System that stores Topology and State
//================================================
#[derive(Debug, Default)]
pub struct System {
    pub(super) top: Topology,
    pub(super) st: State,
}

impl System {
    pub fn new(top: Topology, st: State) -> Result<Self, SelectionError> {
        check_topology_state_sizes(&top, &st)?;
        Ok(Self { top, st })
    }

    pub fn from_file(fname: impl AsRef<Path>) -> Result<Self, SelectionError> {
        let mut fh = FileHandler::open(fname)?;
        let (top, st) = fh.read()?;
        Ok(Self::new(top, st)?)
    }

    /// Create unbound sub-selection
    pub fn sub_select(&self, ind: &impl IndexSliceProvider, def: impl SelectionDef) -> Result<Sel, SelectionError> {
        unsafe {Ok(self.try_bind_sorted_slice(ind.get_index_slice())?.select(def)?)}
    }

    /// Create all detached
    pub fn select_all(&self) -> Sel {
        Sel(unsafe { SVec::from_sorted((0..self.len()).into_iter().collect()) })
    }

    pub fn select_all_bound(&self) -> SelOwnBound<'_> {
        SelOwnBound {
            sys: self,
            index: unsafe { SVec::from_sorted((0..self.len()).into_iter().collect()) },
        }
    }

    pub fn select_bound_mut(
        &mut self,
        def: impl SelectionDef,
    ) -> Result<SelOwnBoundMut<'_>, SelectionError> {
        let index = def.into_sel_index(self, None)?;
        Ok(SelOwnBoundMut { sys: self, index })
    }

    pub fn select_all_bound_mut(&mut self) -> SelOwnBoundMut<'_> {
        let index = unsafe { SVec::from_sorted((0..self.len()).into_iter().collect()) };
        SelOwnBoundMut { sys: self, index }
    }

    // Internal function. Slice is supposed to be sorted!
    unsafe fn try_bind_sorted_slice<'a>(&'a self, sel: &'a [usize]) -> Result<SelBound<'a>, SelectionError> {
        // No need to check for empty index since it's guaranteed to be non-empty
        let last = unsafe { *sel.get_unchecked(sel.len() - 1) };
        if last >= self.top.len() {
            Err(SelectionError::IndexValidation(
                *sel.first().unwrap(),
                last,
                self.top.len() - 1,
            ))
        } else {
            Ok(SelBound {
                sys: self,
                index: sel,
            })
        }
    }

    /// Binds detached selection index to make borrowed selection.
    /// `sel`  is not consumed.
    pub fn bind<'a>(&'a self, sel: &'a Sel) -> SelBound<'a> {
        // No need to check for empty index since it's guaranteed to be non-empty
        let last = unsafe { *sel.0.get_unchecked(sel.0.len() - 1) };
        if last >= self.top.len() {
            panic!("selection is out of bounds");
        } else {
            SelBound {
                sys: self,
                index: sel.0.as_slice(),
            }
        }
    }

    /// Binds detached selection index to make borrowed selection.
    /// `sel`  is not consumed.
    pub fn try_bind<'a>(&'a self, sel: &'a Sel) -> Result<SelBound<'a>, SelectionError> {
        // No need to check for empty index since it's guaranteed to be non-empty
        let last = unsafe { *sel.0.get_unchecked(sel.0.len() - 1) };
        if last >= self.top.len() {
            Err(SelectionError::IndexValidation(
                *sel.0.first().unwrap(),
                last,
                self.top.len() - 1,
            ))
        } else {
            Ok(SelBound {
                sys: self,
                index: sel.0.as_slice(),
            })
        }
    }

    /// Mutably binds detached selection index to make borrowed selection.
    /// `sel`  is not consumed.
    pub fn bind_mut<'a>(&'a mut self, sel: &'a Sel) -> SelBoundMut<'a> {
        // No need to check for empty index since it's guaranteed to be non-empty
        let last = unsafe { *sel.0.get_unchecked(sel.0.len() - 1) };
        if last >= self.top.len() {
            panic!("selection is out of bounds");
        } else {
            SelBoundMut {
                sys: self,
                index: sel.0.as_slice(),
            }
        }
    }

    /// Mutably binds detached selection index to make borrowed selection.
    /// `sel`  is not consumed.
    pub fn try_bind_mut<'a>(&'a mut self, sel: &'a Sel) -> Result<SelBoundMut<'a>, SelectionError> {
        // No need to check for empty index since it's guaranteed to be non-empty
        let last = unsafe { *sel.0.get_unchecked(sel.0.len() - 1) };
        if last >= self.top.len() {
            Err(SelectionError::IndexValidation(
                *sel.0.first().unwrap(),
                last,
                self.top.len() - 1,
            ))
        } else {
            Ok(SelBoundMut {
                sys: self,
                index: sel.0.as_slice(),
            })
        }
    }

    /// Returns mutable parallel iterator over parallel selections.
    pub fn iter_par_split_mut<'a>(
        &'a mut self,
        par: &'a ParSplit,
    ) -> impl IndexedParallelIterator<Item = SelParMut<'a>> {
        par.check_bounds(self);
        par.selections
            .par_iter()
            .map(|sel| SelParMut::new(self, &sel.0))
    }

    /// Returns parallel iterator over parallel selections.
    pub fn iter_par_split<'a>(
        &'a mut self,
        par: &'a ParSplit,
    ) -> impl IndexedParallelIterator<Item = SelPar<'a>> {
        par.check_bounds(self);
        par.selections
            .par_iter()
            .map(|sel| SelPar::new(self, &sel.0))
    }

    pub fn split_bound<'a, RT, F>(&'a self, func: F) -> impl Iterator<Item = SelOwnBound<'a>>
    where
        RT: Default + std::cmp::PartialEq + 'a,
        F: Fn(Particle) -> Option<RT> + 'a,
    {
        self.split(func).map(|sel| SelOwnBound {
            sys: self,
            index: sel.0,
        })
    }

    pub fn split_resindex_bound(&self) -> impl Iterator<Item = SelOwnBound<'_>> {
        self.split_bound(|p| Some(p.atom.resindex))
    }

    pub fn set_state(&mut self, st: State) -> Result<State, SelectionError> {
        if self.len() != st.len() {
            Err(SelectionError::IncompatibleState)
        } else {
            Ok(std::mem::replace(&mut self.st, st))
        }
    }

    pub fn set_topology(&mut self, top: Topology) -> Result<Topology, SelectionError> {
        if self.len() != top.len() {
            Err(SelectionError::IncompatibleTopology)
        } else {
            Ok(std::mem::replace(&mut self.top, top))
        }
    }

    pub fn release(self) -> (Topology, State) {
        (self.top, self.st)
    }

    //===============
    // Modifying
    //===============

    /// Append selection derived from self
    pub fn append_from_self(&mut self, sel: &Sel) -> Result<Sel, SelectionError> {
        let old_last = self.len() - 1;
        let last_ind = unsafe { *sel.0.get_unchecked(sel.0.len() - 1) };
        if last_ind >= self.top.len() {
            return Err(SelectionError::IndexValidation(
                sel.0[0],
                last_ind,
                self.top.len(),
            ))?;
        }
        let pos: Vec<_> = sel.0.iter().map(|i| &self.st.coords[*i]).cloned().collect();
        let atoms: Vec<_> = sel.0.iter().map(|i| &self.top.atoms[*i]).cloned().collect();
        self.st.add_coords(pos.into_iter());
        self.top.add_atoms(atoms.into_iter());
        Ok(Sel::from_iter(old_last + 1..self.len())?)
    }

    pub fn append_atoms_coords<'a>(
        &mut self,
        atoms: impl AtomIterator<'a>,
        coords: impl PosIterator<'a>,
    ) -> Result<Sel, SelectionError> {
        let old_last = self.len() - 1;
        self.st.add_coords(coords.cloned());
        self.top.add_atoms(atoms.cloned());
        check_topology_state_sizes(&self.top, &self.st)?;
        Ok(Sel::from_iter(old_last + 1..self.len())?)
    }

    pub fn append(
        &mut self,
        data: &(impl AtomIterProvider + PosIterProvider),
    ) -> Result<Sel, SelectionError> {
        let old_last = self.len() - 1;
        self.st.add_coords(data.iter_pos().cloned());
        self.top.add_atoms(data.iter_atoms().cloned());
        check_topology_state_sizes(&self.top, &self.st)?;
        Ok(Sel::from_iter(old_last + 1..self.len())?)
    }

    pub fn remove(
        &mut self,
        removed: impl Iterator<Item = usize> + Clone,
    ) -> Result<(), BuilderError> {
        self.st.remove_coords(removed.clone())?;
        self.top.remove_atoms(removed)?;
        Ok(())
    }

    pub fn set_box_from(&mut self, src: &impl BoxProvider) {
        self.st.pbox = src.get_box().cloned();
    }

    pub fn multiply_periodically(&mut self, nbox: [usize; 3]) -> Result<(), SelectionError> {
        if self.get_box().is_none() {
            return Err(PeriodicBoxError::NoPbc)?;
        }
        let m = self.require_box()?.get_matrix();
        let all = self.select_all();
        for x in 0..=nbox[0] {
            for y in 0..=nbox[1] {
                for z in 0..=nbox[2] {
                    if x == 0 && y == 0 && z == 0 {
                        continue;
                    }
                    let added = self.append_from_self(&all)?;
                    let shift =
                        m.column(0) * x as f32 + m.column(1) * y as f32 + m.column(2) * z as f32;
                    self.select_bound_mut(added)?.translate(&shift);
                }
            }
        }
        // Scale the box
        self.get_box_mut().unwrap().scale_vectors([
            nbox[0] as f32,
            nbox[1] as f32,
            nbox[2] as f32,
        ])?;

        // Re-assign resindex
        self.top.assign_resindex();
        Ok(())
    }

    pub fn assign_resindex(&mut self) {
        self.top.assign_resindex();
    }
}

impl Selectable for System {
    fn select(&self, def: impl SelectionDef) -> Result<Sel, SelectionError> {
        Ok(Sel(def.into_sel_index(self, None)?))
    }
}

impl SelectableBound for System {
    fn select_bound(&self, def: impl SelectionDef) -> Result<SelOwnBound<'_>, SelectionError> {
        Ok(SelOwnBound {
            index: def.into_sel_index(self, None)?,
            sys: unsafe{&*self.get_system_ptr()},
        })
    }
}

impl SaveTopology for System {
    fn iter_atoms_dyn(&self) -> Box<dyn Iterator<Item = &Atom> + '_> {
        Box::new(self.iter_atoms())
    }
}

impl SaveState for System {
    fn iter_pos_dyn<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Pos> + 'a> {
        Box::new(self.iter_pos())
    }
}

impl SaveTopologyState for System {}

impl SystemProvider for System {
    fn get_system_ptr(&self) -> *const System {
        self
    }
}

impl SystemMutProvider for System {}

impl LenProvider for System {
    fn len(&self) -> usize {
        self.top.atoms.len()
    }
}

impl IndexProvider for System {
    unsafe fn get_index_unchecked(&self, i: usize) -> usize {
        i
    }

    fn iter_index(&self) -> impl Iterator<Item = usize> {
        (0..self.len()).into_iter()
    }
}

impl IndexParProvider for System {
    fn par_iter_index(&self) -> impl IndexedParallelIterator<Item = usize> {
        (0..self.len()).into_par_iter()
    }
}

// Analog of `bind` with `&sel >> &system` syntax
impl<'a> std::ops::Shr<&'a System> for &'a Sel {
    type Output = SelBound<'a>;
    fn shr(self, rhs: &'a System) -> Self::Output {
        rhs.bind(self)
    }
}

// Analog of `bind_mut` with `&sel >> &mut system` syntax
impl<'a> std::ops::Shr<&'a mut System> for &'a Sel {
    type Output = SelBoundMut<'a>;
    fn shr(self, rhs: &'a mut System) -> Self::Output {
        rhs.bind_mut(self)
    }
}