use std::{cmp::Ordering, mem, ops::{Index, IndexMut}};

use num_traits::Zero;

use crate::{get_bounds, mmheap::MmHeap, KdPoint, KdRegion, WalkDecision};

/// Represents an inner / outer limit for a query within the tree ([`KdTree::k_closest`]).

/// See [`QueryOptions`] for more info

pub enum QueryBound<'a, R: KdRegion> {

    SqDist(<R::Point as KdPoint>::Distance),

    Region(&'a R),

    None

}

/// Fine grained control over a query within the tree ([`KdTree::k_closest`]).

pub struct QueryOptions<'a, R: KdRegion> {

    /// Points within this bound will be EXCLUDED from the query result.

    /// Points ON the inner boundary will also be EXCLUDED (ie the excluded region is a closed set).

    /// Useful to exclude points that are "too close" to the query point,

    /// for example when iteratively increasing a sampling radius in a shortest path algorithm

    /// or some other situation where the number of k_closest points needed is not know a prior.

    /// For `QueryBound::SqDist(d)`, this means points whose sqdist from the query point is <= d will be excluded.

    /// Since the query point itself is included by default if it is in the tree, passing `QueryBound::SqDist(Distance::zero())`

    /// here can be used to exclude it.

    /// For `QueryBound::Region(&r)`, this means points within `r` (where `r.min_sqdist` is `Distance::zero()`) will excluded.

    /// NB this cannot be conservative like an unbounded search.

    /// That is, if `R::min_sqdist` can be an underestimate, then `QueryBound::Region` should not be used.

    /// For `QueryBound::None`, no points will be excluded for being "too close" to the query point.

    pub inner_bound: QueryBound<'a, R>,

    /// Points outside this bound will be EXCLUDED from the query result.

    /// Points ON the outer boundary will be INCLUDED (ie the excluded region is an open set).

    /// This makes the overall included region a half open set, excluding the inner bound but including

    /// the outer bound.  This is to make it easy to have successive queries cover additional area without duplicate results.

    /// For `QueryBound::SqDist(d)`, points whose sqdist from the query point is > d will be excluded.

    /// For `QueryBount::Region(&r)`, points outside `r` (where `r.min_sqdist` is not `Distance::zero()`) will be excluded.

    /// Unlike the inner bound, this will work correctly with a conservative `R::min_sqdist` (ie if min_sqdist is an underestimate

    /// it can only cause the result to include extra points, not miss points that should be included).

    /// For `QueryBound::None`, no points will be excluded for being "too far" from the query point.

    pub outer_bound: QueryBound<'a, R>,

    /// If false, if multiple points are tied for being the kth closest, which one will be returned is arbitrary,

    /// but exactly k points will be returned.

    /// If true, if multiple points are tied for being the kth closest, all of them will be returned,

    /// and more than k points will be returned in this case.

    /// Setting this to true is necessary to correctly call [`KdTree::k_closest`] with an iteratively increasing sampling

    /// radius.

    pub keep_ties: bool

}

impl<'a, R: KdRegion> QueryOptions<'a, R> {

    /// Default `QueryOptions` to have the included region be the entire tree,

    /// and return all points which are tied for being the kth-closest

    pub const ALL_KEEP_TIES: Self = Self{inner_bound: QueryBound::None, outer_bound: QueryBound::None, keep_ties: true};

    /// Default `QueryOptions` to have the included region be the entire tree,

    /// and arbitrarily break ties for the kth-closest point so no more than k points are ever returned.

    pub const ALL_NO_TIES: Self = Self{inner_bound: QueryBound::None, outer_bound: QueryBound::None, keep_ties: false};

    /// Returns true if the included region contains the point `pt`, where `point` is the center of the query,

    /// or false if `pt` is in the excluded region.

        }

        }

    /// Returns true if the included region might overlap with `bounds` and

    /// might contain additional points within `max_sqdist` of the center of the query, `point`.

    /// When some of the required functions of `<R as KdRegion>` are implemented conservatively,

    /// this function will return true in cases where the query region and `bounds` can't

    /// actually contain any additional points in the result set.

    /// However, when this function returns false, the query region will always be disjoint from `bounds`,

    /// UNLESS `self.inner_bound` is not `QueryBound::None` and a conservative implementation is used for [`KdRegion`].

            }

        }

        }

}

/// A KD tree represents a collection of points in space, with the ability to answer distance related queries, mainly:

/// - What are the k closest points to a given point?  (The first k points in the tree, ordered by increasing squared

///   distance from the given point, including the given point if it is in the tree).  (Sometimes called a nearest

///   neighbors query)

/// - What are all the points within a distance d of a given point?  (All points whose squared distance from a given

///   point is at most d^2).  (Sometimes called a ball query)

/// - What are all the points within a given region?  (Sometimes called a cuboid query)

/// This implementation uses an implicit tree, meaning all points are stored in one contiguous buffer and no dereferences

/// are needed to traverse the tree.  This is good for lookup performance, but unfortunately it means that adding/removing

/// points can't be done currently without rebuilding the tree.

/// Besides the three basic distance related queries, KD trees can be used to some extent to help with more complicated

