{-# 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 Show


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


-- | A node with `TNil`. for all its branches.
emptyNode :: QuadTree a
emptyNode = TNode TNil TNil TNil 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 tree
 = case tree of
        TNil            -> Nothing
        TLeaf{}         -> Nothing
        TNode nw ne sw se
         -> case quad of
                NW      -> Just nw
                NE      -> Just ne
                SW      -> Just sw
                SE      -> Just 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 f tree
 = case tree of
        TNil            -> tree
        TLeaf{}         -> tree
        TNode nw ne sw se
         -> case quad of
                NW      -> TNode (f nw) ne sw se
                NE      -> TNode nw (f ne) sw se
                SW      -> TNode nw ne (f sw) se
                SE      -> TNode nw ne sw (f 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 [] x _
        = TLeaf x
        
insertByPath (q:qs) x tree
 = case tree of
        TNil    -> liftToQuad q (insertByPath qs x) emptyNode
        TLeaf{} -> tree
        TNode{} -> liftToQuad q (insertByPath qs x) 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 x tree
 = do   path    <- pathToCoord extent coord
        return  $  insertByPath path x tree


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

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


-- | Lookup an element based given a path to it.
lookupByPath :: [Quad] -> QuadTree a -> Maybe a
lookupByPath path tree
 = case lookupNodeByPath path tree of
        Just (TLeaf x)  -> Just x
        _               -> 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 tree
 = do   path    <- pathToCoord extent coord
        lookupByPath path 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 extentInit treeInit
 = flatten' extentInit treeInit
 where  flatten' extent tree
         = case tree of
                TNil    -> []
                TLeaf x 
                 -> let (_, s, _, w) = takeExtent extent
                    in  [((w, s), x)]

                TNode{} -> concat $ map (flattenQuad extent tree) allQuads
                        
        flattenQuad extent tree quad
         = let  extent'         = cutQuadOfExtent quad extent
                Just tree'      = takeQuadOfTree quad tree
           in   flatten' extent' tree'


-- | Flatten a QuadTree into a list of its contained values, with coordinates.
flattenQuadTreeWithExtents
        :: forall a
        .  Extent       -- ^ Extent that covers the whole tree.
        -> QuadTree a 
        -> [(Extent, a)]
        
flattenQuadTreeWithExtents extentInit treeInit
 = flatten' extentInit treeInit
 where  flatten' extent tree
         = case tree of
                TNil    -> []
                TLeaf x 
                 -> [(extent, x)]

                TNode{} -> concat $ map (flattenQuad extent tree) allQuads
                        
        flattenQuad extent tree quad
         = let  extent'         = cutQuadOfExtent quad extent
                Just tree'      = takeQuadOfTree quad tree
           in   flatten' extent' tree'