module Data.KdTree.Dynamic
       ( -- * Usage

         -- $usage

         -- * Reference

         -- ** Types
         PointAsListFn
       , SquaredDistanceFn
       , KdTree
         -- ** Dynamic /k/-d tree construction
       , empty
       , singleton
       , emptyWithDist
       , singletonWithDist
         -- ** Insertion
       , insert
         -- ** Query
       , nearest
       , inRadius
       , kNearest
       , inRange
       , toList
       , null
       , size
         -- ** Utilities
       , defaultSqrDist
       ) where

import Prelude hiding (null)

import qualified Data.Foldable as F

import qualified Data.KdMap.Dynamic as DKDM
import Data.KdMap.Dynamic (PointAsListFn, SquaredDistanceFn, defaultSqrDist)

-- $usage
--
-- The 'KdTree' is a dynamic variant of
-- @Data.KdTree.Static.@'Data.KdTree.Static.KdTree' that allows for
-- insertion of new points into an existing 'KdTree'. This algorithm
-- was implemented using a
-- <http://repository.cmu.edu/cgi/viewcontent.cgi?article=3453&context=compsci static-to-dynamic transformation>.
--
-- Here's an example of interleaving 3D point insertions and point
-- queries using 'KdTree':
--
-- @
-- >>> let dkdt = 'singleton' point3dAsList (Point3D 0.0 0.0 0.0)
--
-- >>> let dkdt' = 'insert' dkdt (Point3D 1.0 1.0 1.0)
--
-- >>> 'nearest' dkdt' (Point3D 0.4 0.4 0.4)
-- Point3D {x = 0.0, y = 0.0, z = 0.0}
--
-- >>> let dkdt'' = 'insert' dkdt' (Point3D 0.5 0.5 0.5)
--
-- >>> 'nearest' dkdt'' (Point3D 0.4 0.4 0.4)
-- Point3D {x = 0.5, y = 0.5, z = 0.5}
-- @
--
-- Check out @Data.KdMap.Dynamic.@'Data.KdMap.Dynamic.KdMap' if you
-- want to associate a value with each point in your tree structure.

-- | A dynamic /k/-d tree structure that stores points of type @p@
-- with axis values of type @a@.
newtype KdTree a p = KdTree (DKDM.KdMap a p ())

instance F.Foldable (KdTree a) where
  foldr :: (a -> b -> b) -> b -> KdTree a a -> b
foldr a -> b -> b
f b
z (KdTree KdMap a a ()
dkdMap) = ((a, ()) -> b -> b) -> b -> KdMap a a () -> b
forall p v b a. ((p, v) -> b -> b) -> b -> KdMap a p v -> b
DKDM.foldrWithKey (a -> b -> b
f (a -> b -> b) -> ((a, ()) -> a) -> (a, ()) -> b -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, ()) -> a
forall a b. (a, b) -> a
fst) b
z KdMap a a ()
dkdMap

instance (Show a, Show p) => Show (KdTree a p) where
  show :: KdTree a p -> String
show (KdTree KdMap a p ()
kdm) = String
"KdTree " String -> ShowS
forall a. [a] -> [a] -> [a]
++ KdMap a p () -> String
forall a. Show a => a -> String
show KdMap a p ()
kdm

-- | Generates an empty 'KdTree' with a user-specified distance function.
emptyWithDist :: Real a => PointAsListFn a p -> SquaredDistanceFn a p -> KdTree a p
emptyWithDist :: PointAsListFn a p -> SquaredDistanceFn a p -> KdTree a p
emptyWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2 = KdMap a p () -> KdTree a p
forall a p. KdMap a p () -> KdTree a p
KdTree (KdMap a p () -> KdTree a p) -> KdMap a p () -> KdTree a p
forall a b. (a -> b) -> a -> b
$ PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p ()
forall a p v.
PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v
DKDM.emptyWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2

