{-# LANGUAGE UndecidableInstances #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Geometry.RangeTree
-- Copyright   :  (C) Frank Staals
-- License     :  see the LICENSE file
-- Maintainer  :  Frank Staals
--------------------------------------------------------------------------------
module Data.Geometry.RangeTree where

import           Control.Lens hiding (element)
import           Data.Ext
import qualified Data.Foldable as F
import           Data.Geometry.Point
import qualified Data.Geometry.RangeTree.Generic as GRT
import           Data.Geometry.RangeTree.Measure
import           Data.Geometry.Vector
import           Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NonEmpty
import           Data.Measured.Class
import           Data.Range
import           GHC.TypeLits
import           Prelude hiding (last,init,head)

--------------------------------------------------------------------------------

type RangeTree d = RT d d

newtype RT i d v p r =
  RangeTree { _unRangeTree :: GRT.RangeTree (Assoc i d v p r) (Leaf i d v p r) r }

deriving instance (Show r, Show (Assoc i d v p r), Show (Leaf i d v p r)) => Show (RT i d v p r)
deriving instance (Eq r,   Eq   (Assoc i d v p r), Eq   (Leaf i d v p r)) => Eq   (RT i d v p r)

newtype Leaf i d v p r = Leaf { _getPts :: [Point d r :+ p]} deriving (Semigroup,Monoid)

deriving instance (Show r, Show p, Arity d) => Show (Leaf i d v p r)
deriving instance (Eq r, Eq p,     Arity d) => Eq   (Leaf i d v p r)


type family AssocT i d v p r where
  AssocT 1 d v p r = v (Point d r :+ p)
  AssocT 2 d v p r = Maybe (RT 1 d v p r)

newtype Assoc i d v p r = Assoc { unAssoc :: AssocT i d v p r }

deriving instance Show (AssocT i d v p r) => Show (Assoc i d v p r)
deriving instance Eq   (AssocT i d v p r) => Eq   (Assoc i d v p r)


type RTMeasure v d p r = (LabeledMeasure v, Semigroup (v (Point d r :+ p)))

instance RTMeasure v d p r => Semigroup (Assoc 1 d v p r) where
  (Assoc l) <> (Assoc r) = Assoc $ l <> r

instance (RTMeasure v d p r, Ord r, 1 <= d, Arity d) => Semigroup (Assoc 2 d v p r) where
  (Assoc l) <> (Assoc r) = Assoc . createRangeTree'' $ toList l <> toList r
    where
      toList = maybe [] (F.toList . toAscList)
      createRangeTree'' = fmap createRangeTree1 . NonEmpty.nonEmpty



instance (RTMeasure v d p r, Ord r, 1 <= d, Arity d) => Monoid (Assoc 2 d v p r) where
  mempty = Assoc Nothing

----------------------------------------

instance ( RTMeasure v d p r
         ) => Measured (Assoc 1 d v p r) (Leaf 1 d v p r) where
  measure (Leaf pts) = Assoc . labeledMeasure $ pts

instance ( RTMeasure v d p r, Ord r, 1 <= d, Arity d
         ) => Measured (Assoc 2 d v p r) (Leaf 2 d v p r) where
  measure (Leaf pts) = Assoc . createRangeTree'' $ pts
    where
      createRangeTree'' = fmap createRangeTree1 . NonEmpty.nonEmpty

----------------------------------------

createRangeTree' :: (Ord r, RTMeasure v d p r
                   -- , Arity d, Arity (d+1), d ~ (d' + 1), Arity d'
                   -- , Measured (Assoc d v p r) (Leaf d v p r)
                   )
                 => [Point d r :+ p] -> Maybe (RT i d v p r)
createRangeTree' = fmap createRangeTree . NonEmpty.nonEmpty


createRangeTree :: (Ord r, RTMeasure v d p r
                   -- , Arity d, Arity (d+1), d ~ (d' + 1), Arity d'
                   -- , Measured (Assoc d v p r) (Leaf d v p r)
                   )
                => NonEmpty (Point d r :+ p) -> RT i d v p r
createRangeTree = undefined
-- RangeTree . GRT.createTree
--                 . fmap (\p -> last (p^.core.vector) :+ Leaf [p])


--------------------------------------------------------------------------------

-- | Gets all points in the range tree
toAscList :: RT i d v p r -> [Point d r :+ p]
toAscList = concatMap (^.extra.to _getPts) . F.toList . GRT.toAscList . _unRangeTree


--------------------------------------------------------------------------------

createRangeTree1 :: (Ord r, RTMeasure v d p r, 1 <= d, Arity d)
                 => NonEmpty (Point d r :+ p) -> RT 1 d v p r
createRangeTree1 = RangeTree . GRT.createTree
                . fmap (\p -> head (p^.core.vector) :+ Leaf [p])

createRangeTree2 :: forall v d r p. (Ord r, RTMeasure v d p r, Arity d, 2 <= d
                                    , 1 <= d -- this one is kind of silly
                 ) => NonEmpty (Point d r :+ p) -> RT 2 d v p r
createRangeTree2 = RangeTree . GRT.createTree
                 . fmap (\p -> p^.core.coord @2 :+ Leaf [p])

--------------------------------------------------------------------------------
-- * Querying


search   :: ( Ord r, Monoid (v (Point d r :+ p)), Query i d)
         => Vector d (Range r) -> RT i d v p r -> v (Point d r :+ p)
search r = mconcat . search' r


class (i <= d, Arity d) => Query i d where
  search' :: Ord r => Vector d (Range r) -> RT i d v p r -> [v (Point d r :+ p)]

instance (1 <= d, Arity d) => Query 1 d where
  search' qr = map unAssoc . GRT.search' r . _unRangeTree
    where
      r = qr^.element @0

instance ( 1 <= d, i <= d, Query (i-1) d, Arity d
         , i ~ 2
         ) => Query 2 d where
  search' qr = concatMap (maybe [] (search' qr) . unAssoc) . GRT.search' r . _unRangeTree
    where
      r = qr^.element @(i-1)