vptree.hs

module Linear.VPTree
  ( fromList
  , nearest
  ) where

import Linear
import Data.Word8
import Data.List (partition, sortBy)
import Data.Ord (comparing)
import Data.Function (on)

mid ::  Ord a => [a] -> a
mid xs = xs !! (length xs `div` 2)

type Point = V3 Word8

data VPTree = VPTree
  { _point   :: Point
  , _middle  :: Double
  , _lbound  :: Double
  , _left    :: VPTree
  , _right   :: VPTree
  }
  | VPEmpty
  deriving (Show)

distance' :: Point -> Point -> Double
distance' = distance `on` fmap fromIntegral

fromList :: [Point] -> VPTree
fromList []     = VPEmpty
fromList [p]    = VPTree p 0      0      VPEmpty           VPEmpty
fromList (p:ps) = VPTree p middle lbound (fromList inners) (fromList outers) where
  fromP = distance' p
  byDist = sortBy (comparing fromP) ps
  middle = fromP $ mid byDist
  lbound = fromP $ head byDist
  (inners, outers) = partition ((< middle) . fromP) ps

nearest :: VPTree -> Point -> Maybe Point
nearest VPEmpty                         _ = Nothing
nearest (VPTree p' _ _ VPEmpty VPEmpty) _ = Just p'
nearest (VPTree p' d l ins outs) p
  | distance' p p' < l = Just p'
  | distance' p p' < d = nearest ins p
  | otherwise          = nearest outs p