content: {-# 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 Data.Maybe
  import System.IO
  import Vec2
  
  
  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 :: Vec2 -> Vec2 -> Vec2 -> Bool
  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
  