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

import           Control.Lens             hiding (Empty, element, imap, (:<))
import           Data.BinaryTree
import           Data.Ext
import qualified Data.Foldable            as F
import           Data.Geometry.Box
import           Data.Geometry.Point
import           Data.Geometry.Properties
import           Data.Geometry.Vector
import           Data.LSeq                (LSeq, pattern (:<|))
import qualified Data.LSeq                as LSeq
import qualified Data.List.NonEmpty       as NonEmpty
import           Data.Proxy
import           Data.Util
import qualified Data.Vector.Fixed        as FV
import           GHC.TypeLits
import           Prelude                  hiding (replicate)
import           Unsafe.Coerce            (unsafeCoerce)

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

newtype Coord (d :: Nat) = Coord { Coord d -> Int
unCoord ::  Int}

instance KnownNat d => Eq (Coord d) where
  (Coord Int
i) == :: Coord d -> Coord d -> Bool
== (Coord Int
j) = (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
d) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
j Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
d)
    where
      d :: Int
d = Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Int) -> (Proxy d -> Integer) -> Proxy d -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy d -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal (Proxy d -> Int) -> Proxy d -> Int
forall a b. (a -> b) -> a -> b
$ (Proxy d
forall k (t :: k). Proxy t
Proxy :: Proxy d)

instance KnownNat d => Show (Coord d) where
  show :: Coord d -> String
show (Coord Int
i) = Int -> String
forall a. Show a => a -> String
show (Int -> String) -> Int -> String
forall a b. (a -> b) -> a -> b
$ Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
d)
    where
      d :: Int
d = Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Int) -> (Proxy d -> Integer) -> Proxy d -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy d -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal (Proxy d -> Int) -> Proxy d -> Int
forall a b. (a -> b) -> a -> b
$ (Proxy d
forall k (t :: k). Proxy t
Proxy :: Proxy d)

instance KnownNat d => Enum (Coord d) where
  toEnum :: Int -> Coord d
toEnum Int
i = Int -> Coord d
forall (d :: Nat). Int -> Coord d
Coord (Int -> Coord d) -> Int -> Coord d
forall a b. (a -> b) -> a -> b
$ Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
i Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
d)
    where
      d :: Int
d = Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Int) -> (Proxy d -> Integer) -> Proxy d -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy d -> Integer
forall (n :: Nat) (proxy :: Nat -> *).
KnownNat n =>
proxy n -> Integer
natVal (Proxy d -> Int) -> Proxy d -> Int
forall a b. (a -> b) -> a -> b
$ (Proxy d
forall k (t :: k). Proxy t
Proxy :: Proxy d)
  fromEnum :: Coord d -> Int
fromEnum = Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1 (Int -> Int) -> (Coord d -> Int) -> Coord d -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Coord d -> Int
forall (d :: Nat). Coord d -> Int
unCoord


data Split d r = Split !(Coord d) !r !(Box d () r)

deriving instance (Show r, Arity d, KnownNat d) => Show (Split d r)
deriving instance (Eq r, Arity d, KnownNat d)   => Eq (Split d r)


type Split' d r = SP (Coord d) r

newtype KDTree' d p r = KDT { KDTree' d p r -> BinLeafTree (Split d r) (Point d r :+ p)
unKDT :: BinLeafTree (Split d r) (Point d r :+ p) }

