{-# LANGUAGE RankNTypes #-}

-- | A QuadTree can be used to recursively divide up 2D space into quadrants.
--   The smallest division corresponds to an unit `Extent`, so the total depth
--   of the tree will depend on what sized `Extent` you start with.
module Graphics.Gloss.Data.QuadTree
        ( QuadTree (..)
        , emptyTree
        , emptyNode
        , takeQuadOfTree
        , liftToQuad
        , insertByPath
        , insertByCoord
        , lookupNodeByPath
        , lookupByPath
        , lookupByCoord
        , flattenQuadTree
        , flattenQuadTreeWithExtents)
where
import Graphics.Gloss.Data.Quad
import Graphics.Gloss.Data.Extent

-- | The quad tree structure.
data QuadTree a
        -- | An empty node.
        = TNil

        -- | A leaf containint some value.
        | TLeaf a

        -- | A node with four children.
        | TNode (QuadTree a) (QuadTree a)       -- NW NE
                (QuadTree a) (QuadTree a)       -- SW SE
        deriving Int -> QuadTree a -> ShowS
[QuadTree a] -> ShowS
QuadTree a -> String
(Int -> QuadTree a -> ShowS)
-> (QuadTree a -> String)
-> ([QuadTree a] -> ShowS)
-> Show (QuadTree a)
forall a. Show a => Int -> QuadTree a -> ShowS
forall a. Show a => [QuadTree a] -> ShowS
forall a. Show a => QuadTree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [QuadTree a] -> ShowS
$cshowList :: forall a. Show a => [QuadTree a] -> ShowS
show :: QuadTree a -> String
$cshow :: forall a. Show a => QuadTree a -> String
showsPrec :: Int -> QuadTree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> QuadTree a -> ShowS
Show


-- | A `TNil` tree.
emptyTree :: QuadTree a
emptyTree :: QuadTree a
emptyTree = QuadTree a
forall a. QuadTree a
TNil


-- | A node with `TNil`. for all its branches.
emptyNode :: QuadTree a
emptyNode :: QuadTree a
emptyNode = QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
forall a.
QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
TNode QuadTree a
forall a. QuadTree a
TNil QuadTree a
forall a. QuadTree a
TNil QuadTree a
forall a. QuadTree a
TNil QuadTree a
forall a. QuadTree a
TNil


-- | Get a quadrant from a node.
--   If the tree does not have an outer node then `Nothing`.
takeQuadOfTree
        :: Quad
        -> QuadTree a
        -> Maybe (QuadTree a)

takeQuadOfTree :: Quad -> QuadTree a -> Maybe (QuadTree a)
takeQuadOfTree Quad
quad QuadTree a
tree
 = case QuadTree a
tree of
        QuadTree a
TNil            -> Maybe (QuadTree a)
forall a. Maybe a
Nothing
        TLeaf{}         -> Maybe (QuadTree a)
forall a. Maybe a
Nothing
        TNode QuadTree a
nw QuadTree a
ne QuadTree a
sw QuadTree a
se
         -> case Quad
quad of
                Quad
NW      -> QuadTree a -> Maybe (QuadTree a)
forall a. a -> Maybe a
Just QuadTree a
nw
                Quad
NE      -> QuadTree a -> Maybe (QuadTree a)
forall a. a -> Maybe a
Just QuadTree a
ne
                Quad
SW      -> QuadTree a -> Maybe (QuadTree a)
forall a. a -> Maybe a
Just QuadTree a
sw
                Quad
SE      -> QuadTree a -> Maybe (QuadTree a)
forall a. a -> Maybe a
Just QuadTree a
se


-- | Apply a function to a quadrant of a node.
--   If the tree does not have an outer node then return the original tree.
liftToQuad
        :: Quad
        -> (QuadTree a -> QuadTree a)
        -> QuadTree a  -> QuadTree a

liftToQuad :: Quad -> (QuadTree a -> QuadTree a) -> QuadTree a -> QuadTree a
liftToQuad Quad
quad QuadTree a -> QuadTree a
f QuadTree a
tree
 = case QuadTree a
tree of
        QuadTree a
TNil            -> QuadTree a
tree
        TLeaf{}         -> QuadTree a
tree
        TNode QuadTree a
nw QuadTree a
ne QuadTree a
sw QuadTree a
se
         -> case Quad
quad of
                Quad
NW      -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
forall a.
QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
TNode (QuadTree a -> QuadTree a
f QuadTree a
nw) QuadTree a
ne QuadTree a
sw QuadTree a
se
                Quad
NE      -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
forall a.
QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
TNode QuadTree a
nw (QuadTree a -> QuadTree a
f QuadTree a
ne) QuadTree a
sw QuadTree a
se
                Quad
SW      -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
forall a.
QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
TNode QuadTree a
nw QuadTree a
ne (QuadTree a -> QuadTree a
f QuadTree a
sw) QuadTree a
se
                Quad
SE      -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
forall a.
QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a -> QuadTree a
TNode QuadTree a
nw QuadTree a
ne QuadTree a
sw (QuadTree a -> QuadTree a
f QuadTree a
se)


-- | Insert a value into the tree at the position given by a path.
--   If the path intersects an existing `TLeaf` then return the original tree.
insertByPath :: [Quad] -> a -> QuadTree a -> QuadTree a

insertByPath :: [Quad] -> a -> QuadTree a -> QuadTree a
insertByPath [] a
x QuadTree a
_
        = a -> QuadTree a
forall a. a -> QuadTree a
TLeaf a
x

insertByPath (Quad
q:[Quad]
qs) a
x QuadTree a
tree
 = case QuadTree a
tree of
        QuadTree a
TNil    -> Quad -> (QuadTree a -> QuadTree a) -> QuadTree a -> QuadTree a
forall a.
Quad -> (QuadTree a -> QuadTree a) -> QuadTree a -> QuadTree a
liftToQuad Quad
q ([Quad] -> a -> QuadTree a -> QuadTree a
forall a. [Quad] -> a -> QuadTree a -> QuadTree a
insertByPath [Quad]
qs a
x) QuadTree a
forall a. QuadTree a
emptyNode
        TLeaf{} -> QuadTree a
tree
        TNode{} -> Quad -> (QuadTree a -> QuadTree a) -> QuadTree a -> QuadTree a
forall a.
Quad -> (QuadTree a -> QuadTree a) -> QuadTree a -> QuadTree a
liftToQuad Quad
q ([Quad] -> a -> QuadTree a -> QuadTree a
forall a. [Quad] -> a -> QuadTree a -> QuadTree a
insertByPath [Quad]
qs a
x) QuadTree a
tree


-- | Insert a value into the node containing this coordinate.
--   The node is created at maximum depth, corresponding to an unit `Extent`.
insertByCoord :: Extent -> Coord -> a -> QuadTree a -> Maybe (QuadTree a)
insertByCoord :: Extent -> Coord -> a -> QuadTree a -> Maybe (QuadTree a)
insertByCoord Extent
extent Coord
coord a
x QuadTree a
tree
 = do   [Quad]
path    <- Extent -> Coord -> Maybe [Quad]
pathToCoord Extent
extent Coord
coord
        QuadTree a -> Maybe (QuadTree a)
forall (m :: * -> *) a. Monad m => a -> m a
return  (QuadTree a -> Maybe (QuadTree a))
-> QuadTree a -> Maybe (QuadTree a)
forall a b. (a -> b) -> a -> b
$  [Quad] -> a -> QuadTree a -> QuadTree a
forall a. [Quad] -> a -> QuadTree a -> QuadTree a
insertByPath [Quad]
path a
x QuadTree a
tree


-- | Lookup a node based on a path to it.
lookupNodeByPath
        :: [Quad]
        -> QuadTree a
        -> Maybe (QuadTree a)

lookupNodeByPath :: [Quad] -> QuadTree a -> Maybe (QuadTree a)
lookupNodeByPath [] QuadTree a
tree
        = QuadTree a -> Maybe (QuadTree a)
forall a. a -> Maybe a
Just QuadTree a
tree

lookupNodeByPath (Quad
q:[Quad]
qs) QuadTree a
tree
 = case QuadTree a
tree of
        QuadTree a
TNil    -> Maybe (QuadTree a)
forall a. Maybe a
Nothing
        TLeaf{} -> Maybe (QuadTree a)
forall a. Maybe a
Nothing
        TNode{}
         -> let Just QuadTree a
quad       = Quad -> QuadTree a -> Maybe (QuadTree a)
forall a. Quad -> QuadTree a -> Maybe (QuadTree a)
takeQuadOfTree Quad
q QuadTree a
tree
            in  [Quad] -> QuadTree a -> Maybe (QuadTree a)
forall a. [Quad] -> QuadTree a -> Maybe (QuadTree a)
lookupNodeByPath [Quad]
qs QuadTree a
quad


-- | Lookup an element based given a path to it.
lookupByPath :: [Quad] -> QuadTree a -> Maybe a
lookupByPath :: [Quad] -> QuadTree a -> Maybe a
lookupByPath [Quad]
path QuadTree a
tree
 = case [Quad] -> QuadTree a -> Maybe (QuadTree a)
forall a. [Quad] -> QuadTree a -> Maybe (QuadTree a)
lookupNodeByPath [Quad]
path QuadTree a
tree of
        Just (TLeaf a
x)  -> a -> Maybe a
forall a. a -> Maybe a
Just a
x
        Maybe (QuadTree a)
_               -> Maybe a
forall a. Maybe a
Nothing


-- | Lookup a node if a tree given a coordinate which it contains.
lookupByCoord
        :: forall a
        .  Extent       -- ^ Extent that covers the whole tree.
        -> Coord        -- ^ Coordinate of the value of interest.
        -> QuadTree a
        -> Maybe a
lookupByCoord :: Extent -> Coord -> QuadTree a -> Maybe a
lookupByCoord Extent
extent Coord
coord QuadTree a
tree
 = do   [Quad]
path    <- Extent -> Coord -> Maybe [Quad]
pathToCoord Extent
extent Coord
coord
        [Quad] -> QuadTree a -> Maybe a
forall a. [Quad] -> QuadTree a -> Maybe a
lookupByPath [Quad]
path QuadTree a
tree


-- | Flatten a QuadTree into a list of its contained values, with coordinates.
flattenQuadTree
        :: forall a
        .  Extent       -- ^ Extent that covers the whole tree.
        -> QuadTree a
        -> [(Coord, a)]

flattenQuadTree :: Extent -> QuadTree a -> [(Coord, a)]
flattenQuadTree Extent
extentInit QuadTree a
treeInit
 = Extent -> QuadTree a -> [(Coord, a)]
forall b. Extent -> QuadTree b -> [(Coord, b)]
flatten' Extent
extentInit QuadTree a
treeInit
 where  flatten' :: Extent -> QuadTree b -> [(Coord, b)]
flatten' Extent
extent QuadTree b
tree
         = case QuadTree b
tree of
                QuadTree b
TNil    -> []
                TLeaf b
x
                 -> let (Int
_, Int
s, Int
_, Int
w) = Extent -> (Int, Int, Int, Int)
takeExtent Extent
extent
                    in  [((Int
w, Int
s), b
x)]

                TNode{} -> [[(Coord, b)]] -> [(Coord, b)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Coord, b)]] -> [(Coord, b)]) -> [[(Coord, b)]] -> [(Coord, b)]
