{-# LANGUAGE BangPatterns #-}
-- KDTree code
-- by Matthew Sottile <matt@galois.com> <mjsottile@computer.org>
--
module KDTree2d (
KDTreeNode(..),
newKDTree,
kdtAddPoints,
kdtAddPoint,
kdtRangeSearch,
kdtCollisionDetect,
kdtInBounds,
dumpKDTree,
mapKDTree,
kdtreeToList
) where
import Vec2
import Data.Maybe
import System.IO
data KDTreeNode a
= Empty
| Node !(KDTreeNode a) !Vec2 !a !(KDTreeNode a)
deriving Show
-- | An empty KDTree
newKDTree :: KDTreeNode a
newKDTree = Empty
-- | Flatten out a KDTree to a list.
kdtreeToList :: KDTreeNode a -> [a]
kdtreeToList Empty = []
kdtreeToList (Node l _ x r) = [x] ++ kdtreeToList l ++ kdtreeToList r
-- | Apply a worker function to all elements of a KDTree.
mapKDTree :: KDTreeNode a -> (a -> b) -> [b]
mapKDTree Empty _ = []
mapKDTree (Node l p n r) f = f n : (mapKDTree l f ++ mapKDTree r f)
kdtAddWithDepth :: KDTreeNode a -> Vec2 -> a -> Int -> KDTreeNode a
kdtAddWithDepth Empty pos dat _
= Node Empty pos dat Empty
kdtAddWithDepth (Node left npos ndata right) pos dat d
| vecDimSelect pos d < vecDimSelect npos d
= Node (kdtAddWithDepth left pos dat d') npos ndata right
| otherwise
= Node left npos ndata (kdtAddWithDepth right pos dat d')
where d' = if (d == 1) then 0 else 1
kdtAddPoint :: KDTreeNode a -> Vec2 -> a -> KDTreeNode a
kdtAddPoint t p d
= kdtAddWithDepth t p d 0
kdtInBounds p bMin bMax
= vecLessThan p bMax && vecGreaterThan p bMin
-- X dimension
kdtRangeSearchRecX :: KDTreeNode a -> Vec2 -> Vec2 -> [(Vec2,a)]
kdtRangeSearchRecX Empty _ _ = []
kdtRangeSearchRecX (Node left npos ndata right) bMin bMax
| nc < mnc
= nextfun right bMin bMax
| nc > mxc
= nextfun left bMin bMax
| kdtInBounds npos bMin bMax
= (npos, ndata)
: (nextfun right bMin bMax ++ nextfun left bMin bMax)
| otherwise
= nextfun right bMin bMax ++ nextfun left bMin bMax
where Vec2 nc _ = npos
Vec2 mnc _ = bMin
Vec2 mxc _ = bMax
nextfun = kdtRangeSearchRecY
-- Y dimension
kdtRangeSearchRecY :: (KDTreeNode a) -> Vec2 -> Vec2 -> [(Vec2,a)]
kdtRangeSearchRecY Empty _ _ = []
kdtRangeSearchRecY (Node left npos ndata right) bMin bMax
| nc < mnc
= nextfun right bMin bMax
| nc > mxc
= nextfun left bMin bMax
| (kdtInBounds npos bMin bMax)
= (npos, ndata)
: (nextfun right bMin bMax ++ nextfun left bMin bMax)
| otherwise
= nextfun right bMin bMax ++ nextfun left bMin bMax
where Vec2 _ nc = npos
Vec2 _ mnc = bMin
Vec2 _ mxc = bMax
nextfun = kdtRangeSearchRecX
kdtRangeSearch :: (KDTreeNode a) -> Vec2 -> Vec2 -> [(Vec2,a)]
kdtRangeSearch t bMin bMax
= kdtRangeSearchRecX t bMin bMax
kdtAddPoints :: [(Vec2,a)] -> (KDTreeNode a) -> (KDTreeNode a)
kdtAddPoints [] t = t
kdtAddPoints ((pt, dat):ps) t
= kdtAddPoints ps $ kdtAddPoint t pt dat
singleCollision :: Vec2 -> Vec2 -> Vec2 -> Double -> a -> Maybe (Vec2, a)
singleCollision pt start a eps dat
| sqrd_dist < eps * eps
= Just (vecAdd start p, dat)
| otherwise
= Nothing
where b = vecSub pt start
xhat = (vecDot a b) / (vecDot a a)
p = vecScale a xhat
e = vecSub p b
sqrd_dist = vecDot e e
kdtCollisionDetect :: KDTreeNode a -> Vec2 -> Vec2 -> Double -> [(Vec2,a)]
kdtCollisionDetect root !start !end !eps
= colls
where Vec2 sx sy = start
Vec2 ex ey = end
rmin = Vec2 (min sx ex - eps) (min sy ey - eps)
rmax = Vec2 (max sx ex + eps) (max sy ey + eps)
pts = kdtRangeSearch root rmin rmax
a = vecSub end start
colls = mapMaybe (\(pt,dat) -> singleCollision pt start a eps dat) pts
-- Dumping --------------------------------------------------------------------
-- | Dump a KDTree to a file
dumpKDTree :: KDTreeNode Int -> FilePath -> IO ()
dumpKDTree kdt name
= do h <- openFile name WriteMode
hPutStrLn h "n x y z"
dumpKDTreeInner kdt h
hClose h
-- | Dump a KDTree to a handle.
dumpKDTreeInner :: KDTreeNode Int -> Handle -> IO ()
dumpKDTreeInner kdt h
= case kdt of
Empty -> return ()
Node l v d r
-> do printVec v h d
dumpKDTreeInner l h
dumpKDTreeInner r h
-- | Print a vector to a handle.
printVec :: Vec2 -> Handle -> Int -> IO ()
printVec (Vec2 x y) h i
= hPutStrLn h $ show i ++ " " ++ show x ++ " " ++ show y