diff --git a/Changelog.md b/Changelog.md
--- a/Changelog.md
+++ b/Changelog.md
@@ -1,3 +1,8 @@
+0.4
+
+- add function to compute the forest construction parameters from the dataset dimensions (rpTreeCfg)
+- add type parameter to RPTree to label tree branching points
+
 0.3.6
 
 - relax lower dependency bounds to ensure it builds with GHC 8.6.5 (== stackage lts 14.27) as well
diff --git a/app/Main.hs b/app/Main.hs
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -4,13 +4,15 @@
 module Main where
 
 import Control.Monad (replicateM)
-import Data.Foldable (fold)
+import Data.Bitraversable (Bitraversable(..))
+import Data.Foldable (fold, toList)
 import Data.Functor (void)
+import GHC.Stack (HasCallStack)
 
 -- conduit
 import qualified Data.Conduit as C (ConduitT, runConduit, yield, await, transPipe)
 import Data.Conduit ((.|))
-import qualified Data.Conduit.Combinators as C (map, mapM, scanl, scanlM, last, print)
+import qualified Data.Conduit.Combinators as C (map, mapM, scanl, scanlM, last, print, sinkVector, sinkList)
 import qualified Data.Conduit.List as C (chunksOf, unfold, unfoldM)
 -- containers
 import qualified Data.IntMap as IM (IntMap, fromList, insert, lookup, map, mapWithKey, traverseWithKey, foldlWithKey, foldrWithKey)
@@ -21,71 +23,61 @@
 -- splitmix-distributions
 import System.Random.SplitMix.Distributions (Gen, GenT, sample, sampleT, bernoulli, normal)
 -- transformers
-import Control.Monad.Trans.State.Lazy (State, get, put, evalState)
+import Control.Monad.Trans.State.Strict (State, get, put, evalState)
 import Control.Monad.Trans.Class (MonadTrans(..))
 -- vector
-import qualified Data.Vector as V (Vector, toList, fromList, replicate, zip)
+import qualified Data.Vector as V (Vector, toList, fromList, replicate, zip, zipWith)
 
 import Control.Monad (replicateM)
-import Data.RPTree (knn, candidates, Embed(..), Inner(..), RPTree, RPForest, SVector, fromListSv, DVector, fromListDv, dense, writeCsv, forest, dataSource, sparse, normal2, normalSparse2, datS, datD)
+import Data.RPTree (knn, candidates, rpTreeCfg, RPTreeConfig(..), Embed(..), Inner(..), RPTree, RPForest, SVector, fromListSv, DVector, fromListDv, dense, writeCsv, tree, forest, dataSource, sparse, normal2, normalSparse2, datS, datD, circle2d, leaves, levels, treeSize, leafSizes, writeDot)
 -- import Data.RPTree.Internal.Testing (datS, datD)
 
 main :: IO ()
-main = do -- putStrLn "hello!"
+main = do
   let
-    n = 1000
-    maxd = 3
+    n = 30000
+    maxd = 5
     minl = 10
-    ntree = 10
-    d = 100
-    pnz = 0.3
-    chunk = 20
-    src = datS n d pnz .| C.map (\ x -> Embed x ())
-    seed = 1234
-  (q, tts) <- sampleT seed $ do
-    tts <- C.runConduit $
-           forest seed maxd minl ntree chunk pnz d (liftC src)
-    q <- sparse 0.3 d (normal 0.1 0.6)
-    pure (q, tts)
-  let
-    res = knn metricL2 1 tts q
-  print res
-
-
-liftC = C.transPipe lift
+    chunk = 500
+    dim = 2
+    -- cfg = rpTreeCfg n dim
+    cfg = RPCfg maxd minl 3 chunk 1.0
+  csvTree0 n cfg
+  tree0dot n cfg
 
 
 
--- renderTree0 :: Int -> IO ()
-renderTree0 tt = do
-  let csvrows = V.toList $ fold $ flip evalState A $ traverse labeledV tt -- (tree0 n)
-  writeCsv "r/scatter_data.csv" csvrows
+tree0dot :: Int -> RPTreeConfig -> IO ()
+tree0dot n (RPCfg maxd minl _ chunk _) =
+  writeDot f fpath "tree0" $ tree0 n maxd minl chunk
+  where
+    f = show . length
+    fpath = "tree0.dot"
 
--- renderTree1 :: Int -> IO ()
-renderTree1 tt = do
+csvTree0 :: Int -> RPTreeConfig -> IO ()
+csvTree0 n (RPCfg maxd minl _ chunk _) = do
   let
-    -- csvrows :: [(DVector Double, Pal5)]
-    csvrows = fold $ flip evalState A $ traverse labeledV tt -- (tree1 n)
-  writeCsv "r/scatter_data_rt2.csv" $ V.toList csvrows
-
-labeled :: (Enum i) =>
-           [a] -> State i [(a, i)]
-labeled xs = do
-  i <- get
-  put (succ i)
-  let n = length xs
-  pure $ zip xs (replicate n i)
+    tt = tree0 n maxd minl chunk
+    ttlab = prep tt
+  writeCsv "r/scatter_data_2.csv" ttlab
 
+prep :: (Traversable t) => t (V.Vector (Embed v e a)) -> t (V.Vector (v e, Pal))
+prep = flip evalState A . traverse labeled
 
-labeledV :: Enum i => V.Vector a -> State i (V.Vector (a, i))
-labeledV xs = do
+labeled :: (Enum b) =>
+             V.Vector (Embed v e a)
+          -> State b (V.Vector (v e, b))
+labeled xs = do
   i <- get
   put (succ i)
-  let n = length xs
-  pure $ V.zip xs (V.replicate n i)
+  let
+    n = length xs
+    f (Embed x _) ii = (x, ii)
+  pure $ V.zipWith f xs (V.replicate n i)
 
-data Pal5 = A | B | C | D | E deriving (Eq, Show)
-instance Enum Pal5 where
+-- color palette
+data Pal = A | B | C | D | E deriving (Eq, Show)
+instance Enum Pal where
   toEnum = \case
     0 -> A
     1 -> B
@@ -100,20 +92,80 @@
     D -> 3
     E -> 4
 
--- tree0 :: Int -> RPTree Double (V.Vector (DVector Double))
--- tree0 n = evalGen 1234 $ tree 10 1.0 2 (dataset n)
 
--- tree1 :: Int -> RT SVector Double (V.Vector (DVector Double))
--- tree1 n = evalGen 1234 $ treeRT 10 20 1.0 2 (dataset n)
 