forall a b. (a -> b) -> a -> b
$ (Quad -> [(Coord, b)]) -> [Quad] -> [[(Coord, b)]]
forall a b. (a -> b) -> [a] -> [b]
map (Extent -> QuadTree b -> Quad -> [(Coord, b)]
flattenQuad Extent
extent QuadTree b
tree) [Quad]
allQuads

        flattenQuad :: Extent -> QuadTree b -> Quad -> [(Coord, b)]
flattenQuad Extent
extent QuadTree b
tree Quad
quad
         = let  extent' :: Extent
extent'         = Quad -> Extent -> Extent
cutQuadOfExtent Quad
quad Extent
extent
                Just QuadTree b
tree'      = Quad -> QuadTree b -> Maybe (QuadTree b)
forall a. Quad -> QuadTree a -> Maybe (QuadTree a)
takeQuadOfTree Quad
quad QuadTree b
tree
           in   Extent -> QuadTree b -> [(Coord, b)]
flatten' Extent
extent' QuadTree b
tree'


-- | Flatten a QuadTree into a list of its contained values, with extents.
flattenQuadTreeWithExtents
        :: forall a
        .  Extent       -- ^ Extent that covers the whole tree.
        -> QuadTree a
        -> [(Extent, a)]

flattenQuadTreeWithExtents :: Extent -> QuadTree a -> [(Extent, a)]
flattenQuadTreeWithExtents Extent
extentInit QuadTree a
treeInit
 = Extent -> QuadTree a -> [(Extent, a)]