deriving instance (Show p, Show r, Arity d, KnownNat d) => Show (KDTree' d p r)
deriving instance (Eq p, Eq r, Arity d, KnownNat d)     => Eq   (KDTree' d p r)


data KDTree d p r = Empty
                  | Tree (KDTree' d p r)

deriving instance (Show p, Show r, Arity d, KnownNat d) => Show (KDTree d p r)
deriving instance (Eq p, Eq r, Arity d, KnownNat d)     => Eq   (KDTree d p r)


toMaybe          :: KDTree d p r -> Maybe (KDTree' d p r)
toMaybe :: KDTree d p r -> Maybe (KDTree' d p r)
toMaybe KDTree d p r
Empty    = Maybe (KDTree' d p r)
forall a. Maybe a
Nothing
toMaybe (Tree KDTree' d p r
t) = KDTree' d p r -> Maybe (KDTree' d p r)
forall a. a -> Maybe a
Just KDTree' d p r
t


-- | Expects the input to be a set, i.e. no duplicates
--
-- running time: \(O(n \log n)\)
buildKDTree :: (Arity d, 1 <= d, Ord r)
            => [Point d r :+ p] -> KDTree d p r
buildKDTree :: [Point d r :+ p] -> KDTree d p r
buildKDTree = KDTree d p r
-> (NonEmpty (Point d r :+ p) -> KDTree d p r)
-> Maybe (NonEmpty (Point d r :+ p))
-> KDTree d p r
forall b a. b -> (a -> b) -> Maybe a -> b
maybe KDTree d p r
forall (d :: Nat) p r. KDTree d p r
Empty (KDTree' d p r -> KDTree d p r
forall (d :: Nat) p r. KDTree' d p r -> KDTree d p r
Tree (KDTree' d p r -> KDTree d p r)
-> (NonEmpty (Point d r :+ p) -> KDTree' d p r)
-> NonEmpty (Point d r :+ p)
-> KDTree d p r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty (Point d r :+ p) -> KDTree' d p r
forall (d :: Nat) r p.
(Arity d, 1 <= d, Ord r) =>
NonEmpty (Point d r :+ p) -> KDTree' d p r
buildKDTree') (Maybe (NonEmpty (Point d r :+ p)) -> KDTree d p r)
-> ([Point d r :+ p] -> Maybe (NonEmpty (Point d r :+ p)))
-> [Point d r :+ p]
-> KDTree d p r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Point d r :+ p] -> Maybe (NonEmpty (Point d r :+ p))
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty

buildKDTree' :: (Arity d, 1 <= d, Ord r)
             => NonEmpty.NonEmpty (Point d r :+ p) -> KDTree' d p r
buildKDTree' :: NonEmpty (Point d r :+ p) -> KDTree' d p r
buildKDTree' = BinLeafTree (Split d r) (Point d r :+ p) -> KDTree' d p r
forall (d :: Nat) p r.
BinLeafTree (Split d r) (Point d r :+ p) -> KDTree' d p r
KDT (BinLeafTree (Split d r) (Point d r :+ p) -> KDTree' d p r)
-> (NonEmpty (Point d r :+ p)
    -> BinLeafTree (Split d r) (Point d r :+ p))
-> NonEmpty (Point d r :+ p)
-> KDTree' d p r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinLeafTree (SP (Coord d) r) (Point d r :+ p)
-> BinLeafTree (Split d r) (Point d r :+ p)
forall core extra.
(ImplicitPeano (Peano (Dimension core)),
 ArityPeano (Peano (FromPeano (Peano (Dimension core)))),
 KnownNat (FromPeano (Peano (Dimension core))),
 KnownNat (Dimension core), IsBoxable core, Ord (NumType core),
 Peano (FromPeano (Peano (Dimension core)) + 1)
 ~ 'S (Peano (FromPeano (Peano (Dimension core))))) =>
BinLeafTree
  (SP (Coord (Dimension core)) (NumType core)) (core :+ extra)
-> BinLeafTree
     (Split (Dimension core) (NumType core)) (core :+ extra)
addBoxes (BinLeafTree (SP (Coord d) r) (Point d r :+ p)
 -> BinLeafTree (Split d r) (Point d r :+ p))
-> (NonEmpty (Point d r :+ p)
    -> BinLeafTree (SP (Coord d) r) (Point d r :+ p))
-> NonEmpty (Point d r :+ p)
-> BinLeafTree (Split d r) (Point d r :+ p)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Coord d
-> PointSet (LSeq 1) d p r
-> BinLeafTree (SP (Coord d) r) (Point d r :+ p)
forall (d :: Nat) r p.
(1 <= d, Arity d, Ord r) =>
Coord d
-> PointSet (LSeq 1) d p r
-> BinLeafTree (Split' d r) (Point d r :+ p)
build (Int -> Coord d
forall (d :: Nat). Int -> Coord d
Coord Int
1) (PointSet (LSeq 1) d p r
 -> BinLeafTree (SP (Coord d) r) (Point d r :+ p))
-> (NonEmpty (Point d r :+ p) -> PointSet (LSeq 1) d p r)
-> NonEmpty (Point d r :+ p)
-> BinLeafTree (SP (Coord d) r) (Point d r :+ p)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LSeq 1 (Point d r :+ p) -> PointSet (LSeq 1) d p r
forall (d :: Nat) r (n :: Nat) p.
(Arity d, Ord r) =>
LSeq n (Point d r :+ p) -> PointSet (LSeq n) d p r
toPointSet (LSeq 1 (Point d r :+ p) -> PointSet (LSeq 1) d p r)
-> (NonEmpty (Point d r :+ p) -> LSeq 1 (Point d r :+ p))
-> NonEmpty (Point d r :+ p)
-> PointSet (LSeq 1) d p r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty (Point d r :+ p) -> LSeq 1 (Point d r :+ p)
forall a. NonEmpty a -> LSeq 1 a
LSeq.fromNonEmpty
  where     -- compute one tree with bounding boxes, then merge them together
    addBoxes :: BinLeafTree
  (SP (Coord (Dimension core)) (NumType core)) (core :+ extra)
-> BinLeafTree
     (Split (Dimension core) (NumType core)) (core :+ extra)
addBoxes BinLeafTree
  (SP (Coord (Dimension core)) (NumType core)) (core :+ extra)
t = let bbt :: BinLeafTree
  (Box (Dimension core) () (NumType core)) (core :+ extra)
bbt = (Box (Dimension core) () (NumType core)
 -> SP (Coord (Dimension core)) (NumType core)
 -> Box (Dimension core) () (NumType core)
 -> Box (Dimension core) () (NumType core))
-> ((core :+ extra) -> Box (Dimension core) () (NumType core))
-> BinLeafTree
     (SP (Coord (Dimension core)) (NumType core)) (core :+ extra)
-> BinLeafTree
     (Box (Dimension core) () (NumType core)) (core :+ extra)
forall w v a.
(w -> v -> w -> w)
-> (a -> w) -> BinLeafTree v a -> BinLeafTree w a
foldUpData (\Box (Dimension core) () (NumType core)
l SP (Coord (Dimension core)) (NumType core)
_ Box (Dimension core) () (NumType core)
r -> [Box (Dimension core) () (NumType core)]
-> Box
     (Dimension (Box (Dimension core) () (NumType core)))
     ()
     (NumType (Box (Dimension core) () (NumType core)))
forall g (c :: * -> *).
(IsBoxable g, Foldable c, Ord (NumType g), Arity (Dimension g)) =>
c g -> Box (Dimension g) () (NumType g)
boundingBoxList' [Box (Dimension core) () (NumType core)
l,Box (Dimension core) () (NumType core)
r])
                                      (core -> Box (Dimension core) () (NumType core)
forall g.
(IsBoxable g, Ord (NumType g)) =>
g -> Box (Dimension g) () (NumType g)
boundingBox (core -> Box (Dimension core) () (NumType core))
-> ((core :+ extra) -> core)
-> (core :+ extra)
-> Box (Dimension core) () (NumType core)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((core :+ extra) -> Getting core (core :+ extra) core -> core
forall s a. s -> Getting a s a -> a
^.Getting core (core :+ extra) core
forall core extra core'.
Lens (core :+ extra) (core' :+ extra) core core'
core)) BinLeafTree
  (SP (Coord (Dimension core)) (NumType core)) (core :+ extra)
t
                 in (SP (Coord (Dimension core)) (NumType core)
 -> Box (Dimension core) () (NumType core)
 -> Split (Dimension core) (NumType core))
-> ((core :+ extra) -> (core :+ extra) -> core :+ extra)
-> BinLeafTree
     (SP (Coord (Dimension core)) (NumType core)) (core :+ extra)
-> BinLeafTree
     (Box (Dimension core) () (NumType core)) (core :+ extra)
-> BinLeafTree
     (Split (Dimension core) (NumType core)) (core :+ extra)
forall u v w a b c.
(u -> v -> w)
-> (a -> b -> c)
-> BinLeafTree u a
-> BinLeafTree v b
-> BinLeafTree w c
zipExactWith (\(SP Coord (Dimension core)
c NumType core
m) Box (Dimension core) () (NumType core)
b -> Coord (Dimension core)
-> NumType core
-> Box (Dimension core) () (NumType core)
-> Split (Dimension core) (NumType core)
forall (d :: Nat) r. Coord d -> r -> Box d () r -> Split d r
Split Coord (Dimension core)
c NumType core
m Box (Dimension core) () (NumType core)
b) (core :+ extra) -> (core :+ extra) -> core :+ extra
forall a b. a -> b -> a
const BinLeafTree
  (SP (Coord (Dimension core)) (NumType core)) (core :+ extra)
t BinLeafTree
  (Box (Dimension core) () (NumType core)) (core :+ extra)
bbt


-- | Nub by sorting first
ordNub :: Ord a => NonEmpty.NonEmpty a -> NonEmpty.NonEmpty a
ordNub :: NonEmpty a -> NonEmpty a
ordNub = (NonEmpty a -> a) -> NonEmpty (NonEmpty a) -> NonEmpty a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap NonEmpty a -> a
forall a. NonEmpty a -> a
NonEmpty.head (NonEmpty (NonEmpty a) -> NonEmpty a)
-> (NonEmpty a -> NonEmpty (NonEmpty a))
-> NonEmpty a
-> NonEmpty a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty a -> NonEmpty (NonEmpty a)
forall a. Eq a => NonEmpty a -> NonEmpty (NonEmpty a)
NonEmpty.group1 (NonEmpty a -> NonEmpty (NonEmpty a))
-> (NonEmpty a -> NonEmpty a)
-> NonEmpty a
-> NonEmpty (NonEmpty a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty a -> NonEmpty a
forall a. Ord a => NonEmpty a -> NonEmpty a
NonEmpty.sort



toPointSet :: (Arity d, Ord r)
           => LSeq n (Point d r :+ p) -> PointSet (LSeq n) d p r
toPointSet :: LSeq n (Point d r :+ p) -> PointSet (LSeq n) d p r
toPointSet = (Int -> LSeq n (Point d r :+ p) -> LSeq n (Point d r :+ p))
-> PointSet (LSeq n) d p r -> PointSet (LSeq n) d p r
forall (v :: * -> *) a b.
(Vector v a, Vector v b) =>
(Int -> a -> b) -> v a -> v b
FV.imap Int -> LSeq n (Point d r :+ p) -> LSeq n (Point d r :+ p)
forall (d :: Nat) r (n :: Nat) e.
(ImplicitPeano (Peano d), Ord r,
 ArityPeano (Peano (FromPeano (Peano d))),
 KnownNat (FromPeano (Peano d)), KnownNat d,
 Peano (FromPeano (Peano d) + 1)
 ~ 'S (Peano (FromPeano (Peano d)))) =>
Int -> LSeq n (Point d r :+ e) -> LSeq n (Point d r :+ e)
sort (PointSet (LSeq n) d p r -> PointSet (LSeq n) d p r)
-> (LSeq n (Point d r :+ p) -> PointSet (LSeq n) d p r)
-> LSeq n (Point d r :+ p)
-> PointSet (LSeq n) d p r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LSeq n (Point d r :+ p) -> PointSet (LSeq n) d p r
forall (v :: * -> *) a. Vector v a => a -> v a
FV.replicate
  where
    sort :: Int -> LSeq n (Point d r :+ e) -> LSeq n (Point d r :+ e)
sort Int
i = ((Point d r :+ e) -> (Point d r :+ e) -> Ordering)
-> LSeq n (Point d r :+ e) -> LSeq n (Point d r :+ e)
forall a (n :: Nat). (a -> a -> Ordering) -> LSeq n a -> LSeq n a
LSeq.unstableSortBy (Int -> (Point d r :+ e) -> (Point d r :+ e) -> Ordering
forall r (d :: Nat) e.
(Ord r, Arity d) =>
Int -> (Point d r :+ e) -> (Point d r :+ e) -> Ordering
compareOn (Int -> (Point d r :+ e) -> (Point d r :+ e) -> Ordering)
-> Int -> (Point d r :+ e) -> (Point d r :+ e) -> Ordering
forall a b. (a -> b) -> a -> b
$ Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i)


compareOn       :: (Ord r, Arity d)
                => Int -> Point d r :+ e -> Point d r :+ e -> Ordering
compareOn :: Int -> (Point d r :+ e) -> (Point d r :+ e) -> Ordering
compareOn Int
i Point d r :+ e
p Point d r :+ e
q = let f :: (Point d r :+ e) -> r
f = ((Point d r :+ e) -> Getting r (Point d r :+ e) r -> r
forall s a. s -> Getting a s a -> a
^.(Point d r -> Const r (Point d r))
-> (Point d r :+ e) -> Const r (Point d r :+ e)
forall core extra core'.
Lens (core :+ extra) (core' :+ extra) core core'
core((Point d r -> Const r (Point d r))
 -> (Point d r :+ e) -> Const r (Point d r :+ e))
-> ((r -> Const r r) -> Point d r -> Const r (Point d r))
-> Getting r (Point d r :+ e) r
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Int -> Lens' (Point d r) r
forall (d :: Nat) (p :: Nat -> * -> *) r.
(Arity d, AsAPoint p) =>
Int -> Lens' (p d r) r
unsafeCoord Int
i)
                  in ((Point d r :+ e) -> r
f Point d r :+ e
p, Point d r :+ e
p(Point d r :+ e)
-> Getting (Point d r) (Point d r :+ e) (Point d r) -> Point d r
forall s a. s -> Getting a s a -> a
^.Getting (Point d r) (Point d r :+ e) (Point d r)
forall core extra core'.
Lens (core :+ extra) (core' :+ extra) core core'
core) (r, Point d r) -> (r, Point d r) -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` ((Point d r :+ e) -> r
f Point d r :+ e
q, Point d r :+ e
q(Point d r :+ e)
-> Getting (Point d r) (Point d r :+ e) (Point d r) -> Point d r
forall s a. s -> Getting a s a -> a
^.Getting (Point d r) (Point d r :+ e) (Point d r)
forall core extra core'.
Lens (core :+ extra) (core' :+ extra) core core'
core)


build      :: (1 <= d, Arity d, Ord r)
           => Coord d
           -> PointSet (LSeq 1) d p r
           -> BinLeafTree (Split' d r) (Point d r :+ p)
build :: Coord d
-> PointSet (LSeq 1) d p r
-> BinLeafTree (Split' d r) (Point d r :+ p)
build Coord d
i PointSet (LSeq 1) d p r
ps = case PointSet (LSeq 1) d p r
-> Either (Point d r :+ p) (PointSet (LSeq 2) d p r)
forall (d :: Nat) p r.
(1 <= d, Arity d) =>
PointSet (LSeq 1) d p r
-> Either (Point d r :+ p) (PointSet (LSeq 2) d p r)
asSingleton PointSet (LSeq 1) d p r
ps of
    Left Point d r :+ p
p    -> (Point d r :+ p) -> BinLeafTree (Split' d r) (Point d r :+ p)
forall v a. a -> BinLeafTree v a
Leaf Point d r :+ p
p
    Right PointSet (LSeq 2) d p r
ps' -> let (PointSet (LSeq 1) d p r
l,Split' d r
m,PointSet (LSeq 1) d p r
r) = Coord d
-> PointSet (LSeq 2) d p r
-> (PointSet (LSeq 1) d p r, Split' d r, PointSet (LSeq 1) d p r)
forall (d :: Nat) r p.
(Arity d, KnownNat d, Ord r) =>
Coord d
-> PointSet (LSeq 2) d p r
-> (PointSet (LSeq 1) d p r, Split' d r, PointSet (LSeq 1) d p r)
splitOn Coord d
i PointSet (LSeq 2) d p r
ps'
                     j :: Coord d
j       = Coord d -> Coord d
forall a. Enum a => a -> a
succ Coord d
i
                   -- the pattern match proves tha tthe seq has >= 2 elements
                 in BinLeafTree (Split' d r) (Point d r :+ p)
-> Split' d r
-> BinLeafTree (Split' d r) (Point d r :+ p)
-> BinLeafTree (Split' d r) (Point d r :+ p)
forall v a.
BinLeafTree v a -> v -> BinLeafTree v a -> BinLeafTree v a
Node (Coord d
-> PointSet (LSeq 1) d p r
-> BinLeafTree (Split' d r) (Point d r :+ p)
forall (d :: Nat) r p.
(1 <= d, Arity d, Ord r) =>
Coord d
-> PointSet (LSeq 1) d p r
-> BinLeafTree (Split' d r) (Point d r :+ p)
build Coord d
j PointSet (LSeq 1) d p r
l) Split' d r
m (Coord d
-> PointSet (LSeq 1) d p r
-> BinLeafTree (Split' d r) (Point d r :+ p)
forall (d :: Nat) r p.
(1 <= d, Arity d, Ord r) =>
Coord d
-> PointSet (LSeq 1) d p r
-> BinLeafTree (Split' d r) (Point d r :+ p)
build Coord d
j PointSet (LSeq 1) d p r
r)


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

reportSubTree :: KDTree' d p r -> NonEmpty.NonEmpty (Point d r :+ p)
reportSubTree :: KDTree' d p r -> NonEmpty (Point d r :+ p)
reportSubTree = [Point d r :+ p] -> NonEmpty (Point d r :+ p)
forall a. [a] -> NonEmpty a
NonEmpty.fromList ([Point d r :+ p] -> NonEmpty (Point d r :+ p))
-> (KDTree' d p r -> [Point d r :+ p])
-> KDTree' d p r
-> NonEmpty (Point d r :+ p)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BinLeafTree (Split d r) (Point d r :+ p) -> [Point d r :+ p]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList (BinLeafTree (Split d r) (Point d r :+ p) -> [Point d r :+ p])
-> (KDTree' d p r -> BinLeafTree (Split d r) (Point d r :+ p))
-> KDTree' d p r
-> [Point d r :+ p]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KDTree' d p r -> BinLeafTree (Split d r) (Point d r :+ p)
forall (d :: Nat) p r.
KDTree' d p r -> BinLeafTree (Split d r) (Point d r :+ p)
unKDT

-- | Searches in a KDTree
--
-- running time: \(O(n^{(d-1)/d} + k)\)
searchKDTree    :: (Arity d, Ord r)
                => Box d q r -> KDTree d p r -> [Point d r :+ p]
searchKDTree :: Box d q r -> KDTree d p r -> [Point d r :+ p]
searchKDTree Box d q r
qr = [Point d r :+ p]
-> (KDTree' d p r -> [Point d r :+ p])
-> Maybe (KDTree' d p r)
-> [Point d r :+ p]
forall b a. b -> (a -> b) -> Maybe a -> b
maybe [] (Box d q r -> KDTree' d p r -> [Point d r :+ p]
forall (d :: Nat) r q p.
(Arity d, Ord r) =>
Box d q r -> KDTree' d p r -> [Point d r :+ p]
searchKDTree' Box d q r
qr) (Maybe (KDTree' d p r) -> [Point d r :+ p])
-> (KDTree d p r -> Maybe (KDTree' d p r))
-> KDTree d p r
-> [Point d r :+ p]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KDTree d p r -> Maybe (KDTree' d p r)
forall (d :: Nat) p r. KDTree d p r -> Maybe (KDTree' d p r)
toMaybe

searchKDTree'                  :: (Arity d, Ord r)
                              => Box d q r -> KDTree' d p r -> [Point d r :+ p]
searchKDTree' :: Box d q r -> KDTree' d p r -> [Point d r :+ p]
searchKDTree' Box d q r
qr = BinLeafTree (Split d r) (Point d r :+ p) -> [Point d r :+ p]
search (BinLeafTree (Split d r) (Point d r :+ p) -> [Point d r :+ p])
-> (KDTree' d p r -> BinLeafTree (Split d r) (Point d r :+ p))
-> KDTree' d p r
-> [Point d r :+ p]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. KDTree' d p r -> BinLeafTree (Split d r) (Point d r :+ p)
forall (d :: Nat) p r.
KDTree' d p r -> BinLeafTree (Split d r) (Point d r :+ p)
unKDT
  where
    search :: BinLeafTree (Split d r) (Point d r :+ p) -> [Point d r :+ p]
search (Leaf Point d r :+ p
p)
      | (Point d r :+ p
p(Point d r :+ p)
-> Getting (Point d r) (Point d r :+ p) (Point d r) -> Point d r
forall s a. s -> Getting a s a -> a
^.Getting (Point d r) (Point d r :+ p) (Point d r)
forall core extra core'.
Lens (core :+ extra) (core' :+ extra) core core'
core) Point d r -> Box d q r -> Bool
forall g h. IsIntersectableWith g h => g -> h -> Bool
`intersects` Box d q r
qr = [Point d r :+ p
p]
      | Bool
otherwise                 = []
    search t :: BinLeafTree (Split d r) (Point d r :+ p)
t@(Node BinLeafTree (Split d r) (Point d r :+ p)
l (Split Coord d
_ r
_ Box d () r
b) BinLeafTree (Split d r) (Point d r :+ p)
r)
      | Box d () r
b Box d () r -> Box d q r -> Bool
forall (d :: Nat) r q p.
(Arity d, Ord r) =>
Box d q r -> Box d p r -> Bool
`containedIn` Box d q r
qr        = BinLeafTree (Split d r) (Point d r :+ p) -> [Point d r :+ p]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList BinLeafTree (Split d r) (Point d r :+ p)
t
      | Bool
otherwise                 = [Point d r :+ p]
l' [Point d r :+ p] -> [Point d r :+ p] -> [Point d r :+ p]
forall a. [a] -> [a] -> [a]
++ [Point d r :+ p]
r'
      where
        l' :: [Point d r :+ p]
l' = if Box d q r
qr Box d q r -> Box d () r -> Bool
forall g h. IsIntersectableWith g h => g -> h -> Bool
`intersects` BinLeafTree (Split d r) (Point d r :+ p) -> Box d () r
forall (d :: Nat) r p.
(Arity d, Ord r) =>
BinLeafTree (Split d r) (Point d r :+ p) -> Box d () r
boxOf BinLeafTree (Split d r) (Point d r :+ p)
l then BinLeafTree (Split d r) (Point d r :+ p) -> [Point d r :+ p]
search BinLeafTree (Split d r) (Point d r :+ p)
l else []
        r' :: [Point d r :+ p]
r' = if Box d q r
qr Box d q r -> Box d () r -> Bool
forall g h. IsIntersectableWith g h => g -> h -> Bool
`intersects` BinLeafTree (Split d r) (Point d r :+ p) -> Box d () r
forall (d :: Nat) r p.
(Arity d, Ord r) =>
BinLeafTree (Split d r) (Point d r :+ p) -> Box d () r
boxOf BinLeafTree (Split d r) (Point d r :+ p)
r then BinLeafTree (Split d r) (Point d r :+ p) -> [Point d r :+ p]
search BinLeafTree (Split d r) (Point d r :+ p)
r else []


boxOf :: (Arity d, Ord r) => BinLeafTree (Split d r) (Point d r :+ p) -> Box d () r
boxOf :: BinLeafTree (Split d r) (Point d r :+ p) -> Box d () r
boxOf (Leaf Point d r :+ p
p)                 = Point d r -> Box (Dimension (Point d r)) () (NumType (Point d r))
forall g.
(IsBoxable g, Ord (NumType g)) =>
g -> Box (Dimension g) () (NumType g)
boundingBox (Point d r :+ p
p(Point d r :+ p)
-> Getting (Point d r) (Point d r :+ p) (Point d r) -> Point d r
forall s a. s -> Getting a s a -> a
^.Getting (Point d r) (Point d r :+ p) (Point d r)
forall core extra core'.
Lens (core :+ extra) (core' :+ extra) core core'
core)
boxOf (Node BinLeafTree (Split d r) (Point d r :+ p)
_ (Split Coord d
_ r
_ Box d () r
b) BinLeafTree (Split d r) (Point d r :+ p)
_) = Box d () r
b

containedIn :: (Arity d, Ord r) => Box d q r -> Box d p r -> Bool
(Box (CWMin Point d r
p :+ q
_) (CWMax Point d r
q :+ q
_)) containedIn :: Box d q r -> Box d p r -> Bool
`containedIn` Box d p r
b = (Point d r -> Bool) -> [Point d r] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Point d r -> Box d p r -> Bool
forall g h. IsIntersectableWith g h => g -> h -> Bool
`intersects` Box d p r
b) [Point d r
p,Point d r
q]

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


type PointSet seq d p r = Vector d (seq (Point d r :+ p))

-- | running time: \(O(n)\)
splitOn                 :: (Arity d, KnownNat d, Ord r)
                        => Coord d
                        -> PointSet (LSeq 2) d p r
                        -> ( PointSet (LSeq 1) d p r
                           , Split' d r
                           , PointSet (LSeq 1) d p r)
splitOn :: Coord d
-> PointSet (LSeq 2) d p r
-> (PointSet (LSeq 1) d p r, Split' d r, PointSet (LSeq 1) d p r)
splitOn c :: Coord d
c@(Coord Int
i) PointSet (LSeq 2) d p r
pts = (PointSet (LSeq 1) d p r
l, Coord d -> r -> Split' d r
forall a b. a -> b -> SP a b
SP Coord d
c (Point d r :+ p
m(Point d r :+ p) -> Getting r (Point d r :+ p) r -> r
forall s a. s -> Getting a s a -> a
^.(Point d r -> Const r (Point d r))
-> (Point d r :+ p) -> Const r (Point d r :+ p)
forall core extra core'.
Lens (core :+ extra) (core' :+ extra) core core'
core((Point d r -> Const r (Point d r))
 -> (Point d r :+ p) -> Const r (Point d r :+ p))
-> ((r -> Const r r) -> Point d r -> Const r (Point d r))
-> Getting r (Point d r :+ p) r
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Int -> Lens' (Point d r) r
forall (d :: Nat) (p :: Nat -> * -> *) r.
(Arity d, AsAPoint p) =>
Int -> Lens' (p d r) r
unsafeCoord Int
i), PointSet (LSeq 1) d p r
r)
  where
    -- i = traceShow (c,j) j

    m :: Point d r :+ p
m = let xs :: LSeq 2 (Point d r :+ p)
xs = PointSet (LSeq 2) d p r
ptsPointSet (LSeq 2) d p r
-> Getting
     (Endo (LSeq 2 (Point d r :+ p)))
     (PointSet (LSeq 2) d p r)
     (LSeq 2 (Point d r :+ p))
-> LSeq 2 (Point d r :+ p)
forall s a. HasCallStack => s -> Getting (Endo a) s a -> a
^?!Int
-> Traversal' (PointSet (LSeq 2) d p r) (LSeq 2 (Point d r :+ p))
forall (d :: Nat) r. Arity d => Int -> Traversal' (Vector d r) r
element' (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
        in LSeq 2 (Point d r :+ p)
xs LSeq 2 (Point d r :+ p) -> Int -> Point d r :+ p
forall (n :: Nat) a. LSeq n a -> Int -> a
`LSeq.index` (LSeq 2 (Point d r :+ p) -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
F.length LSeq 2 (Point d r :+ p)
xs Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)

    -- Since the input seq has >= 2 elems, F.length xs / 2 >= 1. It follows
    -- that the both sets thus have at least one elemnt.
    -- f :: LSeq 2 _ -> (LSeq 1 _, LSeq 1 _)
    f :: LSeq 2 (Point d r :+ p)
-> (LSeq 1 (Point d r :+ p), LSeq 1 (Point d r :+ p))
f = (LSeq 0 (Point d r :+ p) -> LSeq 1 (Point d r :+ p))
-> (LSeq 0 (Point d r :+ p) -> LSeq 1 (Point d r :+ p))
-> (LSeq 0 (Point d r :+ p), LSeq 0 (Point d r :+ p))
-> (LSeq 1 (Point d r :+ p), LSeq 1 (Point d r :+ p))
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap LSeq 0 (Point d r :+ p) -> LSeq 1 (Point d r :+ p)
forall (m :: Nat) (n :: Nat) a. LSeq m a -> LSeq n a
LSeq.promise LSeq 0 (Point d r :+ p) -> LSeq 1 (Point d r :+ p)
forall (m :: Nat) (n :: Nat) a. LSeq m a -> LSeq n a
LSeq.promise
      ((LSeq 0 (Point d r :+ p), LSeq 0 (Point d r :+ p))
 -> (LSeq 1 (Point d r :+ p), LSeq 1 (Point d r :+ p)))
-> (LSeq 2 (Point d r :+ p)
    -> (LSeq 0 (Point d r :+ p), LSeq 0 (Point d r :+ p)))
-> LSeq 2 (Point d r :+ p)
-> (LSeq 1 (Point d r :+ p), LSeq 1 (Point d r :+ p))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Point d r :+ p) -> Bool)
-> LSeq 2 (Point d r :+ p)
-> (LSeq 0 (Point d r :+ p), LSeq 0 (Point d r :+ p))
forall a (n :: Nat).
(a -> Bool) -> LSeq n a -> (LSeq 0 a, LSeq 0 a)
LSeq.partition (\Point d r :+ p
p -> Int -> (Point d r :+ p) -> (Point d r :+ p) -> Ordering
forall r (d :: Nat) e.
(Ord r, Arity d) =>
Int -> (Point d r :+ e) -> (Point d r :+ e) -> Ordering
compareOn Int
i Point d r :+ p
p Point d r :+ p
m Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
LT)

    (PointSet (LSeq 1) d p r
l,PointSet (LSeq 1) d p r
r) = Vector d (LSeq 1 (Point d r :+ p), LSeq 1 (Point d r :+ p))
-> (PointSet (LSeq 1) d p r, PointSet (LSeq 1) d p r)
forall r r. Vector d (r, r) -> (Vector d r, Vector d r)
unzip' (Vector d (LSeq 1 (Point d r :+ p), LSeq 1 (Point d r :+ p))
 -> (PointSet (LSeq 1) d p r, PointSet (LSeq 1) d p r))
-> (PointSet (LSeq 2) d p r
    -> Vector d (LSeq 1 (Point d r :+ p), LSeq 1 (Point d r :+ p)))
-> PointSet (LSeq 2) d p r
-> (PointSet (LSeq 1) d p r, PointSet (LSeq 1) d p r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (LSeq 2 (Point d r :+ p)
 -> (LSeq 1 (Point d r :+ p), LSeq 1 (Point d r :+ p)))
-> PointSet (LSeq 2) d p r
-> Vector d (LSeq 1 (Point d r :+ p), LSeq 1 (Point d r :+ p))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap LSeq 2 (Point d r :+ p)
-> (LSeq 1 (Point d r :+ p), LSeq 1 (Point d r :+ p))
f (PointSet (LSeq 2) d p r
 -> (PointSet (LSeq 1) d p r, PointSet (LSeq 1) d p r))
-> PointSet (LSeq 2) d p r
-> (PointSet (LSeq 1) d p r, PointSet (LSeq 1) d p r)
forall a b. (a -> b) -> a -> b
$ PointSet (LSeq 2) d p r
pts

    -- unzip' :: Vector d (a,b) -> (Vector d a, Vector d b)
    unzip' :: Vector d (r, r) -> (Vector d r, Vector d r)
unzip' = ([r] -> Vector d r)
-> ([r] -> Vector d r) -> ([r], [r]) -> (Vector d r, Vector d r)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap [r] -> Vector d r
forall (d :: Nat) r. Arity d => [r] -> Vector d r
vectorFromListUnsafe [r] -> Vector d r
forall (d :: Nat) r. Arity d => [r] -> Vector d r
vectorFromListUnsafe (([r], [r]) -> (Vector d r, Vector d r))
-> (Vector d (r, r) -> ([r], [r]))
-> Vector d (r, r)
-> (Vector d r, Vector d r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(r, r)] -> ([r], [r])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(r, r)] -> ([r], [r]))
-> (Vector d (r, r) -> [(r, r)]) -> Vector d (r, r) -> ([r], [r])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector d (r, r) -> [(r, r)]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList


asSingleton   :: (1 <= d, Arity d)
              => PointSet (LSeq 1) d p r
              -> Either (Point d r :+ p) (PointSet (LSeq 2) d p r)
asSingleton :: PointSet (LSeq 1) d p r
-> Either (Point d r :+ p) (PointSet (LSeq 2) d p r)
asSingleton PointSet (LSeq 1) d p r
v = case PointSet (LSeq 1) d p r
vPointSet (LSeq 1) d p r
-> Getting
     (LSeq 1 (Point d r :+ p))
     (PointSet (LSeq 1) d p r)
     (LSeq 1 (Point d r :+ p))
-> LSeq 1 (Point d r :+ p)
forall s a. s -> Getting a s a -> a
^.C 0 -> Lens' (PointSet (LSeq 1) d p r) (LSeq 1 (Point d r :+ p))
forall (proxy :: Nat -> *) (i :: Nat) (d :: Nat) r.
(Arity d, KnownNat i, (i + 1) <= d) =>
proxy i -> Lens' (Vector d r) r
element (C 0
forall (n :: Nat). C n
C :: C 0) of
                  (p :<| s) | LSeq 0 (Point d r :+ p) -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null LSeq 0 (Point d r :+ p)
s -> (Point d r :+ p)
-> Either (Point d r :+ p) (PointSet (LSeq 2) d p r)
forall a b. a -> Either a b
Left Point d r :+ p
p -- only one lement
                  LSeq 1 (Point d r :+ p)
_                  -> PointSet (LSeq 2) d p r
-> Either (Point d r :+ p) (PointSet (LSeq 2) d p r)
forall a b. b -> Either a b
Right (PointSet (LSeq 2) d p r
 -> Either (Point d r :+ p) (PointSet (LSeq 2) d p r))
-> PointSet (LSeq 2) d p r
-> Either (Point d r :+ p) (PointSet (LSeq 2) d p r)
forall a b. (a -> b) -> a -> b
$ PointSet (LSeq 1) d p r -> PointSet (LSeq 2) d p r
forall a b. a -> b
unsafeCoerce PointSet (LSeq 1) d p r
v