-- | Generates an empty 'KdTree' with the default distance function.
empty :: Real a => PointAsListFn a p -> KdTree a p
empty :: PointAsListFn a p -> KdTree a p
empty PointAsListFn a p
p2l = PointAsListFn a p -> SquaredDistanceFn a p -> KdTree a p
forall a p.
Real a =>
PointAsListFn a p -> SquaredDistanceFn a p -> KdTree a p
emptyWithDist PointAsListFn a p
p2l (SquaredDistanceFn a p -> KdTree a p)
-> SquaredDistanceFn a p -> KdTree a p
forall a b. (a -> b) -> a -> b
$ PointAsListFn a p -> SquaredDistanceFn a p
forall a p. Num a => PointAsListFn a p -> SquaredDistanceFn a p
defaultSqrDist PointAsListFn a p
p2l

-- | Returns whether the 'KdTree' is empty.
null :: KdTree a p -> Bool
null :: KdTree a p -> Bool
null (KdTree KdMap a p ()
dkdMap) = KdMap a p () -> Bool
forall a p v. KdMap a p v -> Bool
DKDM.null KdMap a p ()
dkdMap

-- | Generates a 'KdTree' with a single point using a
-- user-specified distance function.
singletonWithDist :: Real a => PointAsListFn a p
                            -> SquaredDistanceFn a p
                            -> p
                            -> KdTree a p
singletonWithDist :: PointAsListFn a p -> SquaredDistanceFn a p -> p -> KdTree a p
singletonWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2 p
p = KdMap a p () -> KdTree a p
forall a p. KdMap a p () -> KdTree a p
KdTree (KdMap a p () -> KdTree a p) -> KdMap a p () -> KdTree a p
forall a b. (a -> b) -> a -> b
$ PointAsListFn a p
-> SquaredDistanceFn a p -> (p, ()) -> KdMap a p ()
forall a p v.
Real a =>
PointAsListFn a p -> SquaredDistanceFn a p -> (p, v) -> KdMap a p v
DKDM.singletonWithDist PointAsListFn a p
p2l SquaredDistanceFn a p
d2 (p
p, ())

-- | Generates a 'KdTree' with a single point using the default
-- distance function.
singleton :: Real a => PointAsListFn a p -> p -> KdTree a p
singleton :: PointAsListFn a p -> p -> KdTree a p
singleton PointAsListFn a p
p2l = PointAsListFn a p -> SquaredDistanceFn a p -> p -> KdTree a p
forall a p.
Real a =>
PointAsListFn a p -> SquaredDistanceFn a p -> p -> KdTree a p
singletonWithDist PointAsListFn a p
p2l (SquaredDistanceFn a p -> p -> KdTree a p)
-> SquaredDistanceFn a p -> p -> KdTree a p
forall a b. (a -> b) -> a -> b
$ PointAsListFn a p -> SquaredDistanceFn a p
forall a p. Num a => PointAsListFn a p -> SquaredDistanceFn a p
defaultSqrDist PointAsListFn a p
p2l

-- | Adds a given point to a 'KdTree'.
--
-- Average time complexity per insert for /n/ inserts: /O(log^2(n))/.
insert :: Real a => KdTree a p -> p -> KdTree a p
insert :: KdTree a p -> p -> KdTree a p
insert (KdTree KdMap a p ()
dkdMap) p
p = KdMap a p () -> KdTree a p
forall a p. KdMap a p () -> KdTree a p
KdTree (KdMap a p () -> KdTree a p) -> KdMap a p () -> KdTree a p
forall a b. (a -> b) -> a -> b
$ KdMap a p () -> p -> () -> KdMap a p ()
forall a p v. Real a => KdMap a p v -> p -> v -> KdMap a p v
DKDM.insert KdMap a p ()
dkdMap p
p ()

-- | Given a 'KdTree' and a query point, returns the nearest point
-- in the 'KdTree' to the query point.
--
-- Average time complexity: /O(log^2(n))/.
nearest :: Real a => KdTree a p -> p -> p
nearest :: KdTree a p -> p -> p
nearest (KdTree KdMap a p ()
dkdMap) = (p, ()) -> p
forall a b. (a, b) -> a
fst ((p, ()) -> p) -> (p -> (p, ())) -> p -> p
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KdMap a p () -> p -> (p, ())
forall a p v. Real a => KdMap a p v -> p -> (p, v)
DKDM.nearest KdMap a p ()
dkdMap