+tree0 :: Int -- ^ dataset size
+      -> Int -- ^ max tree depth
+      -> Int -- ^ min leaf size
+      -> Int -- ^ chunk size
+      -> RPTree Double () (V.Vector (Embed DVector Double ()))
+tree0 n maxd minl chunk = sample s $ tree s maxd minl chunk 1.0 2 (srcC n .| embedC)
+  where
+    s = 1235137
+
 dataset :: Int -> V.Vector (DVector Double)
 dataset n = V.fromList $ sample 1234 $ replicateM n (dense 2 $ normal 0 1)
 
-
+datasetCircles :: Int -> V.Vector (DVector Double)
+datasetCircles n = V.fromList $ sample 1234 $ C.runConduit $ srcCircles n .| C.sinkList
 
 
 
 srcC :: Monad m => Int -> C.ConduitT i (DVector Double) (GenT m) ()
 srcC n = dataSource n normal2
 
+srcCircles :: Monad m =>
+              Int -> C.ConduitT i (DVector Double) (GenT m) ()
+srcCircles n = dataSource n circle2d2
 
+-- binary mixture of two non-overlapping circles
+circle2d2 :: (Monad m) => GenT m (DVector Double)
+circle2d2 = do
+  let
+    d = fromListDv [2, 3]
+    r = 1
+  b <- bernoulli 0.5
+  if b
+    then circle2d r
+    else (^+^ d) <$> circle2d r
+
+
+
+
+
+-- main :: IO ()
+-- main = do -- putStrLn "hello!"
+--   let
+--     n = 1000
+--     maxd = 3
+--     minl = 10
+--     ntree = 10
+--     d = 100
+--     pnz = 0.3
+--     chunk = 20
+--     src = datS n d pnz .| C.map (\ x -> Embed x ())
+--     -- src = srcCircles n
+--     seed = 1234
+--   (q, tts) <- sampleT seed $ do
+--     tts <- C.runConduit $
+--            forest seed maxd minl ntree chunk pnz d (liftC src)
+--     q <- sparse 0.3 d (normal 0.1 0.6)
+--     pure (q, tts)
+--   let
+--     res = knn (flip metricL2) 1 tts q
+--   print res
+
+
+-- liftC = C.transPipe lift
+
+embedC :: Monad m => C.ConduitT (v e) (Embed v e ()) m ()
+embedC = C.map (\ x -> Embed x ())
+
+
+
+-- -- renderTree1 :: Int -> IO ()
+-- renderTree1 tt = do
+--   let
+--     -- csvrows :: [(DVector Double, Pal5)]
+--     csvrows = fold $ flip evalState A $ traverse labeledV tt -- (tree1 n)
+--   writeCsv "r/scatter_data_rt2.csv" $ V.toList csvrows
diff --git a/r/scatter.png b/r/scatter.png
deleted file mode 100644
Binary files a/r/scatter.png and /dev/null differ
diff --git a/rp-tree.cabal b/rp-tree.cabal
--- a/rp-tree.cabal
+++ b/rp-tree.cabal
@@ -1,5 +1,5 @@
 name:                rp-tree
-version:             0.3.6
+version:             0.4
 synopsis:            Random projection trees
 description:         Random projection trees for approximate nearest neighbor search in high-dimensional vector spaces
                      .
@@ -9,12 +9,12 @@
 license-file:        LICENSE
 author:              Marco Zocca
 maintainer:          ocramz
-copyright:           2021 Marco Zocca
+copyright:           (c) 2021 Marco Zocca, UnfoldML AB
 category:            Data Mining, Data Structures, Machine Learning, Data
 build-type:          Simple
 extra-source-files:  README.md
                      Changelog.md
-extra-doc-files:     r/scatter.png
+-- extra-doc-files:     r/scatter.png
 cabal-version:       1.18
 tested-with:         GHC == 8.6.5, GHC == 8.10.4
 
@@ -35,9 +35,14 @@
                      , conduit >= 1.3.1
                      , containers >= 0.6
                      , deepseq >= 1.4
+                     , microlens
+                     , microlens-th
+                     , mtl >= 2.2.2
+                     , psqueues
                      , serialise >= 0.2
                      , splitmix
                      , splitmix-distributions >= 0.9
+                     , text >= 1.2.3.1
                      , transformers >= 0.5
                      , vector >= 0.12
                      , vector-algorithms >= 0.8
diff --git a/src/Data/RPTree.hs b/src/Data/RPTree.hs
--- a/src/Data/RPTree.hs
+++ b/src/Data/RPTree.hs
@@ -49,51 +49,60 @@
 -}
 module Data.RPTree (
   -- * Construction
-  -- tree,
-  forest
-  -- , defaultParams
+  tree
+  , forest
+  -- ** Parameters
+  , rpTreeCfg, RPTreeConfig(..)
   -- , ForestParams
   -- * Query
   , knn
+  -- , knnPQ
   -- * I/O
   , serialiseRPForest
   , deserialiseRPForest
-  -- * Validation
+  -- * Statistics
   , recallWith
   -- * Access
-  , levels, points, candidates
+  , leaves, levels, points, candidates
+  -- * Validation
+  , treeStats, treeSize, leafSizes
+  , RPTreeStats
   -- * Types
   , Embed(..)
   -- ** RPTree
   , RPTree, RPForest
-  -- *
+  -- * Vector types
+  -- ** Sparse
   , SVector, fromListSv, fromVectorSv
+  -- ** Dense
   , DVector, fromListDv, fromVectorDv
-  -- * Vector space types
+  -- * Vector space typeclasses
   , Inner(..), Scale(..)
-    -- ** helpers for implementing Inner instances
-    -- *** inner product
+    -- ** Helpers for implementing 'Inner' instances
+    -- *** Inner product
   , innerSS, innerSD, innerDD
     -- *** L2 distance
   , metricSSL2, metricSDL2
   -- *** Scale
   , scaleS, scaleD
 
-
   -- * Rendering
-  , draw
-  -- * CSV
+  -- , draw
+  -- ** CSV
   , writeCsv
+  -- ** GraphViz dot
+  , writeDot
   -- * Testing
-  , randSeed, BenchConfig(..), normalSparse2
+  , BenchConfig(..), normalSparse2
   , liftC
   -- ** Random generation
+  , randSeed
   -- *** Conduit
   , dataSource
   , datS, datD
   -- *** Vector data
   , sparse, dense
-  , normal2
+  , normal2, circle2d
   ) where
 
 import Control.Monad (replicateM)
