{-# 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 Brillo.Data.QuadTree (
  QuadTree (..),
  emptyTree,
  emptyNode,
  takeQuadOfTree,
  liftToQuad,
  insertByPath,
  insertByCoord,
  lookupNodeByPath,
  lookupByPath,
  lookupByCoord,
  flattenQuadTree,
  flattenQuadTreeWithExtents,
)
where

import Brillo.Data.Extent
import Brillo.Data.Quad


-- | 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
$cshowsPrec :: forall a. Show a => Int -> QuadTree a -> ShowS
showsPrec :: Int -> QuadTree a -> ShowS
$cshow :: forall a. Show a => QuadTree a -> String
show :: QuadTree a -> String
$cshowList :: forall a. Show a => [QuadTree a] -> ShowS
showList :: [QuadTree a] -> ShowS
Show)


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


-- | A node with `TNil`. for all its branches.
emptyNode :: QuadTree a
emptyNode :: forall a. 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 :: forall a. 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 :: forall a.
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 :: forall a. [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 :: forall a. 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 a. a -> Maybe 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 :: forall a. [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{} ->
      case Quad -> QuadTree a -> Maybe (QuadTree a)
forall a. Quad -> QuadTree a -> Maybe (QuadTree a)
takeQuadOfTree Quad
q QuadTree a
tree of
        Maybe (QuadTree a)
Nothing -> Maybe (QuadTree a)
forall a. Maybe a
Nothing
        Just QuadTree a
quad -> [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 :: forall a. [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 :: forall a. 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 :: forall a. Extent -> QuadTree a -> [(Coord, a)]
flattenQuadTree Extent
extentInit QuadTree a
treeInit =
  Extent -> QuadTree a -> [(Coord, a)]
forall a. Extent -> QuadTree a -> [(Coord, a)]
flatten' Extent
extentInit QuadTree a
treeInit
  where
    flatten' :: Extent -> QuadTree a -> [(Coord, a)]
flatten' Extent
extent QuadTree a
tree =
      case QuadTree a
tree of
        QuadTree a
TNil -> []
        TLeaf a
x ->
          let (Int
_, Int
s, Int
_, Int
w) = Extent -> (Int, Int, Int, Int)
takeExtent Extent
extent
          in  [((Int
w, Int
s), a
x)]
        TNode{} -> [[(Coord, a)]] -> [(Coord, a)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Coord, a)]] -> [(Coord, a)]) -> [[(Coord, a)]] -> [(Coord, a)]
forall a b. (a -> b) -> a -> b
$ (Quad -> [(Coord, a)]) -> [Quad] -> [[(Coord, a)]]
forall a b. (a -> b) -> [a] -> [b]
map (Extent -> QuadTree a -> Quad -> [(Coord, a)]
flattenQuad Extent
extent QuadTree a
tree) [Quad]
allQuads

    flattenQuad :: Extent -> QuadTree a -> Quad -> [(Coord, a)]
flattenQuad Extent
extent QuadTree a
tree Quad
quad = do
      let extent' :: Extent
extent' = Quad -> Extent -> Extent
cutQuadOfExtent Quad
quad Extent
extent
      case Quad -> QuadTree a -> Maybe (QuadTree a)
forall a. Quad -> QuadTree a -> Maybe (QuadTree a)
takeQuadOfTree Quad
quad QuadTree a
tree of
        Maybe (QuadTree a)
Nothing -> []
        Just QuadTree a
tree' -> Extent -> QuadTree a -> [(Coord, a)]
flatten' Extent
extent' QuadTree a
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 :: forall a. Extent -> QuadTree a -> [(Extent, a)]
flattenQuadTreeWithExtents Extent
extentInit QuadTree a
treeInit =
  Extent -> QuadTree a -> [(Extent, a)]
forall a. Extent -> QuadTree a -> [(Extent, a)]
flatten' Extent
extentInit QuadTree a
treeInit
  where
    flatten' :: Extent -> QuadTree a -> [(Extent, a)]
flatten' Extent
extent QuadTree a
tree =
      case QuadTree a
tree of
        QuadTree a
TNil -> []
        TLeaf a
x ->
          [(Extent
extent, a
x)]
        TNode{} -> [[(Extent, a)]] -> [(Extent, a)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Extent, a)]] -> [(Extent, a)])
-> [[(Extent, a)]] -> [(Extent, a)]
forall a b. (a -> b) -> a -> b
$ (Quad -> [(Extent, a)]) -> [Quad] -> [[(Extent, a)]]
forall a b. (a -> b) -> [a] -> [b]
map (Extent -> QuadTree a -> Quad -> [(Extent, a)]
flattenQuad Extent
extent QuadTree a
tree) [Quad]
allQuads

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