forall b. Extent -> QuadTree b -> [(Extent, b)]
flatten' Extent
extentInit QuadTree a
treeInit
 where  flatten' :: Extent -> QuadTree b -> [(Extent, b)]
flatten' Extent
extent QuadTree b
tree
         = case QuadTree b
tree of
                QuadTree b
TNil    -> []
                TLeaf b
x
                 -> [(Extent
extent, b
x)]

                TNode{} -> [[(Extent, b)]] -> [(Extent, b)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Extent, b)]] -> [(Extent, b)])
-> [[(Extent, b)]] -> [(Extent, b)]
forall a b. (a -> b) -> a -> b
$ (Quad -> [(Extent, b)]) -> [Quad] -> [[(Extent, b)]]
forall a b. (a -> b) -> [a] -> [b]
map (Extent -> QuadTree b -> Quad -> [(Extent, b)]
flattenQuad Extent
extent QuadTree b
tree) [Quad]
allQuads

        flattenQuad :: Extent -> QuadTree b -> Quad -> [(Extent, b)]
flattenQuad Extent
extent QuadTree b
tree Quad
quad
         = let  extent' :: Extent
extent'         = Quad -> Extent -> Extent
cutQuadOfExtent Quad
quad Extent
extent
                Just QuadTree b
tree'      = Quad -> QuadTree b -> Maybe (QuadTree b)
forall a. Quad -> QuadTree a -> Maybe (QuadTree a)
takeQuadOfTree Quad
quad QuadTree b
tree
           in   Extent -> QuadTree b -> [(Extent, b)]
flatten' Extent
extent' QuadTree b
tree'