/// distance related queries, like finding the closest pairs of points.

#[derive(Debug)]

pub struct KdTree<R: KdRegion, V = ()> {

    pub bounds: Option<R>,

    points: Box<[(R::Point, V)]>

}

impl<R: KdRegion, V> KdTree<R, V> {

    /// Get the number of points in the kdtree

	/// Iterate over all point, value pairs in the tree in depth first order,

	/// calling a visitor function on each.  The visitor function gets a const reference

	/// to the point, the value, and the bounds of the subtree corresponding to the point,

	/// and may return a [`WalkDecision`] to instruct the traversal to skip the

	/// subtree or to stop the traversal entirely.

        }

    /// Iterate over all point, value pairs in the tree in depth first order,

    /// calling a visitor function on each.  The visitor function gets a mutable reference

    /// to the value, but a const reference

	/// to the point and the bounds of the subtree corresponding to the point,

	/// and may return a [`WalkDecision`] to instruct the traversal to skip the

	/// subtree or to stop the traversal entirely.

    pub fn walk_mut<'a>(&'a mut self, visitor: &mut impl FnMut(&R, &'a R::Point, &mut V) -> WalkDecision) {

        let Some(bounds) = self.bounds.clone() else { return };

        let mut todo = vec![(bounds, 0, self.len(), 0)];

        while let Some((bounds, a, b, layer)) = todo.pop() {

            let mid_idx = (a + b)/2;

            // safe because a < b <= self.points.len() so mid_idx < self.points.len()

            let point = &self.points[mid_idx].0;

            let value = &mut self.points[mid_idx].1;

            match visitor(&bounds, point, value) {

                WalkDecision::Stop => return,

                WalkDecision::SkipChildren => continue,

                WalkDecision::Continue => ()

            }

            let (sub0, sub1) = bounds.split(point, layer);

            if a < mid_idx { todo.push((sub0, a, mid_idx, layer + 1)) }

            if mid_idx + 1 < b { todo.push((sub1, mid_idx + 1, b, layer + 1)) }

        }

    }

	/// Return the k points in the tree which are the closest to a given point.

    /// Behavior can be fine-tuned using `cfg`:

    /// - Can restrict the result set to only include points within a certain sqdist of the query point

    ///   or withn a certain region, as well as to exclude points within another smaller sqdist or subregion.

    /// - Can specify how any points tied for kth closest should be handled (keep all or keep one arbitrarily).

	/// If there are fewer than k points in the tree, returns all the points.

    /// If k is 0, returns an empty minmax heap without looking at the tree

                } else {

                        Ordering::Less => {

                                }

                        }

                    }

                }

    /// Borrowing iterator over only references to points.

    /// The order of the result is arbitrary, but all points will be visited exactly once.

    /// Consuming iterator over only points, which are moved out of self.

    /// The order of the result is arbitrarily, but all points will be visited exactly once.

    pub fn into_points(self) -> impl Iterator<Item=R::Point> {

        self.into_iter().map(|(p,_)|p)

    }

    /// Borrowing iterator over only references to values.

    /// The order of the result is arbitrary, but the value for each point will be visited exactly once.

    pub fn iter_values(&self) -> impl Iterator<Item=&V> + '_ {

        self.into_iter().map(|(_,v)|v)

    }

    /// Mutable borrowing iterator over only references to values.

    /// The order of the result is arbitrary, but the value for each point will be visited exactly once.

    /// This is a splitting borrow, so it is safe to store a reference to some value and mutate it even

    /// after subsequent values have been visited.

    pub fn mut_values(&mut self) -> impl Iterator<Item=&mut V> + '_ {

        self.into_iter().map(|(_,v)|v)

    }

    /// Consuming iterator over only values, which are moved out of self.

    /// The order of the result is arbitrary, but the value for each point will be visited exactly once.

    pub fn into_values(self) -> impl Iterator<Item=V> {

        self.into_iter().map(|(_,v)|v)

    }

    /// Return a reference to the value for some point in the tree, or None if the point is not found

    pub fn get(&self, point: &R::Point) -> Option<&V> {

        self.points.get(self.find(point)).map(|(_,v)|v)

    }

    /// Return a reference to a point in the tree and the corresponding value, or None if the point is not found

    pub fn get_point_value(&self, point: &R::Point) -> Option<(&R::Point, &V)> {

        self.points.get(self.find(point)).map(|(p,v)|(p,v))

    }

    /// Return true if the tree contains some point or false otherwise

    pub fn contains_point(&self, point: &R::Point) -> bool {

        self.find(point) != self.len()

    }

    /// Return a mutable reference to the value for some point in the tree, or None if the point is not found

    pub fn get_mut(&mut self, point: &R::Point) -> Option<&mut V> {

        self.points.get_mut(self.find(point)).map(|(_,v)|v)

    }

    /// Convert a const reference to a point in the tree into an internal index.

    /// This function is unsafe because it can't be used productively without

    /// [`Self::launder_idx_mut`].  `ent` MUST be a reference to one of the points actually stored in

    /// the tree, NOT an identical point elsewhere, or this function invokes undefined behavior.

    pub unsafe fn launder_point_ref(&self, ent: &R::Point) -> usize {

        (ent as *const R::Point as *const (R::Point, V))

            .offset_from(self.points.as_ptr()) as usize

    }

    /// Convert a const reference to a value in the tree into an internal index.

    /// This function is unsafe because it can't be used productively without

    /// [`Self::launder_idx_point`].  `ent` MUST be a reference to one of the values actually stored in

    /// the tree, NOT an identical point elsewhere, or this function invokes undefined behavior.

    pub unsafe fn launder_value_ref(&self, ent: &V) -> usize {

        ((ent as *const V).byte_sub(mem::offset_of!((R::Point, V), 1))

            as *const (R::Point, V)

        ).offset_from(self.points.as_ptr()) as usize

    }

    /// Convert an internal index into a reference to a point in the tree.

    /// The internal index must have come from [`Self::launder_point_ref`] or [`Self::launder_value_ref`]

    /// called on the same tree.

    /// The intent of this function is to allow finding the points corresponding to values

    /// given a value reference, like for example if some of the values are made into

    /// an intrusive linked data structure.

    pub unsafe fn launder_idx_point(&self, idx: usize) -> &R::Point {

        &self.points[idx].0

    }

    /// Convert an internal index into a mutable reference to a value in the tree.

    /// The internal index must have come from [`Self::launder_point_ref`] or [`Self::launder_value_ref`]

    /// called on the same tree.

    /// The intent of this function is to allow mutating the values of the points in the

    /// result set of `k_closest` etc.

    pub unsafe fn launder_idx_mut(&mut self, idx: usize) -> &mut V {

        &mut self.points[idx].1

    }

    fn find(&self, point: &R::Point) -> usize {

        self.find_r(point, 0, self.len(), 0)

    }

    fn find_r(&self, point: &R::Point, mut a: usize, mut b: usize, mut layer: usize) -> usize {

        while a < b {

            let mid_idx = (a + b)/2;

            let p = &self.points[mid_idx].0;

            match point.cmp(p, layer) {

                Ordering::Less => b = mid_idx,

                Ordering::Greater => a = mid_idx + 1,

                Ordering::Equal => {

                    if point == p { return mid_idx }

                    a = self.find_r(point, a, mid_idx, layer + 1);

                    if a != self.len() { return a }

                    a = mid_idx + 1

                }

            }

            layer += 1

        }

        self.len()

    }

    #[cfg(test)]

        }

        }

    #[cfg(test)]

        }

    #[cfg(test)]

    #[cfg(test)]

}

