/*
 * // 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::usize3::USize3;
use crate::isize3::ISize3;
use crate::vector3::Vector3D;
use crate::point_neighbor_searcher3::*;
use std::sync::{RwLock, Arc};
use std::cmp::Ordering;
use rayon::prelude::*;
use rayon::iter::ParallelIterator;
use log::info;

///
/// # Parallel version of hash grid-based 3-D point searcher.
///
/// This class implements parallel version of 3-D point searcher by using hash
/// grid for its internal acceleration data structure. Each point is recorded to
/// its corresponding bucket where the hashing function is 3-D grid mapping.
///
pub struct PointParallelHashGridSearcher3 {
    _grid_spacing: f64,
    _resolution: ISize3,
    _points: Vec<Vector3D>,
    _keys: Vec<usize>,
    _start_index_table: Vec<usize>,
    _end_index_table: Vec<usize>,
    _sorted_indices: Vec<usize>,
}

impl PointParallelHashGridSearcher3 {
    pub fn new_default() -> PointParallelHashGridSearcher3 {
        return PointParallelHashGridSearcher3 {
            _grid_spacing: 1.0,
            _resolution: ISize3::new(1, 1, 1),
            _points: vec![],
            _keys: vec![],
            _start_index_table: vec![],
            _end_index_table: vec![],
            _sorted_indices: vec![],
        };
    }

    ///
    /// # Constructs hash grid with given resolution and grid spacing.
    ///
    /// This constructor takes hash grid resolution and its grid spacing as
    /// its input parameters. The grid spacing must be 3x or greater than
    /// search radius.
    ///
    /// - parameter:  resolution  The resolution.
    /// - parameter:  grid_spacing The grid spacing.
    ///
    pub fn new_vec(resolution: USize3, grid_spacing: f64) -> PointParallelHashGridSearcher3 {
        return PointParallelHashGridSearcher3::new(resolution.x,
                                                   resolution.y,
                                                   resolution.z, grid_spacing);
    }

    ///
    /// # Constructs hash grid with given resolution and grid spacing.
    ///
    /// This constructor takes hash grid resolution and its grid spacing as
    /// its input parameters. The grid spacing must be 3x or greater than
    /// search radius.
    ///
    /// - parameter:  resolution_x The resolution x.
    /// - parameter:  resolution_y The resolution y.
    /// - parameter:  resolution_z The resolution z.
    /// - parameter:  grid_spacing The grid spacing.
    ///
    pub fn new(resolution_x: usize,
               resolution_y: usize,
               resolution_z: usize,
               grid_spacing: f64) -> PointParallelHashGridSearcher3 {
        let resolution = ISize3::new(resolution_x as isize,
                                     resolution_y as isize,
                                     resolution_z as isize);
        return PointParallelHashGridSearcher3 {
            _grid_spacing: grid_spacing,
            _resolution: ISize3::new(resolution_x as isize,
                                     resolution_y as isize,
                                     resolution_z as isize),
            _points: vec![],
            _keys: vec![],
            _start_index_table: vec![0; (resolution.x * resolution.y * resolution.z) as usize],
            _end_index_table: vec![0; (resolution.x * resolution.y * resolution.z) as usize],
            _sorted_indices: vec![],
        };
    }

    /// Returns builder fox PointHashGridSearcher3.
    pub fn builder() -> Builder {
        return Builder::new();
    }

    ///
    /// #  Creates a new instance of the object with same properties than original.
    ///
    /// - return     Copy of this object.
    ///
    pub fn clone(&self) -> PointParallelHashGridSearcher3Ptr {
        let mut searcher = PointParallelHashGridSearcher3::new_default();
        searcher.set(self);
        return PointParallelHashGridSearcher3Ptr::new(RwLock::new(searcher));
    }

    /// Copy from the other instance.
    pub fn set(&mut self, other: &PointParallelHashGridSearcher3) {
        self._grid_spacing = other._grid_spacing;
        self._resolution = other._resolution;
        self._points = other._points.clone();
        self._keys = other._keys.clone();
        self._start_index_table = other._start_index_table.clone();
        self._end_index_table = other._end_index_table.clone();
        self._sorted_indices = other._sorted_indices.clone();
    }
}

impl PointParallelHashGridSearcher3 {
    ///
    /// # Returns the hash key list.
    ///
    /// The hash key list maps sorted point index i to its hash key value.
    /// The sorting order is based on the key value itself.
    ///
    /// \return     The hash key list.
    ///
    pub fn keys(&self) -> &Vec<usize> {
        return &self._keys;
    }

    ///
    /// # Returns the start index table.
    ///
    /// The start index table maps the hash grid bucket index to starting index
    /// of the sorted point list. Assume the hash key list looks like:
    ///
    /// \code
    /// [5|8|8|10|10|10]
    /// \endcode
    ///
    /// Then start_index_table and end_index_table should be like:
    ///
    /// \code
    /// [.....|0|...|1|..|3|..]
    /// [.....|1|...|3|..|6|..]
    ///       ^5    ^8   ^10
    /// \endcode
    ///
    /// So that end_index_table(i) - start_index_table(i) is the number points
    /// in i-th table bucket.
    ///
    /// \return     The start index table.
    ///
    pub fn start_index_table(&self) -> &Vec<usize> {
        return &self._start_index_table;
    }

    ///
    /// # Returns the end index table.
    ///
    /// The end index table maps the hash grid bucket index to starting index
    /// of the sorted point list. Assume the hash key list looks like:
    ///
    /// \code
    /// [5|8|8|10|10|10]
    /// \endcode
    ///
    /// Then start_index_table and end_index_table should be like:
    ///
    /// \code
    /// [.....|0|...|1|..|3|..]
    /// [.....|1|...|3|..|6|..]
    ///       ^5    ^8   ^10
    /// \endcode
    ///
    /// So that end_index_table(i) - start_index_table(i) is the number points
    /// in i-th table bucket.
    ///
    /// \return     The end index table.
    ///
    pub fn end_index_table(&self) -> &Vec<usize> {
        return &self._end_index_table;
    }

    ///
    /// # Returns the sorted indices of the points.
    ///
    /// When the hash grid is built, it sorts the points in hash key order. But
    /// rather than sorting the original points, this class keeps the shuffled
    /// indices of the points. The list this function returns maps sorted index
    /// i to original index j.
    ///
    /// \return     The sorted indices of the points.
    ///
    pub fn sorted_indices(&self) -> &Vec<usize> {
        return &self._sorted_indices;
    }

    ///
    /// Returns the hash value for given 3-D bucket index.
    ///
    /// - parameter:  bucket_index The bucket index.
    ///
    /// \return     The hash key from bucket index.
    ///
    pub fn get_hash_key_from_bucket_index(&self, bucket_index: &ISize3) -> usize {
        let mut wrapped_index = *bucket_index;
        wrapped_index.x = bucket_index.x % self._resolution.x;
        wrapped_index.y = bucket_index.y % self._resolution.y;
        wrapped_index.z = bucket_index.z % self._resolution.z;
        if wrapped_index.x < 0 {
            wrapped_index.x += self._resolution.x;
        }
        if wrapped_index.y < 0 {
            wrapped_index.y += self._resolution.y;
        }
        if wrapped_index.z < 0 {
            wrapped_index.z += self._resolution.z;
        }
        return ((wrapped_index.z * self._resolution.y + wrapped_index.y) * self._resolution.x
            + wrapped_index.x) as usize;
    }

    ///
    /// Gets the bucket index from a point.
    ///
    /// - parameter:  position The position of the point.
    ///
    /// - return     The bucket index.
    ///
    pub fn get_bucket_index(&self, position: &Vector3D) -> ISize3 {
        let mut bucket_index = ISize3::new_default();
        bucket_index.x = f64::floor(position.x / self._grid_spacing) as isize;
        bucket_index.y = f64::floor(position.y / self._grid_spacing) as isize;
        bucket_index.z = f64::floor(position.z / self._grid_spacing) as isize;
        return bucket_index;
    }

    pub fn get_hash_key_from_position(&self, position: &Vector3D) -> usize {
        let bucket_index = self.get_bucket_index(position);

        return self.get_hash_key_from_bucket_index(&bucket_index);
    }

    pub fn get_nearby_keys(&self, position: &Vector3D, nearby_keys: &mut [usize; 8]) {
        let origin_index = self.get_bucket_index(position);
        let mut nearby_bucket_indices = [ISize3::new_default(); 8];

        for i in 0..8 {
            nearby_bucket_indices[i] = origin_index;
        }

        if (origin_index.x as f64 + 0.5) * self._grid_spacing <= position.x {
            nearby_bucket_indices[4].x += 1;
            nearby_bucket_indices[5].x += 1;
            nearby_bucket_indices[6].x += 1;
            nearby_bucket_indices[7].x += 1;
        } else {
            nearby_bucket_indices[4].x -= 1;
            nearby_bucket_indices[5].x -= 1;
            nearby_bucket_indices[6].x -= 1;
            nearby_bucket_indices[7].x -= 1;
        }

        if (origin_index.y as f64 + 0.5) * self._grid_spacing <= position.y {
            nearby_bucket_indices[2].y += 1;
            nearby_bucket_indices[3].y += 1;
            nearby_bucket_indices[6].y += 1;
            nearby_bucket_indices[7].y += 1;
        } else {
            nearby_bucket_indices[2].y -= 1;
            nearby_bucket_indices[3].y -= 1;
            nearby_bucket_indices[6].y -= 1;
            nearby_bucket_indices[7].y -= 1;
        }

        if (origin_index.z as f64 + 0.5) * self._grid_spacing <= position.z {
            nearby_bucket_indices[1].z += 1;
            nearby_bucket_indices[3].z += 1;
            nearby_bucket_indices[5].z += 1;
            nearby_bucket_indices[7].z += 1;
        } else {
            nearby_bucket_indices[1].z -= 1;
            nearby_bucket_indices[3].z -= 1;
            nearby_bucket_indices[5].z -= 1;
            nearby_bucket_indices[7].z -= 1;
        }

        for i in 0..8 {
            nearby_keys[i] = self.get_hash_key_from_bucket_index(&nearby_bucket_indices[i]);
        }
    }
}

impl PointNeighborSearcher3 for PointParallelHashGridSearcher3 {
    fn type_name() -> String {
        return "PointParallelHashGridSearcher3".parse().unwrap();
    }

    fn build(&mut self, points: &Vec<Vector3D>) {
        self._points.clear();
        self._keys.clear();
        self._start_index_table.clear();
        self._end_index_table.clear();
        self._sorted_indices.clear();

        // Allocate memory chunks
        let number_of_points = points.len();
        let mut temp_keys: Vec<usize> = Vec::new();
        temp_keys.resize(number_of_points, 0);
        self._start_index_table.resize((self._resolution.x * self._resolution.y * self._resolution.z) as usize, usize::MAX);
        self._end_index_table.resize((self._resolution.x * self._resolution.y * self._resolution.z) as usize, usize::MAX);
        self._keys.resize(number_of_points, 0);
        self._sorted_indices.resize(number_of_points, 0);
        self._points.resize(number_of_points, Vector3D::new_default());

        if number_of_points == 0 {
            return;
        }

        (&mut self._sorted_indices, &mut self._points, 0..number_of_points).into_par_iter().for_each(|(x, y, index)| {
            *x = index;
            *y = points[index];
        });

        (&mut temp_keys, 0..number_of_points).into_par_iter().for_each(|(x, index)| {
            *x = self.get_hash_key_from_position(&points[index]);
        });

        self._sorted_indices.par_sort_by(|index_a: &usize, index_b: &usize| {
            match temp_keys[*index_a] < temp_keys[*index_b] {
                true => Ordering::Less,
                false => Ordering::Greater
            }
        });

        (&mut self._points, &mut self._keys, &self._sorted_indices).into_par_iter().for_each(|(x, y, z)| {
            *x = points[*z];
            *y = temp_keys[*z];
        });

        // Now _points and _keys are sorted by points' hash key values.
        // Let's fill in start/end index table with _keys.

        // Assume that _keys array looks like:
        // [5|8|8|10|10|10]
        // Then _startIndexTable and _endIndexTable should be like:
        // [.....|0|...|1|..|3|..]
        // [.....|1|...|3|..|6|..]
        //       ^5    ^8   ^10
        // So that _endIndexTable[i] - _startIndexTable[i] is the number points
        // in i-th table bucket.

        self._start_index_table[self._keys[0]] = 0;
        self._end_index_table[self._keys[number_of_points - 1]] = number_of_points;

        (1..number_of_points).for_each(|i| {
            if self._keys[i] > self._keys[i - 1] {
                self._start_index_table[self._keys[i]] = i;
                self._end_index_table[self._keys[i - 1]] = i;
            }
        });

        let mut sum_number_of_points_per_bucket = 0;
        let mut max_number_of_points_per_bucket = 0;
        let mut number_of_non_empty_bucket = 0;
        for i in 0..self._start_index_table.len() {
            if self._start_index_table[i] != usize::MAX {
                let number_of_points_in_bucket = self._end_index_table[i] - self._start_index_table[i];
                sum_number_of_points_per_bucket += number_of_points_in_bucket;
                max_number_of_points_per_bucket = usize::max(max_number_of_points_per_bucket, number_of_points_in_bucket);
                number_of_non_empty_bucket += 1;
            }
        }

        info!("Average number of points per non-empty bucket: {}", sum_number_of_points_per_bucket as f64 / number_of_non_empty_bucket as f64);
        info!("Max number of points per bucket: {}", max_number_of_points_per_bucket);
    }

    fn for_each_nearby_point<Callback>(&self, origin: &Vector3D, radius: f64, callback: &mut Callback)
        where Callback: ForEachNearbyPointFunc {
        let mut nearby_keys: [usize; 8] = [0; 8];
        self.get_nearby_keys(origin, &mut nearby_keys);

        let query_radius_squared = radius * radius;

        for i in 0..8 {
            let nearby_key = nearby_keys[i];
            let start = self._start_index_table[nearby_key];
            let end = self._end_index_table[nearby_key];

            // Empty bucket -- continue to next bucket
            if start == usize::MAX {
                continue;
            }

            for j in start..end {
                let direction = self._points[j] - *origin;
                let distance_squared = direction.length_squared();
                if distance_squared <= query_radius_squared {
                    callback(self._sorted_indices[j], &self._points[j]);
                }
            }
        }
    }

    fn has_nearby_point(&self, origin: &Vector3D, radius: f64) -> bool {
        let mut nearby_keys: [usize; 8] = [0; 8];
        self.get_nearby_keys(origin, &mut nearby_keys);

        let query_radius_squared = radius * radius;

        for i in 0..8 {
            let nearby_key = nearby_keys[i];
            let start = self._start_index_table[nearby_key];
            let end = self._end_index_table[nearby_key];

            // Empty bucket -- continue to next bucket
            if start == usize::MAX {
                continue;
            }

            for j in start..end {
                let direction = self._points[j] - *origin;
                let distance_squared = direction.length_squared();
                if distance_squared <= query_radius_squared {
                    return true;
                }
            }
        }

        return false;
    }
}

/// Shared pointer for the PointParallelHashGridSearcher3 type.
pub type PointParallelHashGridSearcher3Ptr = Arc<RwLock<PointParallelHashGridSearcher3>>;


///
/// # Front-end to create PointParallelHashGridSearcher3 objects step by step.
///
pub struct Builder {
    _resolution: USize3,
    _grid_spacing: f64,
}

impl Builder {
    /// Returns builder with resolution.
    pub fn with_resolution(&mut self, resolution: USize3) -> &mut Self {
        self._resolution = resolution;
        return self;
    }

    /// Returns builder with grid spacing.
    pub fn with_grid_spacing(&mut self, grid_spacing: f64) -> &mut Self {
        self._grid_spacing = grid_spacing;
        return self;
    }

    /// Builds PointParallelHashGridSearcher3 instance.
    pub fn build(&mut self) -> PointParallelHashGridSearcher3 {
        return PointParallelHashGridSearcher3::new_vec(self._resolution, self._grid_spacing);
    }

    /// Builds shared pointer of PointParallelHashGridSearcher3 instance.
    pub fn make_shared(&mut self) -> PointParallelHashGridSearcher3Ptr {
        return PointParallelHashGridSearcher3Ptr::new(RwLock::new(self.build()));
    }

    /// constructor
    pub fn new() -> Builder {
        return Builder {
            _resolution: USize3::new(64, 64, 64),
            _grid_spacing: 1.0,
        };
    }
}