{-# LANGUAGE TypeFamilies, DeriveFunctor, DeriveFoldable, DeriveTraversable #-}

import Control.Applicative {- from base -}
import Data.Distributive {- from distributive -}
import Data.Key {- from the keys package -}
import Data.Functor.Bind {- from semigroupoids -}
import Data.Functor.Representable {- from representable-functors -}

-- lets define unpacked affine 3-vectors of doubles
data Vector = Vector {-# UNPACK #-} !Double {-# UNPACK #-} !Double {-# UNPACK #-} !Double deriving (Eq, Ord, Show)

-- first lets define a functor that contains eight items
data Oct a = Oct !a !a !a !a !a !a !a !a deriving (Eq, Ord, Show, Functor, Foldable, Traversable)

-- we'll pretend that only ints from 0 to 7 exist to avoid making a more specific key type.
type instance Key Oct = Int

-- we can tabulate a function
instance Representable Oct where
  tabulate f = Oct (f 0) (f 1) (f 2) (f 3) (f 4) (f 5) (f 6) (f 7)

-- index into the structure
instance Indexable Oct where
  index (Oct s _ _ _ _ _ _ _) 0 = s
  index (Oct _ t _ _ _ _ _ _) 1 = t
  index (Oct _ _ u _ _ _ _ _) 2 = u
  index (Oct _ _ _ v _ _ _ _) 3 = v
  index (Oct _ _ _ _ w _ _ _) 4 = w
  index (Oct _ _ _ _ _ x _ _) 5 = x
  index (Oct _ _ _ _ _ _ y _) 6 = y
  index (Oct _ _ _ _ _ _ _ z) 7 = z
  index _ _ = error "index out of range"

-- adjust a single slot
instance Adjustable Oct where
  adjust f 0 (Oct s t u v w x y z) = Oct (f s) t u v w x y z
  adjust f 1 (Oct s t u v w x y z) = Oct s (f t) u v w x y z
  adjust f 2 (Oct s t u v w x y z) = Oct s t (f u) v w x y z
  adjust f 3 (Oct s t u v w x y z) = Oct s t u (f v) w x y z
  adjust f 4 (Oct s t u v w x y z) = Oct s t u v (f w) x y z
  adjust f 5 (Oct s t u v w x y z) = Oct s t u v w (f x) y z
  adjust f 6 (Oct s t u v w x y z) = Oct s t u v w x (f y) z
  adjust f 7 (Oct s t u v w x y z) = Oct s t u v w x y (f z)

-- and because we're representable, we can do lots of things
instance Lookup Oct where
  lookup = lookupDefault

instance Keyed Oct where
  mapWithKey = mapWithKeyRep

instance Zip Oct where
  zipWith = zipWithRep

instance ZipWithKey Oct where
  zipWithKey = zipWithKeyRep

instance Bind Oct where
  (>>-) = bindRep

instance Apply Oct where
  (<.>) = apRep

instance Applicative Oct where
  pure = pureRep
  (<*>) = apRep

instance Monad Oct where
  return = pureRep
  (>>=) = bindRep

instance Distributive Oct where
  distribute = distributeRep

-- then we'll make a primitive octree without annotating it with AABBs
data Tree a
  = Empty
  | Node {-# UNPACK #-} !(Oct (Tree a))
  | Tip {-# UNPACK #-} !Vector a
  deriving (Eq, Ord, Show, Functor, Foldable, Traversable)

-- but we'll build a wrapper around it that contains the AABB.
data Octree a = Octree
  { lo :: {-# UNPACK #-} !Vector
  , hi :: {-# UNPACK #-} !Vector
  , tree :: !(Tree a)
  } deriving (Eq, Ord, Show, Functor, Foldable, Traversable)

-- given this cheesy little convenience function
infix 7 ?
(?) :: Bool -> Int -> Int
False ? _ = 0
True ? x = x

-- we can make a combinator that knows how to calculate what octant the vector is in
-- and return the bounding box of that octant. If we don't use the vectors, laziness will
-- save us the trouble of computing them.
walk :: Vector -> Vector -> Vector -> (Int -> Vector -> Vector -> r) -> r
walk (Vector xl yl zl) (Vector xh yh zh) (Vector x y z) k =
  k (xq ? 1 + yq ? 2 + zq ? 4) (Vector xl' yl' zl') (Vector xh' yh' zh')
  where xm = xl + (xh - xl) / 2
        ym = yl + (yh - yl) / 2
        zm = zl + (zh - zl) / 2
        xq = x <= xm
        yq = y <= ym
        zq = z <= zm
        (xl',xh') = if xq then (xl,xm) else (xm,xh)
        (yl',yh') = if yq then (yl,ym) else (ym,yh)
        (zl',zh') = if zq then (zl,zm) else (zm,zh)

-- Now we are able to define insertWith directly
insertWith :: (a -> a -> a) -> Vector -> a -> Octree a -> Octree a
insertWith f v a (Octree lo0 hi0 t0) = Octree lo0 hi0 $ go lo0 hi0 t0 where
  -- empty node, so just insert
  go _ _ Empty = Tip v a
  -- we've found an octree node, figure out the desired branch and descend
  go lo hi (Node n) = walk lo hi v $ \o lo' hi' -> Node $ adjust (go lo' hi') o n
  -- we're at a leaf
  go lo hi n@(Tip v' a')
     -- if we have a collision, use our combining function
     | v == v'   = Tip v (f a' a)
     -- split the tip, and then continue inserting into the result
     | otherwise = walk lo hi v' $ \o _ _ -> go lo hi $ Node $ replace o n $ pure Empty

-- and we built insert on top of that.
insert :: Vector -> a -> Octree a -> Octree a
insert = insertWith const

-- we can map over the primitive octree with access to the key, but thats about it
type instance Key Tree = Vector

instance Keyed Tree where
  mapWithKey _ Empty = Empty
  mapWithKey f (Node e) = Node $ mapWithKey f <$> e
  mapWithKey f (Tip v a) = Tip v (f v a)

-- but our wrapped octrees have a notion of a key as well, and admit more efficient operations
type instance Key Octree = Vector

-- we can then piggyback some additional functionality from the types in Data.Key
instance Adjustable Octree where
  adjust f v (Octree lo0 hi0 t0) = Octree lo0 hi0 $ go lo0 hi0 t0 where
    go _  _  Empty    = Empty
    go lo hi (Node n) = walk lo hi v $ \o lo' hi' -> Node $ adjust (go lo' hi') o n
    go lo hi n@(Tip v' a)
       | v == v'   = Tip v' (f a)
       | otherwise = Tip v' a

instance Keyed Octree where
  mapWithKey f (Octree lo hi t) = Octree lo hi (mapWithKey t)

instance Lookup Octree where
  lookup v (Octree lo0 hi0 t0) = go lo0 hi0 t0 where
    go _  _  Empty    = Nothing
    go lo hi (Node n) = walk lo hi v $ \o lo' hi' -> go lo hi' (index n o)
    go lo hi (Tip v' a) | v == v' = Just a
                        | otherwise = Nothing

create :: Vector -> Vector -> Octree a
create lo hi = Octree lo hi Empty