@@ -104,17 +113,20 @@
 import Data.List (partition, sortBy)
 import Data.Monoid (Sum(..))
 import Data.Ord (comparing)
+import Data.Semigroup (Min(..))
 import GHC.Generics (Generic)
 import GHC.Word (Word64)
 
 -- containers
 import Data.Sequence (Seq, (|>))
-import qualified Data.Map as M (Map, fromList, toList, foldrWithKey, insert, insertWith)
+import qualified Data.Map as M (Map, fromList, toList, foldrWithKey, insert, insertWith, intersection)
 import qualified Data.Set as S (Set, fromList, intersection, insert)
 -- deepseq
 import Control.DeepSeq (NFData(..))
+-- psqueues
+import qualified Data.IntPSQ as PQ (IntPSQ, findMin, minView, empty, insert, fromList, toList)
 -- transformers
-import Control.Monad.Trans.State (StateT(..), runStateT, evalStateT, State, runState, evalState)
+import Control.Monad.Trans.State (StateT(..), runStateT, evalStateT, State, runState, evalState, get, put)
 import Control.Monad.Trans.Class (MonadTrans(..))
 -- vector
 import qualified Data.Vector as V (Vector, replicateM, fromList)
@@ -124,11 +136,11 @@
 -- vector-algorithms
 import qualified Data.Vector.Algorithms.Merge as V (sortBy)
 
-import Data.RPTree.Conduit (tree, forest, ForestParams, defaultParams, dataSource, liftC)
-import Data.RPTree.Gen (sparse, dense, normal2, normalSparse2)
-import Data.RPTree.Internal (RPTree(..), RPForest, RPT(..), Embed(..), levels, points, Inner(..), Scale(..), scaleS, scaleD, (/.), innerDD, innerSD, innerSS, metricSSL2, metricSDL2, SVector(..), fromListSv, fromVectorSv, DVector(..), fromListDv, fromVectorDv, partitionAtMedian, Margin, getMargin, sortByVG, serialiseRPForest, deserialiseRPForest)
+import Data.RPTree.Conduit (tree, forest, dataSource, liftC, rpTreeCfg, RPTreeConfig(..))
+import Data.RPTree.Gen (sparse, dense, normal2, normalSparse2, circle2d)
+import Data.RPTree.Internal (RPTree(..), RPForest, RPT(..), Embed(..), leaves, levels, points, Inner(..), Scale(..), scaleS, scaleD, (/.), innerDD, innerSD, innerSS, metricSSL2, metricSDL2, SVector(..), fromListSv, fromVectorSv, DVector(..), fromListDv, fromVectorDv, partitionAtMedian, Margin, getMargin, sortByVG, serialiseRPForest, deserialiseRPForest)
 import Data.RPTree.Internal.Testing (BenchConfig(..), randSeed, datS, datD)
-import Data.RPTree.Draw (draw, writeCsv)
+import Data.RPTree.Draw (writeDot, writeCsv)
 
 
 -- | Look up the \(k\) nearest neighbors to a query point
@@ -150,7 +162,21 @@
   where
     cs = VG.map (\xe -> (eEmbed xe `distf` q, xe)) $ VG.take k $ fold $ (`candidates` q) <$> tts
 
+-- | Same as 'knn' but with a (hopefully) faster implementation
+knnPQ :: (Ord p, Inner SVector v, VU.Unbox d, RealFrac d) =>
+         (u d -> v d -> p) -- ^ distance function
+      -> Int -- ^ k neighbors
+      -> RPForest d (V.Vector (Embed u d x)) -- ^ random projection forest
+      -> v d -- ^ query point
+      -> V.Vector (p, Embed u d x)
+knnPQ distf k tts q = sortByVG fst cs
+  where
+    cs = VG.map (\xe -> (eEmbed xe `distf` q, xe)) $ fold cstt
+    cstt = (takeFromPQ nsing) . (`candidatesPQ` q) <$> tts
+    nsing = (k `div` n) `max` 1
+    n = length tts
 
+
 -- | average recall-at-k, computed over a set of trees
 recallWith :: (Inner SVector v, VU.Unbox d, Fractional a1, Ord d, Ord a2, Ord x, Ord (u d), Num d) =>
               (u d -> v d -> a2)
@@ -165,7 +191,7 @@
 
 recallWith1 :: (Inner SVector v, Ord d, VU.Unbox d, Fractional p, Ord a, Ord x, Ord (u d), Num d) =>
               (u d -> v d -> a) -- ^ distance function
-           -> RPTree d (V.Vector (Embed u d x))
+           -> RPTree d l (V.Vector (Embed u d x))
            -> Int -- ^ k : number of nearest neighbors to consider
            -> v d -- ^ query point
            -> p
@@ -187,13 +213,13 @@
 --
 -- in case of a narrow margin, collect both branches of the tree
 candidates :: (Inner SVector v, VU.Unbox d, Ord d, Num d, Semigroup xs) =>
-              RPTree d xs
+              RPTree d l xs
            -> v d -- ^ query point
            -> xs
 candidates (RPTree rvs tt) x = go 0 tt
   where
-    go _     (Tip xs)                     = xs
-    go ixLev (Bin thr margin ltree rtree) = do
+    go _     (Tip _ xs)                     = xs
+    go ixLev (Bin _ thr margin ltree rtree) =
       let
         (mglo, mghi) = getMargin margin
         r = rvs VG.! ixLev
@@ -201,17 +227,93 @@
         i' = succ ixLev
         dl = abs (mglo - proj) -- left margin
         dr = abs (mghi - proj) -- right margin
+      in
+        if | proj < thr &&
+             dl > dr -> go i' ltree <> go i' rtree
+           | proj < thr  -> go i' ltree
+           | proj > thr &&
+             dl < dr -> go i' ltree <> go i' rtree
+           | otherwise   -> go i' rtree
+
+
+-- | like 'candidates' but outputs an ordered 'IntPQ' where the margin to the median projection is interpreted as queue priority
+candidatesPQ :: (Fractional d, Ord d, Inner SVector v, VU.Unbox d) =>
+               RPTree d l xs
+            -> v d -- ^ query point
+            -> PQ.IntPSQ d xs
+candidatesPQ (RPTree rvs tt) x = evalS $ go 0 tt PQ.empty (1/0)
+  where
+    go _ (Tip _ xs) acc dprev =
+      insPQ dprev xs acc
+    go ixLev (Bin _ thr margin ltree rtree) acc dprev = do
+      let
+        (mglo, mghi) = getMargin margin
+        r = rvs VG.! ixLev
+        proj = r `inner` x
+        i' = succ ixLev
+        dl = abs (mglo - proj) -- left margin
+        dr = abs (mghi - proj) -- right margin
       if | proj < thr &&
