/*
 * // 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::particle_system_data3::*;
use crate::vector3::Vector3D;
use crate::sph_kernels3::*;
use crate::point_neighbor_searcher3::PointNeighborSearcher3;
use crate::bounding_box3::BoundingBox3D;
use crate::point_generator3::PointGenerator3;
use rayon::prelude::*;
use crate::bcc_lattice_point_generator::BccLatticePointGenerator;
use std::sync::{RwLock, Arc};

///
/// # 3-D SPH particle system data.
///
/// This class extends ParticleSystemData3 to specialize the data model for SPH.
/// It includes density and pressure array as a default particle attribute, and
/// it also contains SPH utilities such as interpolation operator.
///
pub struct SphSystemData3 {
    /// Target density of this particle system in kg/m^3.
    _target_density: f64,

    /// Target spacing of this particle system in meters.
    _target_spacing: f64,

    /// Relative radius of SPH kernel.
    /// SPH kernel radius divided by target spacing.
    _kernel_radius_over_target_spacing: f64,

    /// SPH kernel radius in meters.
    _kernel_radius: f64,

    _pressure_idx: usize,

    _density_idx: usize,

    _base_data: ParticleSystemData3,
}

impl SphSystemData3 {
    /// Constructs empty SPH system.
    pub fn new_default() -> SphSystemData3 {
        return SphSystemData3 {
            _target_density: crate::constants::K_WATER_DENSITY,
            _target_spacing: 0.1,
            _kernel_radius_over_target_spacing: 1.8,
            _kernel_radius: 0.0,
            _pressure_idx: 0,
            _density_idx: 0,
            _base_data: ParticleSystemData3::new(0),
        };
    }

    /// Constructs SPH system data with given number of particles.
    pub fn new(number_of_particles: usize) -> SphSystemData3 {
        let mut data = SphSystemData3 {
            _target_density: crate::constants::K_WATER_DENSITY,
            _target_spacing: 0.1,
            _kernel_radius_over_target_spacing: 1.8,
            _kernel_radius: 0.0,
            _pressure_idx: 0,
            _density_idx: 0,
            _base_data: ParticleSystemData3::new(number_of_particles),
        };

        data._density_idx = data._base_data.add_scalar_data(None);
        data._pressure_idx = data._base_data.add_scalar_data(None);

        data.set_target_spacing(0.1);

        return data;
    }
}

impl SphSystemData3 {
    ///
    /// \brief      Sets the radius.
    ///
    /// Sets the radius of the particle system. The radius will be interpreted
    /// as target spacing.
    ///
    pub fn set_radius(&mut self, new_radius: f64) {
        // Interpret it as setting target spacing
        self.set_target_spacing(new_radius);
    }

    ///
    /// \brief      Sets the mass of a particle.
    ///
    /// Setting the mass of a particle will change the target density.
    ///
    /// - parameter:  new_mass The new mass.
    ///
    pub fn set_mass(&mut self, new_mass: f64) {
        let inc_ratio = new_mass / self._base_data.mass();
        self._target_density *= inc_ratio;
        self._base_data.set_mass(new_mass);
    }

    /// Returns the density array accessor (immutable).
    pub fn densities(&self) -> &Vec<f64> {
        return self._base_data.scalar_data_at(self._density_idx);
    }

    /// Returns the density array accessor (mutable).
    pub fn densities_mut(&mut self) -> &mut Vec<f64> {
        return self._base_data.scalar_data_at_mut(self._density_idx);
    }

    /// Returns the pressure array accessor (immutable).
    pub fn pressures(&self) -> &Vec<f64> {
        return self._base_data.scalar_data_at(self._pressure_idx);
    }

    /// Returns the pressure array accessor (mutable).
    pub fn pressures_mut(&mut self) -> &mut Vec<f64> {
        return self._base_data.scalar_data_at_mut(self._pressure_idx);
    }

    ///
    /// \brief Updates the density array with the latest particle positions.
    ///
    /// This function updates the density array by recalculating each particle's
    /// latest nearby particles' position.
    ///
    /// \warning You must update the neighbor searcher
    /// (SphSystemData3::build_neighbor_searcher) before calling this function.
    ///
    pub fn update_densities(&mut self) {
        let m = self._base_data.mass();

        let mut density: Vec<f64> = self._base_data.positions().into_par_iter().map(|pos| {
            let sum = self.sum_of_kernel_nearby(pos);
            return m * sum;
        }).collect();
        self.densities_mut().clear();
        self.densities_mut().append(&mut density);
    }

    /// Sets the target density of this particle system.
    pub fn set_target_density(&mut self, target_density: f64) {
        self._target_density = target_density;

        self.compute_mass();
    }

    /// Returns the target density of this particle system.
    pub fn target_density(&self) -> f64 {
        return self._target_density;
    }

    ///
    /// \brief Sets the target particle spacing in meters.
    ///
    /// Once this function is called, hash grid and density should be
    /// updated using updateHashGrid() and update_densities).
    ///
    pub fn set_target_spacing(&mut self, spacing: f64) {
        self._base_data.set_radius(spacing);

        self._target_spacing = spacing;
        self._kernel_radius = self._kernel_radius_over_target_spacing * self._target_spacing;

        self.compute_mass();
    }

    /// Returns the target particle spacing in meters.
    pub fn target_spacing(&self) -> f64 {
        return self._target_spacing;
    }

    ///
    /// \brief Sets the relative kernel radius.
    ///
    /// Sets the relative kernel radius compared to the target particle
    /// spacing (i.e. kernel radius / target spacing).
    /// Once this function is called, hash grid and density should
    /// be updated using updateHashGrid() and update_densities).
    ///
    pub fn set_relative_kernel_radius(&mut self, relative_radius: f64) {
        self._kernel_radius_over_target_spacing = relative_radius;
        self._kernel_radius = self._kernel_radius_over_target_spacing * self._target_spacing;

        self.compute_mass();
    }

    ///
    /// \brief Returns the relative kernel radius.
    ///
    /// Returns the relative kernel radius compared to the target particle
    /// spacing (i.e. kernel radius / target spacing).
    ///
    pub fn relative_kernel_radius(&self) -> f64 {
        return self._kernel_radius_over_target_spacing;
    }

    ///
    /// \brief Sets the absolute kernel radius.
    ///
    /// Sets the absolute kernel radius compared to the target particle
    /// spacing (i.e. relative kernel radius * target spacing).
    /// Once this function is called, hash grid and density should
    /// be updated using updateHashGrid() and update_densities).
    ///
    pub fn set_kernel_radius(&mut self, kernel_radius: f64) {
        self._kernel_radius = kernel_radius;
        self._target_spacing = kernel_radius / self._kernel_radius_over_target_spacing;

        self.compute_mass();
    }

    /// Returns the kernel radius in meters unit.
    pub fn kernel_radius(&self) -> f64 {
        return self._kernel_radius;
    }
    /// Returns sum of kernel function evaluation for each nearby particle.
    pub fn sum_of_kernel_nearby(&self, origin: &Vector3D) -> f64 {
        let mut sum = 0.0;
        let kernel = SphStdKernel3::new(self._kernel_radius);
        self._base_data.neighbor_searcher().read().unwrap().for_each_nearby_point(
            origin, self._kernel_radius, &mut |_: usize, neighbor_position: &Vector3D| {
                let dist = origin.distance_to(*neighbor_position);
                sum += kernel.apply(dist);
            });
        return sum;
    }

    ///
    /// \brief Returns interpolated value at given origin point.
    ///
    /// Returns interpolated scalar data from the given position using
    /// standard SPH weighted average. The data array should match the
    /// particle layout. For example, density or pressure arrays can be
    /// used.
    ///
    /// \warning You must update the neighbor searcher
    /// (SphSystemData3::build_neighbor_searcher) before calling this function.
    ///
    pub fn interpolate_scalar(&self, origin: &Vector3D,
                              values: &Vec<f64>) -> f64 {
        let mut sum = 0.0;
        let d = self.densities();
        let kernel = SphStdKernel3::new(self._kernel_radius);
        let m = self._base_data.mass();

        self._base_data.neighbor_searcher().read().unwrap().for_each_nearby_point(
            origin, self._kernel_radius, &mut |i: usize, neighbor_position: &Vector3D| {
                let dist = origin.distance_to(*neighbor_position);
                let weight = m / d[i] * kernel.apply(dist);
                sum += weight * values[i];
            });

        return sum;
    }

    ///
    /// \brief Returns interpolated vector value at given origin point.
    ///
    /// Returns interpolated vector data from the given position using
    /// standard SPH weighted average. The data array should match the
    /// particle layout. For example, velocity or acceleration arrays can be
    /// used.
    ///
    /// \warning You must update the neighbor searcher
    /// (SphSystemData3::build_neighbor_searcher) before calling this function.
    ///
    pub fn interpolate_vec(&self, origin: &Vector3D,
                           values: &Vec<Vector3D>) -> Vector3D {
        let mut sum = Vector3D::new_default();
        let d = self.densities();
        let kernel = SphStdKernel3::new(self._kernel_radius);
        let m = self._base_data.mass();

        self._base_data.neighbor_searcher().read().unwrap().for_each_nearby_point(
            origin, self._kernel_radius, &mut |i: usize, neighbor_position: &Vector3D| {
                let dist = origin.distance_to(*neighbor_position);
                let weight = m / d[i] * kernel.apply(dist);
                sum += values[i] * weight;
            });

        return sum;
    }

    ///
    /// Returns the gradient of the given values at i-th particle.
    ///
    /// \warning You must update the neighbor lists
    /// (SphSystemData3::build_neighbor_lists) before calling this function.
    ///
    pub fn gradient_at(&self, i: usize,
                       values: &Vec<f64>) -> Vector3D {
        let mut sum = Vector3D::new_default();
        let p = self._base_data.positions();
        let d = self.densities();
        let neighbors = &self._base_data.neighbor_lists()[i];
        let origin = p[i];
        let kernel = SphSpikyKernel3::new(self._kernel_radius);
        let m = self._base_data.mass();

        for j in neighbors {
            let neighbor_position = p[*j];
            let dist = origin.distance_to(neighbor_position);
            if dist > 0.0 {
                let dir = (neighbor_position - origin) / dist;
                sum += kernel.gradient_dir(dist, &dir) * d[i] * m * (values[i] / crate::math_utils::square(d[i])
                    + values[*j] / crate::math_utils::square(d[*j]));
            }
        }

        return sum;
    }

    ///
    /// Returns the laplacian of the given values at i-th particle.
    ///
    /// \warning You must update the neighbor lists
    /// (SphSystemData3::build_neighbor_lists) before calling this function.
    ///
    pub fn laplacian_at_scalar(&self, i: usize,
                               values: &Vec<f64>) -> f64 {
        let mut sum = 0.0;
        let p = self._base_data.positions();
        let d = self.densities();
        let neighbors = &self._base_data.neighbor_lists()[i];
        let origin = p[i];
        let kernel = SphSpikyKernel3::new(self._kernel_radius);
        let m = self._base_data.mass();

        for j in neighbors {
            let neighbor_position = p[*j];
            let dist = origin.distance_to(neighbor_position);
            sum += m * (values[*j] - values[i]) / d[*j] * kernel.second_derivative(dist);
        }

        return sum;
    }

    ///
    /// Returns the laplacian of the given values at i-th particle.
    ///
    /// \warning You must update the neighbor lists
    /// (SphSystemData3::build_neighbor_lists) before calling this function.
    ///
    pub fn laplacian_at_vec(&self, i: usize,
                            values: &Vec<Vector3D>) -> Vector3D {
        let mut sum = Vector3D::new_default();
        let p = self._base_data.positions();
        let d = self.densities();
        let neighbors = &self._base_data.neighbor_lists()[i];
        let origin = p[i];
        let kernel = SphSpikyKernel3::new(self._kernel_radius);
        let m = self._base_data.mass();

        for j in neighbors {
            let neighbor_position = p[*j];
            let dist = origin.distance_to(neighbor_position);
            sum += (values[*j] - values[i]) * m / d[*j] * kernel.second_derivative(dist);
        }

        return sum;
    }

    /// Builds neighbor searcher with kernel radius.
    pub fn build_neighbor_searcher(&mut self) {
        self._base_data.build_neighbor_searcher(self._kernel_radius);
    }

    /// Builds neighbor lists with kernel radius.
    pub fn build_neighbor_lists(&mut self) {
        self._base_data.build_neighbor_lists(self._kernel_radius);
    }

    /// Computes the mass based on the target density and spacing.
    fn compute_mass(&mut self) {
        let mut points: Vec<Vector3D> = Vec::new();
        let points_generator = BccLatticePointGenerator {};
        let sample_bound = BoundingBox3D::new(
            Vector3D::new(-1.5 * self._kernel_radius, -1.5 * self._kernel_radius, -1.5 * self._kernel_radius),
            Vector3D::new(1.5 * self._kernel_radius, 1.5 * self._kernel_radius, 1.5 * self._kernel_radius));

        points_generator.generate(&sample_bound, self._target_spacing, &mut points);

        let mut max_number_density = 0.0;
        let kernel = SphStdKernel3::new(self._kernel_radius);

        for i in 0..points.len() {
            let point = points[i];
            let mut sum = 0.0;

            for j in 0..points.len() {
                let neighbor_point = points[j];
                sum += kernel.apply(neighbor_point.distance_to(point));
            }

            max_number_density = f64::max(max_number_density, sum);
        }

        debug_assert!(max_number_density > 0.0);

        let new_mass = self._target_density / max_number_density;
        self._base_data.set_mass(new_mass);
    }
}

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

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

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

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

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

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

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

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

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

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

/// Shared pointer for the SphSystemData3 type.
pub type SphSystemData3Ptr = Arc<RwLock<SphSystemData3>>;