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