-           dl > dr -> go i' ltree <> go i' rtree
-         | proj < thr  -> go i' ltree
+           dl > dr -> do
+             ll <- go i' ltree acc (min dprev dl)
+             lr <- go i' rtree acc (min dprev dr)
+             pure $ PQ.fromList (PQ.toList ll <> PQ.toList lr)
+         | proj < thr  -> go i' ltree acc (min dprev dl)
          | proj > thr &&
-           dl < dr -> go i' ltree <> go i' rtree
-         | otherwise   -> go i' rtree
+           dl < dr -> do
+             ll <- go i' ltree acc (min dprev dl)
+             lr <- go i' rtree acc (min dprev dr)
+             pure $ PQ.fromList (PQ.toList ll <> PQ.toList lr)
+         | otherwise -> go i' rtree acc (min dprev dr)
 
+takeFromPQ :: (Ord p, Foldable t, Monoid (t a)) =>
+              Int -- ^ number of elements to keep
+           -> PQ.IntPSQ p (t a)
+           -> t a
+takeFromPQ n pq = foldMap snd $ reverse $ go [] 0 pq
+  where
+    go acc nacc q = case PQ.minView q of
+      Nothing -> acc
+      Just (_, p, xs, pqRest) ->
+        let
+          nxs = length xs
+          nacc' = nacc + nxs
+        in if nacc' < n
+           then go ((p, xs) : acc) nacc' pqRest
+           else acc
 
+type S = State Int
+evalS :: S a -> a
+evalS = flip evalState 0
 
+insPQ :: (Ord p) => p -> v -> PQ.IntPSQ p v -> S (PQ.IntPSQ p v)
+insPQ p x pq = do
+  i <- get
+  let
+    pq' = PQ.insert i p x pq
+  put (succ i)
+  pure pq'
 
 
+
+
+data RPTreeStats = RPTreeStats {
+  rptsLength :: Int
+                               } deriving (Eq, Show)
+
+treeStats :: RPTree d l a -> RPTreeStats
+treeStats (RPTree _ tt) = RPTreeStats l
+  where
+    l = length tt
+
+
+-- | How many data items are stored in the 'RPTree'
+treeSize :: (Foldable t) => RPTree d l (t a) -> Int
+treeSize = sum . leafSizes
+
+-- | How many data items are stored in each leaf of the 'RPTree'
+leafSizes :: Foldable t => RPTree d l (t a) -> RPT d l Int
+leafSizes (RPTree _ tt) = length <$> tt 
 
 -- pqSeq :: Ord a => PQ.IntPSQ a b -> Seq (a, b)
 -- pqSeq pqq = go pqq mempty
