module Data.KdMap.Dynamic
(
PointAsListFn
, SquaredDistanceFn
, KdMap
, empty
, singleton
, emptyWithDist
, singletonWithDist
, insert
, insertPair
, batchInsert
, nearest
, inRadius
, kNearest
, inRange
, assocs
, keys
, elems
, null
, size
, foldrWithKey
, defaultSqrDist
, subtreeSizes
) where
import Prelude hiding (null)
import Control.Applicative hiding (empty)
import Data.Bits
import Data.Foldable
import Data.Function
import Data.List as L hiding (insert, null)
import qualified Data.List (null)
import Data.Traversable
import Control.DeepSeq
import Control.DeepSeq.Generics (genericRnf)
import GHC.Generics
import qualified Data.KdMap.Static as KDM
import Data.KdMap.Static (PointAsListFn, SquaredDistanceFn, defaultSqrDist)
data KdMap a p v = KdMap
{ _trees :: [KDM.KdMap a p v]
, _pointAsList :: PointAsListFn a p
, _distSqr :: SquaredDistanceFn a p
, _numNodes :: Int
} deriving Generic
instance (NFData a, NFData p, NFData v) => NFData (KdMap a p v) where rnf = genericRnf
instance (Show a, Show p, Show v) => Show (KdMap a p v) where
show kdm = "KdMap " ++ show (_trees kdm)
instance Functor (KdMap a p) where
fmap f dkdMap = dkdMap { _trees = map (fmap f) $ _trees dkdMap }
foldrWithKey :: ((p, v) -> b -> b) -> b -> KdMap a p v -> b
foldrWithKey f z dkdMap = L.foldr (flip $ KDM.foldrWithKey f) z $ _trees dkdMap
instance Foldable (KdMap a p) where
foldr f = foldrWithKey (f . snd)
instance Traversable (KdMap a p) where
traverse f (KdMap t p d n) =
KdMap <$> traverse (traverse f) t <*> pure p <*> pure d <*> pure n
emptyWithDist :: PointAsListFn a p -> SquaredDistanceFn a p -> KdMap a p v
emptyWithDist p2l d2 = KdMap [] p2l d2 0
null :: KdMap a p v -> Bool
null (KdMap [] _ _ _) = True
null _ = False
singletonWithDist :: Real a => PointAsListFn a p
-> SquaredDistanceFn a p
-> (p, v)
-> KdMap a p v
singletonWithDist p2l d2 (k, v) =
KdMap [KDM.buildWithDist p2l d2 [(k, v)]] p2l d2 1
empty :: Real a => PointAsListFn a p -> KdMap a p v
empty p2l = emptyWithDist p2l $ defaultSqrDist p2l
singleton :: Real a => PointAsListFn a p -> (p, v) -> KdMap a p v
singleton p2l = singletonWithDist p2l $ defaultSqrDist p2l
insert :: Real a => KdMap a p v -> p -> v -> KdMap a p v
insert (KdMap trees p2l d2 n) k v =
let bitList = map ((1 .&.) . (n `shiftR`)) [0..]
(onesPairs, theRestPairs) = span ((== 1) . fst) $ zip bitList trees
((_, ones), (_, theRest)) = (unzip onesPairs, unzip theRestPairs)
newTree = KDM.buildWithDist p2l d2 $ (k, v) : L.concatMap KDM.assocs ones
in KdMap (newTree : theRest) p2l d2 $ n + 1
insertPair :: Real a => KdMap a p v -> (p, v) -> KdMap a p v
insertPair t = uncurry (insert t)
nearest :: Real a => KdMap a p v -> p -> (p, v)
nearest (KdMap ts _ d2 _) query =
let nearests = map (`KDM.nearest` query) ts
in if Data.List.null nearests
then error "Called nearest on empty KdMap."
else L.minimumBy (compare `on` (d2 query . fst)) nearests
kNearest :: Real a => KdMap a p v -> Int -> p -> [(p, v)]
kNearest (KdMap trees _ d2 _) k query =
let neighborSets = map (\t -> KDM.kNearest t k query) trees
in take k $ L.foldr merge [] neighborSets
where merge [] ys = ys
merge xs [] = xs
merge xs@(x:xt) ys@(y:yt)
| distX <= distY = x : merge xt ys
| otherwise = y : merge xs yt
where distX = d2 query $ fst x
distY = d2 query $ fst y
inRadius :: Real a => KdMap a p v -> a -> p -> [(p, v)]
inRadius (KdMap trees _ _ _) radius query =
L.concatMap (\t -> KDM.inRadius t radius query) trees
inRange :: Real a => KdMap a p v
-> p
-> p
-> [(p, v)]
inRange (KdMap trees _ _ _) lowers uppers =
L.concatMap (\t -> KDM.inRange t lowers uppers) trees
size :: KdMap a p v -> Int
size (KdMap _ _ _ n) = n
assocs :: KdMap a p v -> [(p, v)]
assocs (KdMap trees _ _ _) = L.concatMap KDM.assocs trees
keys :: KdMap a p v -> [p]
keys = map fst . assocs
elems :: KdMap a p v -> [v]
elems = map snd . assocs
batchInsert :: Real a => KdMap a p v -> [(p, v)] -> KdMap a p v
batchInsert = L.foldl' insertPair
subtreeSizes :: KdMap a p v -> [Int]
subtreeSizes (KdMap trees _ _ _) = map KDM.size trees