-- | Given a 'KdTree', a query point, and a number @k@, returns the
-- @k@ nearest points in the 'KdTree' to the query point.
--
-- Neighbors are returned in order of increasing distance from query
-- point.
--
-- Average time complexity: /log(k) * log^2(n)/ for /k/ nearest
-- neighbors on a structure with /n/ data points.
--
-- Worst case time complexity: /n * log(k)/ for /k/ nearest neighbors
-- on a structure with /n/ data points.
kNearest :: Real a => KdTree a p -> Int -> p -> [p]
kNearest :: KdTree a p -> Int -> p -> [p]
kNearest (KdTree KdMap a p ()
dkdMap) Int
k p
query =
  ((p, ()) -> p) -> [(p, ())] -> [p]
forall a b. (a -> b) -> [a] -> [b]
map (p, ()) -> p
forall a b. (a, b) -> a
fst ([(p, ())] -> [p]) -> [(p, ())] -> [p]
forall a b. (a -> b) -> a -> b
$ KdMap a p () -> Int -> p -> [(p, ())]
forall a p v. Real a => KdMap a p v -> Int -> p -> [(p, v)]
DKDM.kNearest KdMap a p ()
dkdMap Int
k p
query

-- | Given a 'KdTree', a query point, and a radius, returns all
-- points in the 'KdTree' that are within the given radius of the
-- query points.
--
-- Points are not returned in any particular order.
--
-- Worst case time complexity: /O(n)/ for /n/ data points.
inRadius :: Real a => KdTree a p -> a -> p -> [p]
inRadius :: KdTree a p -> a -> p -> [p]
inRadius (KdTree KdMap a p ()
dkdMap) a
radius p
query =
  ((p, ()) -> p) -> [(p, ())] -> [p]
forall a b. (a -> b) -> [a] -> [b]
map (p, ()) -> p
forall a b. (a, b) -> a
fst ([(p, ())] -> [p]) -> [(p, ())] -> [p]
forall a b. (a -> b) -> a -> b
$ KdMap a p () -> a -> p -> [(p, ())]
forall a p v. Real a => KdMap a p v -> a -> p -> [(p, v)]
DKDM.inRadius KdMap a p ()
dkdMap a
radius p
query

-- | Finds all points in a 'KdTree' with points within a given range,
-- where the range is specified as a set of lower and upper bounds.
--
-- Points are not returned in any particular order.
--
-- Worst case time complexity: /O(n)/ for n data points and a range
-- that spans all the points.
inRange :: Real a => KdTree a p
                     -> p -- ^ lower bounds of range
                     -> p -- ^ upper bounds of range
                     -> [p] -- ^ all points within given range
inRange :: KdTree a p -> p -> p -> [p]
inRange (KdTree KdMap a p ()
dkdMap) p
lowers p
uppers =
  ((p, ()) -> p) -> [(p, ())] -> [p]
forall a b. (a -> b) -> [a] -> [b]
map (p, ()) -> p
forall a b. (a, b) -> a
fst ([(p, ())] -> [p]) -> [(p, ())] -> [p]
forall a b. (a -> b) -> a -> b
$ KdMap a p () -> p -> p -> [(p, ())]
forall a p v. Real a => KdMap a p v -> p -> p -> [(p, v)]
DKDM.inRange KdMap a p ()
dkdMap p
lowers p
uppers

-- | Returns the number of elements in the 'KdTree'.
--
-- Time complexity: /O(1)/
size :: KdTree a p -> Int
size :: KdTree a p -> Int
size (KdTree KdMap a p ()
dkdMap) = KdMap a p () -> Int
forall a p v. KdMap a p v -> Int
DKDM.size KdMap a p ()
dkdMap

-- | Returns a list of all the points in the 'KdTree'.
--
-- Time complexity: /O(n)/
toList :: KdTree a p -> [p]
toList :: KdTree a p -> [p]
toList (KdTree KdMap a p ()
dkdMap) = KdMap a p () -> [p]
forall a p v. KdMap a p v -> [p]
DKDM.keys KdMap a p ()
dkdMap