diff --git a/src/Data/RPTree/Conduit.hs b/src/Data/RPTree/Conduit.hs
--- a/src/Data/RPTree/Conduit.hs
+++ b/src/Data/RPTree/Conduit.hs
@@ -3,13 +3,13 @@
 {-# language BangPatterns #-}
 {-# options_ghc -Wno-unused-imports #-}
 {-# options_ghc -Wno-unused-top-binds #-}
+{-# options_ghc -Wno-type-defaults #-}
 module Data.RPTree.Conduit
   (
     tree,
   forest,
-  ForestParams,
-  fpMaxTreeDepth,
-  defaultParams
+  RPTreeConfig(..),
+  rpTreeCfg
   -- ** helpers
   , dataSource
   , liftC
@@ -19,6 +19,7 @@
 import Control.Monad (replicateM)
 import Data.Functor (void)
 import GHC.Word (Word64)
+import GHC.Stack (HasCallStack)
 
 -- conduit
 import qualified Data.Conduit as C (ConduitT, runConduit, yield, await, transPipe)
@@ -62,7 +63,7 @@
      -> Double -- ^ nonzero density of projection vectors
      -> Int -- ^ dimension of projection vectors
      -> C.ConduitT () (Embed v Double x) m () -- ^ data source
-     -> m (RPTree Double (V.Vector (Embed v Double x)))
+     -> m (RPTree Double () (V.Vector (Embed v Double x)))
 tree seed maxDepth minLeaf n pnz dim src = do
   let
     rvs = sample seed $ V.replicateM maxDepth (sparse pnz dim stdNormal)
@@ -84,10 +85,10 @@
            (Embed v d x)
            o
            m
-           (RPT d (V.Vector (Embed v d x))) 
+           (RPT d () (V.Vector (Embed v d x))) 
 insertC maxDepth minLeaf n rvs = chunkedAccum n z (insert maxDepth minLeaf rvs)
   where
-    z = Tip mempty
+    z = Tip () mempty
 
 
 
@@ -119,23 +120,35 @@
                        insertMultiC maxd minl chunksize rvss
   pure $ IM.intersectionWith RPTree rvss ts
 
-data ForestParams = CP {
+data RPTreeConfig = RPCfg {
   fpMaxTreeDepth :: Int -- ^ max tree depth \(l > 1\) 
   , fpMinLeafSize :: Int -- ^ min leaf size 
   , fpNumTrees :: Int -- ^ number of trees \(n_t > 1\)
   , fpDataChunkSize :: Int -- ^ data chunk size
   , fpProjNzDensity :: Double -- ^ nonzero density of projection vectors \(p_{nz} \in (0, 1)\)
-  , fpProjDimension :: Int -- ^ dimension of projection vectors \(d > 1\)
                           } deriving (Show)
 
-defaultParams :: Int  -- ^ dimension of projection vectors \(d > 1\)
-              -> ForestParams
-defaultParams d = CP 5 10 3 100 0.5 d
+defaultParams :: RPTreeConfig
+defaultParams = RPCfg 5 10 3 100 0.5
 
+-- | Configure the rp-tree forest construction process with some natural defaults
+rpTreeCfg :: Integral a =>
+             a -- ^ data size
+          -> Int -- ^ vector dimension
+          -> RPTreeConfig
+rpTreeCfg n d = RPCfg maxd minl ntree nchunk pnz
+  where
+    minl = 10
+    maxd = ceiling $ logBase 2 (fromIntegral n / fromIntegral minl)
+    ntree = 3
+    nchunk = ceiling $ fromIntegral n / 100
+    pnzMin = 1 / logBase 10 (fromIntegral d)
+    pnz = pnzMin `min` 1.0
 
 
 
 
+
 insertMultiC :: (Monad m, Ord d, Inner u v, VU.Unbox d, Fractional d, VG.Vector v1 (u d)) =>
                 Int  -- ^ max tree depth
              -> Int -- ^ min leaf size
@@ -145,11 +158,11 @@
                 (Embed v d x)
                 o
                 m
-                (IM.IntMap (RPT d (V.Vector (Embed v d x))))
+                (IM.IntMap (RPT d () (V.Vector (Embed v d x))))
 insertMultiC maxd minl n rvss = chunkedAccum n im0 (insertMulti maxd minl rvss)
   where
     im0 = IM.map (const z) rvss
-    z = Tip mempty
+    z = Tip () mempty
 
 
 {-# SCC insertMulti #-}
@@ -157,9 +170,9 @@
                Int
             -> Int
             -> IM.IntMap (v1 (u d)) -- ^ projection vectors
-            -> IM.IntMap (RPT d (V.Vector (Embed v d x))) -- ^ accumulator of subtrees
+            -> IM.IntMap (RPT d () (V.Vector (Embed v d x))) -- ^ accumulator of subtrees
             -> V.Vector (Embed v d x) -- ^ data chunk
-            -> IM.IntMap (RPT d (V.Vector (Embed v d x)))
+            -> IM.IntMap (RPT d () (V.Vector (Embed v d x)))
 insertMulti maxd minl rvss tacc xs =
   flip IM.mapWithKey tacc $ \ !i !t -> case IM.lookup i rvss of
                                       Just !rvs -> insert maxd minl rvs t xs
@@ -170,41 +183,44 @@
           Int -- ^ max tree depth
        -> Int -- ^ min leaf size
        -> v1 (u d) -- ^ projection vectors
-       -> RPT d (V.Vector (Embed v d x)) -- ^ accumulator
+       -> RPT d () (V.Vector (Embed v d x)) -- ^ accumulator
        -> V.Vector (Embed v d x) -- ^ data chunk
-       -> RPT d (V.Vector (Embed v d x))
+       -> RPT d () (V.Vector (Embed v d x))
 insert maxDepth minLeaf rvs = loop 0
   where
-    z = Tip mempty
+    z = Tip () mempty
     loop ixLev !tt xs =
       let
-        r = rvs VG.! ixLev
+        r = rvs VG.! ixLev -- proj vector for current level
       in
         case tt of
 
-          b@(Bin thr0 margin0 tl0 tr0) ->
-            if ixLev >= maxDepth || length xs <= minLeaf
+          b@(Bin _ thr0 margin0 tl0 tr0) ->
+            if ixLev >= maxDepth
               then b -- return current subtree
               else
-                let
-                  (thr, margin, ll, rr) =
-                    partitionAtMedian r xs
-                  margin' = margin0 <> margin
-                  thr' = (thr0 + thr) / 2
-                  tl = loop (ixLev + 1) tl0 ll
-                  tr = loop (ixLev + 1) tr0 rr
-                in Bin thr' margin' tl tr
+              case partitionAtMedian r xs of
+                Nothing -> Tip () mempty
+                Just (thr, margin, ll, rr) -> Bin () thr' margin' tl tr
+                  where
+                    margin' = margin0 <> margin
+                    thr' = (thr0 + thr) / 2
+                    tl = loop (ixLev + 1) tl0 ll
+                    tr = loop (ixLev + 1) tr0 rr
 
-          Tip xs0 -> do
+          Tip _ xs0 -> do
             let xs' = xs <> xs0
-            if ixLev >= maxDepth || length xs <= minLeaf
-              then Tip xs' -- concat data in leaf
+            if ixLev >= maxDepth || length xs' <= minLeaf
+              then Tip () xs' -- concat data in leaf
               else
-                let
-                  (thr, margin, ll, rr) = partitionAtMedian r xs'
-                  tl = loop (ixLev + 1) z ll
-                  tr = loop (ixLev + 1) z rr
-                in Bin thr margin tl tr
+              case partitionAtMedian r xs' of
+                Nothing -> Tip () mempty
+                Just (thr, margin, ll, rr) -> Bin () thr margin tl tr
+                  where
+                    tl = loop (ixLev + 1) z ll
+                    tr = loop (ixLev + 1) z rr
+
+
 
 
 -- | Aggregate the input stream in chunks of a given size (semantics of 'C.chunksOf'), and fold over the resulting stream building up an accumulator structure (e.g. a tree)
diff --git a/src/Data/RPTree/Draw.hs b/src/Data/RPTree/Draw.hs
--- a/src/Data/RPTree/Draw.hs
+++ b/src/Data/RPTree/Draw.hs
@@ -1,7 +1,18 @@
+{-# LANGUAGE FlexibleContexts #-}
 {-# language LambdaCase #-}
 {-# options_ghc -Wno-unused-imports #-}
-module Data.RPTree.Draw where
+{-# options_ghc -Wno-unused-top-binds #-}
+module Data.RPTree.Draw (
+  -- * CSV
+  writeCsv
+  -- * GraphViz dot
+  , writeDot
+  -- , draw
+                        )where
 
+import Data.Bifoldable (Bifoldable(..))
+import Data.Bifunctor (Bifunctor(..))
+import Data.Bitraversable (Bitraversable(..))
 import Data.List (intercalate)
 import Text.Printf (PrintfArg, printf)
 
@@ -10,8 +21,16 @@
 -- bytestring
 import qualified Data.ByteString.Lazy    as LBS (ByteString, writeFile)
 import qualified Data.ByteString.Builder as BSB (Builder, toLazyByteString, string7, charUtf8)
--- -- mtl
--- import Control.Monad.State (MonadState(..))
+-- containers
+import qualified Data.Set as S (Set, insert, fromList)
+-- mtl
+import Control.Monad.State (MonadState(..), modify)
+-- text
+import qualified Data.Text.Lazy as TL (Text)
+import qualified Data.Text.Lazy.Builder as TLB (Builder, toLazyText, fromString)
+import qualified Data.Text.Lazy.IO as TL (writeFile)
+-- transformers
+import Control.Monad.Trans.State (State, evalState)
 -- vector
 import qualified Data.Vector as V (Vector, replicateM)
 import qualified Data.Vector.Generic as VG (Vector(..), map, sum, unfoldr, unfoldrM, length, replicateM, (!))
@@ -23,24 +42,122 @@
 
 
 -- | Encode dataset as CSV and save into file
-writeCsv :: (Show a, Show b, VU.Unbox a) =>
-            FilePath
-         -> [(DVector a, b)] -- ^ data point, label
+writeCsv :: (Foldable t, VU.Unbox a, Show a, Show b) =>
+            FilePath -- ^ path of output file
+         -> t (V.Vector (DVector a, b)) -- ^ data point, label
          -> IO ()
-writeCsv fp ds = LBS.writeFile fp $ BSB.toLazyByteString $ toCsv ds
+writeCsv fp ds = TL.writeFile fp $ TLB.toLazyText $ toCsvV ds
 
+toCsvV :: (Foldable t, VU.Unbox a, Show a, Show b) =>
+          t (V.Vector (DVector a, b)) -> TLB.Builder
+toCsvV = foldMap (\v -> foldMap (\(r, i) -> toCsvRow r i <> newline ) v)
+
 toCsvRow :: (Show a, Show b, VU.Unbox a) =>
             DVector a
          -> b
-         -> BSB.Builder
-toCsvRow dv i = BSB.string7 $ intercalate "," [show x, show y, show i]
+         -> TLB.Builder
+toCsvRow dv i = TLB.fromString $ intercalate "," [show x, show y, show i]
   where
     (x:y:_) = toListDv dv
 
-toCsv :: (Show a, Show b, VU.Unbox a) =>
-         [(DVector a, b)] -> BSB.Builder
-toCsv rs = mconcat [toCsvRow r i <> BSB.charUtf8 '\n' | (r, i) <- rs]
+newline :: TLB.Builder
+newline = TLB.fromString "\n"
 
+-- toCsv :: (Foldable t, Show a, Show b, VU.Unbox a) =>
+--          t (DVector a, b) -> TLB.Builder
+-- toCsv = foldMap (\(r, i) -> toCsvRow r i <> newline )
+
+
+-- | tree to graphviz dot format
+writeDot :: Ord t =>
+            (t -> String) -- ^ how to render the node content
+         -> FilePath -- ^ path of output file
+         -> String -- ^ graph name
+         -> RPTree d x t
+         -> IO ()
+writeDot f fp name tt = TL.writeFile fp (toDot f name tt) 
+
+toDot :: Ord a => (a -> String) -> String -> RPTree d x a -> TL.Text
+toDot f name (RPTree _ tt) = TLB.toLazyText $ open <> x <> close
+  where
+    x = foldl insf mempty $ toEdges tt
+      where
+        insf acc = \case
+          Edge i1 i2 ->
+            acc <> TLB.fromString (unwords [show i1, "->", show i2]) <> newline
+          Node i xs -> acc <> TLB.fromString (unwords [show i, nlab ] ) <> newline
+            where
+              nlab = unwords ["[", "label=\"", f xs,"\"]"]
+          BNode i -> acc <> TLB.fromString (unwords [show i, blab]) <> newline
+            where
+              blab = unwords ["[", "shape=point", "]"]
+    open = TLB.fromString $ unwords ["digraph" , name, "{\n"]
+    close = TLB.fromString "}"
+
+data G a = Edge Int Int
+         | Node Int a -- tip nodes
+         | BNode Int -- branching point nodes
+         deriving (Eq, Ord, Show)
+
+
+toEdges :: (Ord a) => RPT d x a -> S.Set (G a)
+toEdges = S.fromList . go [] [] . labelBranches
+  where
+    go s acc = \case
+      Tip i1 x ->
+        let
+          n = Node i1 x
+          acc' = push n acc
+          acc'' = maybe acc' (\i0 -> push (Edge i0 i1) acc') (pop s)
+        in
+          acc''
+      Bin i1 _ _ tl tr ->
+        let
+          b1 = BNode i1
+          acc' = push b1 acc
+          acc'' = case pop s of
+            Nothing -> acc'
+            Just i0 ->
+              let
+                e = Edge i0 i1
+                b0 = BNode i0
+              in push e (push b0 acc')
+          s' = push i1 s
+        in
+          go s' acc'' tl <> go s' acc tr
+
+
+
+labelBranches :: Bitraversable t => t x d -> t Int d
+labelBranches = flip evalState 0 . bitraverse counter pure
+
+counter :: (MonadState Int m) => x -> m Int
+counter _ = do
+  i <- get
+  modify succ
+  pure i
+
+type Stack a = [a]
+push :: a -> Stack a -> Stack a
+push = (:)
+pop :: Stack a -> Maybe a
+pop xs
+  | null xs = Nothing
+  | otherwise = Just $ head xs
+
+
+-- tt0 :: RPT Integer () [a]
+tt0 :: RPT Integer [a1] [a2]
+tt0 = Bin [] 0 mempty tl t
+  where
+    tl = Bin [] 1 mempty (Bin [] 2 mempty t t) (Bin [] 3 mempty t t)
+    t = Tip [] []
+
+
+
+-- 
+
+
 -- | Render a tree to stdout
 --
 -- Useful for debugging
@@ -48,20 +165,20 @@
 -- This should be called only for small trees, otherwise the printed result quickly overflows the screen and becomes hard to read.
 --
 -- NB : prints distance information rounded to two decimal digits
-draw :: (Show a, Boxed a, PrintfArg v) => RPTree v a -> IO ()
+draw :: (Show a, Boxed a, PrintfArg v) => RPTree v l a -> IO ()
 draw = drawRPT . _rpTree
 
-drawRPT :: (Show a, Boxed a, PrintfArg v) => RPT v a -> IO ()
+drawRPT :: (Show a, Boxed a, PrintfArg v) => RPT v l a -> IO ()
 drawRPT = putStrLn . toStringRPT
 
-toStringRPT :: (Show a, Boxed a, PrintfArg v) => RPT v a -> String
+toStringRPT :: (Show a, Boxed a, PrintfArg v) => RPT v l a -> String
 toStringRPT = B.render . toBox
 
-toBox :: (Show a, Boxed a, PrintfArg v) => RPT v a -> B.Box
+toBox :: (Show a, Boxed a, PrintfArg v) => RPT v l a -> B.Box
 toBox = \case
-  (Bin thr _ tl tr) ->
+  (Bin _ thr _ tl tr) ->
     txt (node thr) `stack` (toBox tl `byside` toBox tr)
-  Tip xs -> boxed xs -- tipData xs -- txt $ show x
+  Tip _ xs -> boxed xs -- tipData xs -- txt $ show x
   where
     node x = printf "%5.2f" x -- (show x)
 
diff --git a/src/Data/RPTree/Gen.hs b/src/Data/RPTree/Gen.hs
--- a/src/Data/RPTree/Gen.hs
+++ b/src/Data/RPTree/Gen.hs
@@ -10,18 +10,15 @@
 -- containers
 import qualified Data.IntMap as IM (IntMap, insert, toList)
 -- splitmix-distribitions
-import System.Random.SplitMix.Distributions (Gen, GenT, stdUniform, bernoulli, exponential, normal, discrete, categorical)
+import System.Random.SplitMix.Distributions (Gen, GenT, uniformR, stdUniform, bernoulli, exponential, normal, discrete, categorical)
 -- transformers
 import Control.Monad.Trans.Class (MonadTrans(..))
 import Control.Monad.Trans.State (StateT(..), get, put, runStateT, evalStateT, State, runState, evalState)
--- vector
 
-
 import qualified Data.Vector.Generic as VG (Vector(..), unfoldrM, length, replicateM, (!))
 import qualified Data.Vector.Unboxed as VU (Vector, Unbox, fromList)
 
-
-import Data.RPTree.Internal (RPTree(..), RPT(..), SVector(..), fromListSv, DVector(..))
+import Data.RPTree.Internal (RPTree(..), RPT(..), SVector(..), fromListSv, DVector(..), fromListDv)
 
 
 -- | Sample without replacement with a single pass over the data
@@ -114,6 +111,16 @@
           let p = gs !! i
           p
 
+
+circle2d :: (Monad m) => Double -> GenT m (DVector Double)
+circle2d r = go
+  where
+    go = do
+      x <- uniformR (- r) r
+      y <- uniformR (- r) r
+      if x**2 + y**2 <= r
+        then pure $ fromListDv [x, y]
+        else go
 
 normalSparse2 :: Monad m => Double -> Int -> GenT m (SVector Double)
 normalSparse2 pnz d = do
diff --git a/src/Data/RPTree/Internal.hs b/src/Data/RPTree/Internal.hs
--- a/src/Data/RPTree/Internal.hs
+++ b/src/Data/RPTree/Internal.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE RankNTypes #-}
+{-# language DeriveAnyClass #-}
 {-# LANGUAGE DeriveTraversable #-}
 {-# LANGUAGE DeriveFoldable #-}
 -- {-# LANGUAGE UndecidableInstances #-}
@@ -9,7 +10,7 @@
 {-# language LambdaCase #-}
 {-# language MultiParamTypeClasses #-}
 {-# language GeneralizedNewtypeDeriving #-}
--- {-# language TemplateHaskell #-}
+{-# language TemplateHaskell #-}
 {-# LANGUAGE BangPatterns        #-}
 {-# options_ghc -Wno-unused-imports #-}
 module Data.RPTree.Internal where
@@ -17,6 +18,9 @@
 import Control.Exception (Exception(..))
 import Control.Monad.IO.Class (MonadIO(..))
 import Control.Monad.ST (runST)
+import Data.Bifoldable (Bifoldable(..))
+import Data.Bifunctor (Bifunctor(..))
+import Data.Bitraversable (Bitraversable(..))
 import Data.Function ((&))
 import Data.Foldable (fold, foldl', toList)
 import Data.Functor.Identity (Identity(..))
@@ -26,6 +30,7 @@
 import Data.Semigroup (Min(..), Max(..))
 import Data.Typeable (Typeable)
 import GHC.Generics (Generic)
+import GHC.Stack (HasCallStack)
 
 -- bytestring
 import qualified Data.ByteString.Lazy as LBS (ByteString, toStrict, fromStrict)
@@ -33,6 +38,10 @@
 import qualified Data.IntMap.Strict as IM (IntMap, fromList)
 -- deepseq
 import Control.DeepSeq (NFData(..))
+-- microlens
+import Lens.Micro ((^..), Traversal', folded, Getting)
+-- microlens-th
+import Lens.Micro.TH (makeLensesFor)
 -- serialise
 import Codec.Serialise (Serialise(..), serialise, deserialiseOrFail)
 -- vector
@@ -49,7 +58,9 @@
 data Embed v e a = Embed {
   eEmbed :: !(v e) -- ^ vector embedding
   , eData :: !a -- ^ data item
-                       } deriving (Eq, Ord, Show, Generic, Functor)
+                       } deriving (Eq, Ord, Generic, Functor)
+instance (Show (v e), Show e, Show a) => Show (Embed v e a) where
+  show (Embed v dat) = unwords [show v, show dat]
 instance (NFData (v e), NFData a) => NFData (Embed v e a)
 instance (Serialise (v e), Serialise a) => Serialise (Embed v e a)
 
@@ -66,7 +77,9 @@
 data Margin a = Margin {
   cMarginLow :: !(Max a) -- ^ lower bound on the cut point
   , cMarginHigh :: !(Min a) -- ^ upper bound
-                   } deriving (Eq, Show, Generic)
+                   } deriving (Eq, Generic)
+instance Show a => Show (Margin a) where
+  show (Margin lo hi) = unwords ["low", show (getMax lo), "high", show (getMin hi)]
 instance (Serialise a) => Serialise (Margin a)
 getMargin :: Margin a -> (a, a)
 getMargin (Margin ml mh) = (getMax ml, getMin mh)
@@ -74,7 +87,8 @@
 -- | Used for updating in a streaming setting
 instance (Ord a) => Semigroup (Margin a) where
   Margin lo1 hi1 <> Margin lo2 hi2 = Margin (lo1 <> lo2) (hi1 <> hi2)
-
+instance (Ord a, Num a) => Monoid (Margin a) where
+  mempty = Margin (Max 0) (Min 0)
 
 -- | Sparse vectors with unboxed components
 data SVector a = SV { svDim :: {-# UNPACK #-} !Int
@@ -124,37 +138,52 @@
 -- | Internal
 --
 -- one projection vector per tree level (as suggested in https://www.cs.helsinki.fi/u/ttonteri/pub/bigdata2016.pdf )
-data RPT d a =
+data RPT d l a =
   Bin {
-  _rpThreshold ::  !d
+  _rpLabel :: l
+  , _rpThreshold ::  !d
   , _rpMargin :: {-# UNPACK #-} !(Margin d)
-  , _rpL :: !(RPT d a)
-  , _rpR :: !(RPT d a) }
-  | Tip { _rpData :: !a }
+  , _rpL :: !(RPT d l a)
+  , _rpR :: !(RPT d l a) }
+  | Tip {
+      _rpLabel :: l
+    , _rpData :: !a }
   deriving (Eq, Show, Generic, Functor, Foldable, Traversable)
-instance (Serialise a, Serialise d) => Serialise (RPT d a)
--- makeLensesFor [("_rpData", "rpData")] ''RPT
-instance (NFData v, NFData a) => NFData (RPT v a)
+instance Bifunctor (RPT d) where
+  bimap f g = \case
+    Bin x thr mg tl tr -> Bin (f x) thr mg (bimap f g tl) (bimap f g tr)
+    Tip x y -> Tip (f x) (g y)
+instance Bifoldable (RPT d) where
+  bifoldMap f g = \case
+    Bin x _ _ tl tr -> f x <> bifoldMap f g tl <> bifoldMap f g tr
+    Tip x y -> f x <> g y
+instance Bitraversable (RPT d) where
+  bitraverse f g = \case
+    Bin x thr mg tl tr -> Bin <$> f x <*> pure thr <*> pure mg <*> bitraverse f g tl <*> bitraverse f g tr
+    Tip x y -> Tip <$> f x <*> g y
+instance (Serialise a, Serialise l, Serialise d) => Serialise (RPT d l a)
+makeLensesFor [("_rpData", "rpData"), ("_rpLabel", "rpLabel")] ''RPT
+instance (NFData v, NFData l, NFData a) => NFData (RPT v l a)
 
 -- | Random projection trees
 --
--- The first type parameter corresponds to a floating point scalar value, the second is the type of the data collected at the leaves of the tree (e.g. lists of vectors)
+-- The first type parameter corresponds to a floating point scalar value, the second labels every tree part (e.g. for visual rendering) and the third is the type of the data collected at the leaves of the tree (e.g. lists of vectors).
 --
--- We keep them separate to leverage the Functor instance for postprocessing and visualization
+-- We keep them separate to leverage the (Bi-)Functor instance for postprocessing and visualization.
 --
--- One projection vector per tree level (as suggested in https://www.cs.helsinki.fi/u/ttonteri/pub/bigdata2016.pdf )
-data RPTree d a = RPTree {
+-- This implementation uses one projection vector per tree level (as suggested in https://www.cs.helsinki.fi/u/ttonteri/pub/bigdata2016.pdf ).
+data RPTree d l a = RPTree {
   _rpVectors :: V.Vector (SVector d) -- ^ one random projection vector per tree level
-  , _rpTree :: !(RPT d a)
+  , _rpTree :: !(RPT d l a)
                          } deriving (Eq, Show, Functor, Foldable, Traversable, Generic)
-instance (Serialise d, Serialise a, VU.Unbox d) => Serialise (RPTree d a)
--- makeLensesFor [("_rpTree", "rpTree")] ''RPTree
-instance (NFData a, NFData d) => NFData (RPTree d a)
+instance (Serialise d, Serialise l, Serialise a, VU.Unbox d) => Serialise (RPTree d l a)
+makeLensesFor [("_rpTree", "rpTree")] ''RPTree
+instance (NFData a, NFData l, NFData d) => NFData (RPTree d l a)
 
 -- | A random projection forest is an ordered set of 'RPTree's
 --
 -- This supports efficient updates of the ensemble in the streaming/online setting.
-type RPForest d a = IM.IntMap (RPTree d a)
+type RPForest d a = IM.IntMap (RPTree d () a)
 
 -- | Serialise each tree in the 'RPForest' as a separate bytestring
 serialiseRPForest :: (Serialise d, Serialise a, VU.Unbox d) =>
@@ -170,19 +199,19 @@
   Left e -> Left (show e)
   Right xs -> Right $ IM.fromList $ zip [0 ..] xs
 
--- rpTreeData :: Traversal' (RPTree d a) a
--- rpTreeData = rpTree . rpData
+rpTreeData :: Traversal' (RPTree d l a) a
+rpTreeData = rpTree . rpData
 
--- -- | All data buckets stored at the leaves of the tree
--- leaves :: RPTree d a -> [a]
--- leaves = (^.. rpTreeData)
+-- | All data buckets stored at the leaves of the tree
+leaves :: RPTree d l a -> [a]
+leaves = (^.. rpTreeData)
 
 -- | Number of tree levels
-levels :: RPTree d a -> Int
+levels :: RPTree d l a -> Int
 levels (RPTree v _) = VG.length v
 
 -- | Set of data points used to construct the index
-points :: Monoid m => RPTree d m -> m
+points :: Monoid m => RPTree d l m -> m
 points (RPTree _ t) = fold t
 
 -- -- points in 2d
@@ -366,19 +395,25 @@
 -}
 
 
+type VE v a x = V.Vector (Embed v a x)
 
 
 {-# SCC partitionAtMedian #-}
 -- | Partition the data wrt the median value of the inner product
-partitionAtMedian :: (Ord a, Inner u v, VU.Unbox a, Fractional a) =>
-                     u a -- ^ projection vector
-                  -> V.Vector (Embed v a x) -- ^ dataset (3 or more elements)
-                  -> (a, Margin a, V.Vector (Embed v a x), V.Vector (Embed v a x)) -- ^ median, margin, smaller, larger
-partitionAtMedian r xs = (thr, margin, ll, rr)
+partitionAtMedian ::
+  (Ord a, Inner u v, VU.Unbox a, Fractional a) =>
+  u a -- ^ projection vector
+  -> V.Vector (Embed v a x) -- ^ dataset (3 or more elements)
+  -> Maybe (a, Margin a, VE v a x, VE v a x) -- ^ median, margin, smaller, larger
+partitionAtMedian r xs
+  | n < 1 = Nothing
+  | otherwise = Just (thr, margin, ll, rr)
   where
     (ll, rr) = (VG.take nh xs', VG.drop nh xs')
-    -- (pjl, pjr) = (VG.head inns, VG.last inns) -- (min, max) inner product values
-    (mgl, mgr) = (inns VG.! (nh - 1), inns VG.! (nh + 1))
+    (mgl, mgr)
+      | n >= 3 = (inns VG.! (nh - 1), inns VG.! (nh + 1))
+      | n == 2 = (inns VG.! 0, inns VG.! 1)
+      | otherwise = let z = inns VG.! 0 in (z, z) -- assuming at least 1 element
     margin = Margin (Max mgl) (Min mgr)
     thr = inns VG.! nh -- inner product threshold,  mgl < thr < mgr
     n = VG.length xs -- total data size
@@ -386,13 +421,12 @@
     projs = sortByVG snd $ VG.map (\xe -> (xe, r `inner` (eEmbed xe))) xs
     (xs', inns) = VG.unzip projs
 
+
 sortByVG :: (VG.Vector v a, Ord b) => (a -> b) -> v a -> v a
 sortByVG f v = runST $ do
   vm <- VG.thaw v
   V.sortBy (comparing f) vm
   VG.freeze vm
-
-
 
 
 
