/*
 * // Copyright (c) 2021 Feng Yang
 * //
 * // I am making my contributions/submissions to this project solely in my
 * // personal capacity and am not conveying any rights to any intellectual
 * // property of any third parties.
 */

use crate::vector3::Vector3D;
use crate::point_neighbor_searcher3::*;
use crate::point_parallel_hash_grid_searcher3::*;
use crate::usize3::USize3;
use std::sync::{RwLock, Arc};

///
/// # 3-D particle system data.
///
/// This class is the key data structure for storing particle system data. A
/// single particle has position, velocity, and force attributes by default. But
/// it can also have additional custom scalar or vector attributes.
///
pub struct ParticleSystemData3 {
    _radius: f64,
    _mass: f64,
    _number_of_particles: usize,
    _position_idx: usize,
    _velocity_idx: usize,
    _force_idx: usize,

    _scalar_data_list: Vec<Vec<f64>>,
    _vector_data_list: Vec<Vec<Vector3D>>,

    _neighbor_searcher: PointParallelHashGridSearcher3Ptr,
    _neighbor_lists: Vec<Vec<usize>>,
}

impl ParticleSystemData3 {
    /// Default constructor.
    pub fn new_default() -> ParticleSystemData3 {
        return ParticleSystemData3 {
            _radius: 1.0e-3,
            _mass: 1.0e-3,
            _number_of_particles: 0,
            _position_idx: 0,
            _velocity_idx: 0,
            _force_idx: 0,
            _scalar_data_list: vec![],
            _vector_data_list: vec![],
            _neighbor_searcher: PointParallelHashGridSearcher3::builder()
                .with_resolution(USize3::new(64, 64, 64))
                .with_grid_spacing(2.0 * 1.0e-3)
                .make_shared(),
            _neighbor_lists: vec![],
        };
    }

    /// Constructs particle system data with given number of particles.
    pub fn new(number_of_particles: usize) -> ParticleSystemData3 {
        let mut system = ParticleSystemData3::new_default();
        system._position_idx = system.add_vector_data(None);
        system._velocity_idx = system.add_vector_data(None);
        system._force_idx = system.add_vector_data(None);
        system.resize(number_of_particles);
        return system;
    }
}

impl ParticleSystemData3 {
    ///
    /// \brief      Resizes the number of particles of the container.
    ///
    /// This function will resize internal containers to store newly given
    /// number of particles including custom data layers. However, this will
    /// invalidate neighbor searcher and neighbor lists. It is users
    /// responsibility to call ParticleSystemData3::build_neighbor_searcher and
    /// ParticleSystemData3::build_neighbor_lists to refresh those data.
    ///
    /// - parameter:   new_number_of_particles    New number of particles.
    ///
    pub fn resize(&mut self, new_number_of_particles: usize) {
        self._number_of_particles = new_number_of_particles;

        for attr in &mut self._scalar_data_list {
            attr.resize(new_number_of_particles, 0.0);
        }

        for attr in &mut self._vector_data_list {
            attr.resize(new_number_of_particles, Vector3D::new_default());
        }
    }

    /// Returns the number of particles.
    pub fn number_of_particles(&self) -> usize {
        return self._number_of_particles;
    }

    ///
    /// \brief      Adds a scalar data layer and returns its index.
    ///
    /// This function adds a new scalar data layer to the system. It can be used
    /// for adding a scalar attribute, such as temperature, to the particles.
    ///
    /// - parameter:  initial_val  Initial value of the new scalar data.
    ///
    pub fn add_scalar_data(&mut self, initial_val: Option<f64>) -> usize {
        let attr_idx = self._scalar_data_list.len();
        let mut data: Vec<f64> = Vec::new();
        data.resize(self.number_of_particles(), initial_val.unwrap_or(0.0));
        self._scalar_data_list.push(data);
        return attr_idx;
    }

    ///
    /// \brief      Adds a vector data layer and returns its index.
    ///
    /// This function adds a new vector data layer to the system. It can be used
    /// for adding a vector attribute, such as vortex, to the particles.
    ///
    /// - parameter:  initial_val  Initial value of the new vector data.
    ///
    pub fn add_vector_data(&mut self, initial_val: Option<Vector3D>) -> usize {
        let attr_idx = self._vector_data_list.len();
        let mut data: Vec<Vector3D> = Vec::new();
        data.resize(self.number_of_particles(), initial_val.unwrap_or(Vector3D::new_default()));
        self._vector_data_list.push(data);
        return attr_idx;
    }

    /// Returns the radius of the particles.
    pub fn radius(&self) -> f64 {
        return self._radius;
    }

    /// Sets the radius of the particles.
    pub fn set_radius(&mut self, new_radius: f64) {
        self._radius = f64::max(new_radius, 0.0);
    }

    /// Returns the mass of the particles.
    pub fn mass(&self) -> f64 {
        return self._mass;
    }

    /// Sets the mass of the particles.
    pub fn set_mass(&mut self, new_mass: f64) {
        self._mass = f64::max(new_mass, 0.0);
    }

    /// Returns the position array (immutable).
    pub fn positions(&self) -> &Vec<Vector3D> {
        return self.vector_data_at(self._position_idx);
    }

    /// Returns the position array (mutable).
    pub fn positions_mut(&mut self) -> &mut Vec<Vector3D> {
        return self.vector_data_at_mut(self._position_idx);
    }

    /// Returns the velocity array (immutable).
    pub fn velocities(&self) -> &Vec<Vector3D> {
        return self.vector_data_at(self._velocity_idx);
    }