pub struct Iter<'a, P: KdPoint, V> {

    buf: &'a [(P, V)],

    idx: usize

}

impl<'a, P: KdPoint, V> Iterator for Iter<'a, P, V> {

    type Item = (&'a P, &'a V);

}

pub struct IterMut<'a, P: KdPoint, V> {

    tail: &'a mut [(P, V)],

}

impl<'a, P, V> Iterator for IterMut<'a, P, V>

where P: KdPoint {

    type Item = (&'a P, &'a mut V);

    fn next(&mut self) -> Option<Self::Item> {

        let Some(([(a, b)], tail)) = mem::take(&mut self.tail).split_at_mut_checked(1) else { return None };

        self.tail = tail;

        Some((a, b))

    }

}

impl<'a, R: KdRegion, V> IntoIterator for &'a KdTree<R, V> {

	type Item = (&'a R::Point, &'a V);

	type IntoIter = Iter<'a, R::Point, V>;

}

impl<'a, R: KdRegion, V> IntoIterator for &'a mut KdTree<R, V> {

    type Item = (&'a R::Point, &'a mut V);

    type IntoIter = IterMut<'a, R::Point, V>;

    fn into_iter(self) -> Self::IntoIter {

        IterMut{tail: &mut self.points}

    }

}

impl<R: KdRegion, V> IntoIterator for KdTree<R, V> {

    type Item = (R::Point, V);

    type IntoIter = <Vec<(R::Point, V)> as IntoIterator>::IntoIter;

    fn into_iter(self) -> Self::IntoIter {

        self.points.into_vec().into_iter()

    }

}

impl<R: KdRegion, V, const N: usize> From<[(R::Point, V); N]> for KdTree<R, V> {

    fn from(ents: [(R::Point, V); N]) -> Self {

        let bounds = get_bounds((&ents).into_iter().map(|(p,_)|p));

        let points = ents.into_iter().collect();

        let mut res = Self{points, bounds};

        res.ify();

        res

    }

}

impl<R: KdRegion, V> FromIterator<(R::Point, V)> for KdTree<R, V> {

}

impl<R: KdRegion, V> Index<&R::Point> for KdTree<R, V> {

    type Output = V;

    fn index(&self, point: &R::Point) -> &V {

        self.get(point).unwrap()

    }

}

impl<R: KdRegion, V> IndexMut<&R::Point> for KdTree<R, V> {

    fn index_mut(&mut self, point: &R::Point) -> &mut Self::Output {

        self.get_mut(point).unwrap()

    }

}