use crate::prelude::*;

//================================================
/// Read-only subsystem for non-blocking parallel access to atoms and posisitons
/// Doesn't have access to shared fields such as box and bonds.
//================================================
pub struct SelPar<'a> {
    index: &'a [usize],
    sys: *const System,
}

impl<'a> SelPar<'a> {
    pub(crate) fn new(sys: &'a System, index: &'a [usize]) -> SelPar<'a> {
        Self {
            index,
            sys,
        }
    }
}

unsafe impl Sync for SelPar<'_> {}
unsafe impl Send for SelPar<'_> {}

impl IndexSliceProvider for SelPar<'_> {
    fn get_index_slice(&self) -> &[usize] {
        self.index
    }
}

impl SystemProvider for SelPar<'_> {
    fn get_system_ptr(&self) -> *const System {
        self.sys
    }
}

//================================================
/// Read-write subsystem for non-blocking parallel access to atoms and posisitons
/// Doesn't have access to shared fields such as box and bonds.
//================================================
pub struct SelParMut<'a> {
    index: &'a [usize],
    sys: *mut System,
}

impl<'a> SelParMut<'a> {
    pub(crate) fn new(sys: *mut System, index: &'a [usize]) -> SelParMut<'a> {
        Self {
            index,
            sys,
        }
    }
}

unsafe impl Sync for SelParMut<'_> {}
unsafe impl Send for SelParMut<'_> {}

impl IndexSliceProvider for SelParMut<'_> {
    fn get_index_slice(&self) -> &[usize] {
        self.index
    }
}

impl SystemProvider for SelParMut<'_> {
    fn get_system_ptr(&self) -> *const System {
        self.sys as *const System
    }
}

// PosMutProvider and AtomMutProvider are implemented directly (not via SysMutProvider)
// so that shared fields (box, bonds, molecules) cannot be mutated during parallel use.
impl PosMutProvider for SelParMut<'_> {
    unsafe fn coords_ptr_mut(&mut self) -> *mut Pos {
        (*self.sys).st.coords.as_mut_ptr()
    }
}
impl AtomMutProvider for SelParMut<'_> {
    unsafe fn atoms_ptr_mut(&mut self) -> *mut Atom {
        (*self.sys).top.atoms.as_mut_ptr()
    }
}
// VelMutProvider and ForceMutProvider are implemented directly (same reason as above).
// VelProvider and ForceProvider are covered by blanket impls via SystemProvider.
impl VelMutProvider for SelParMut<'_> {
    unsafe fn vel_ptr_mut(&mut self) -> *mut Vel {
        let v = &mut (*self.sys).st.velocities;
        if v.is_empty() { std::ptr::null_mut() } else { v.as_mut_ptr() }
    }
}
impl ForceMutProvider for SelParMut<'_> {
    unsafe fn force_ptr_mut(&mut self) -> *mut Force {
        let v = &mut (*self.sys).st.forces;
        if v.is_empty() { std::ptr::null_mut() } else { v.as_mut_ptr() }
    }
}

//============================================================================
/// Collection of non-overlapping selections that could be mutated in parallel
/// Selections don't have access to shared fields such as box and bonds.
//============================================================================
pub struct ParSplit {
    pub(crate) selections: Vec<Sel>,
    pub(crate) max_index: usize,
}

impl ParSplit {
    pub(crate) fn check_bounds(&self, sys: &System) {
        if self.max_index >= sys.len() {
            panic!("max index of ParSplit is out of bounds");
        }
    }

    pub fn get_bound_mut<'a>(&'a self, sys: &'a mut System, i: usize) -> SelParMut<'a> {
        self.check_bounds(sys);
        SelParMut::new(sys as *mut System, &self.selections[i].0)
    }

    pub fn get_bound<'a>(&'a self, sys: &'a System, i: usize) -> SelPar<'a> {
        self.check_bounds(sys);
        SelPar::new(sys, &self.selections[i].0)
    }

    pub fn into_selections(self) -> Vec<Sel> {
        self.selections
    }
}


#[cfg(test)]
mod tests {
    use rayon::iter::ParallelIterator;

    use crate::prelude::*;

    #[test]
    fn par_unwrap() -> anyhow::Result<()> {
        // Load file
        let mut sys = System::from_file("tests/membr.gro")?;

        // Make a parallel split for each POPG lipid molecule
        let par = sys.split_par(|p| {
            if p.atom.resname == "POPG" {
                // Whenever new distinct result is returned form this closure
                // new selection is created, so each distinct POPG residue
                // becomes a separate selection.
                Some(p.atom.resindex)
            } else {
                // All other atoms are ignored
                None
            }
        })?;

        // Bind split to a system
        sys.iter_par_split_mut(&par) 
            // Run unwrap on each selection in parallel
            .try_for_each(|mut sel| sel.unwrap_simple())?; 
        Ok(())
    }
}