    /// Returns the velocity array (mutable).
    pub fn velocities_mut(&mut self) -> &mut Vec<Vector3D> {
        return self.vector_data_at_mut(self._velocity_idx);
    }

    /// Returns the force array (immutable).
    pub fn forces(&self) -> &Vec<Vector3D> {
        return self.vector_data_at(self._force_idx);
    }

    /// Returns the force array (mutable).
    pub fn forces_mut(&mut self) -> &mut Vec<Vector3D> {
        return self.vector_data_at_mut(self._force_idx);
    }

    /// Returns custom scalar data layer at given index (immutable).
    pub fn scalar_data_at(&self, idx: usize) -> &Vec<f64> {
        return &self._scalar_data_list[idx];
    }

    /// Returns custom scalar data layer at given index (mutable).
    pub fn scalar_data_at_mut(&mut self, idx: usize) -> &mut Vec<f64> {
        return &mut self._scalar_data_list[idx];
    }

    /// Returns custom vector data layer at given index (immutable).
    pub fn vector_data_at(&self, idx: usize) -> &Vec<Vector3D> {
        return &self._vector_data_list[idx];
    }

    /// Returns custom vector data layer at given index (mutable).
    pub fn vector_data_at_mut(&mut self, idx: usize) -> &mut Vec<Vector3D> {
        return &mut self._vector_data_list[idx];
    }

    ///
    /// \brief      Adds a particle to the data structure.
    ///
    /// This function will add a single particle to the data structure. For
    /// custom data layers, zeros will be assigned for new particles.
    /// However, this will invalidate neighbor searcher and neighbor lists. It
    /// is users responsibility to call
    /// ParticleSystemData3::build_neighbor_searcher and
    /// ParticleSystemData3::build_neighbor_lists to refresh those data.
    ///
    /// - parameter:   new_position The new position.
    /// - parameter:   new_velocity The new velocity.
    /// - parameter:   new_force    The new force.
    ///
    pub fn add_particle(&mut self,
                        new_position: &Vector3D,
                        new_velocity: Option<Vector3D>,
                        new_force: Option<Vector3D>) {
        let new_positions = vec![*new_position];
        let new_velocities = vec![new_velocity.unwrap_or(Vector3D::new_default())];
        let new_forces = vec![new_force.unwrap_or(Vector3D::new_default())];

        self.add_particles(&new_positions,
                           Some(&new_velocities),
                           Some(&new_forces));
    }

    ///
    /// \brief      Adds particles to the data structure.
    ///
    /// This function will add particles to the data structure. For custom data
    /// layers, zeros will be assigned for new particles. However, this will
    /// invalidate neighbor searcher and neighbor lists. It is users
    /// responsibility to call ParticleSystemData3::build_neighbor_searcher and
    /// ParticleSystemData3::build_neighbor_lists to refresh those data.
    ///
    /// - parameter:   new_positions  The new positions.
    /// - parameter:   new_velocities The new velocities.
    /// - parameter:   new_forces     The new forces.
    ///
    pub fn add_particles(&mut self,
                         new_positions: &Vec<Vector3D>,
                         new_velocities: Option<&Vec<Vector3D>>,
                         new_forces: Option<&Vec<Vector3D>>) {
        let old_number_of_particles = self.number_of_particles();
        let new_number_of_particles = old_number_of_particles + new_positions.len();

        self.resize(new_number_of_particles);

        let pos = self.positions_mut();
        pos.append(&mut new_positions.clone());
        let vel = self.velocities_mut();
        vel.append(&mut new_velocities.unwrap_or(&vec![]).clone());
        let frc = self.forces_mut();
        frc.append(&mut new_forces.unwrap_or(&vec![]).clone());
    }

    ///
    /// \brief      Returns neighbor searcher.
    ///
    /// This function returns currently set neighbor searcher object. By
    /// default, PointParallelHashGridSearcher2 is used.
    ///
    /// \return     Current neighbor searcher.
    ///
    pub fn neighbor_searcher(&self) -> PointParallelHashGridSearcher3Ptr {
        return self._neighbor_searcher.clone();
    }

    ///
    /// \brief      Returns neighbor lists.
    ///
    /// This function returns neighbor lists which is available after calling
    /// PointParallelHashGridSearcher3::build_neighbor_lists. Each list stores
    /// indices of the neighbors.
    ///
    /// \return     Neighbor lists.
    ///
    pub fn neighbor_lists(&self) -> &Vec<Vec<usize>> {
        return &self._neighbor_lists;
    }

    /// Builds neighbor searcher with given search radius.
    pub fn build_neighbor_searcher(&mut self, max_search_radius: f64) {
        // Use PointParallelHashGridSearcher3 by default
        self._neighbor_searcher = PointParallelHashGridSearcher3::builder()
            .with_resolution(USize3::new(64, 64, 64))
            .with_grid_spacing(2.0 * max_search_radius)
            .make_shared();

        self._neighbor_searcher.write().unwrap().build(self.positions());
    }

    /// Builds neighbor lists with given search radius.
    pub fn build_neighbor_lists(&mut self, max_search_radius: f64) {
        self._neighbor_lists.resize(self.number_of_particles(), vec![]);

        for i in 0..self.number_of_particles() {
            let origin = self.positions()[i];
            self._neighbor_lists[i].clear();

            self._neighbor_searcher.clone().read().unwrap().for_each_nearby_point(
                &origin, max_search_radius,
                &mut |j: usize, _: &Vector3D| {
                    if i != j {
                        self._neighbor_lists[i].push(j);
                    }
                });
        }
    }
}

/// Shared pointer type of ParticleSystemData3.
pub type ParticleSystemData3Ptr = Arc<RwLock<ParticleSystemData3>>;