<commit_before>module HLearn.DataStructures.KDTree
    where
    
import Control.Monad
import Control.Monad.Random
import qualified Data.Vector as V
import Test.QuickCheck
import Debug.Trace

import qualified Data.Foldable as F
import qualified Data.Traversable as T

import HLearn.Algebra

-------------------------------------------------------------------------------
-- data types

-- type DP = V.Vector Double
-- type family (Dimension dp)
-- 
-- (!) :: DP -> Dimension -> Double
-- (!) = (V.!)

class (Num (DimensionIndex x), Bounded (DimensionIndex x), Ord (DimensionIndex x), Ord (DimensionBase x)) => HasDimensions x where
    type DimensionIndex x 
    type DimensionBase x
    (!) :: x -> DimensionIndex x -> DimensionBase x

data KDTree dp
    = Leaf 
    | Node { val :: dp, splitdim :: DimensionIndex dp, left :: KDTree dp, right :: KDTree dp}
--     deriving (Read,Show,Eq,Ord)

deriving instance (Show dp, Show (DimensionIndex dp)) => Show (KDTree dp)

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

ppshow :: (Show dp) => KDTree dp -> String
ppshow = go [] 
    where
        go i Leaf = "" -- i++"*\n"
        go i t = i ++ " " ++ (show $ val t) ++ "\n" ++ go (i++"l") (left t) ++ go (i++"r") (right t)

basicshow :: (Show dp, Show (DimensionIndex dp)) => KDTree dp -> String
basicshow Leaf = "Leaf"
basicshow t = "Node { val="++show (val t)++", splitdim="++show (splitdim t)++", left="++go (left t)++", right="++go (right t)++" }"
    where
        go Leaf = "Leaf"
        go _ = "Node"


nextSplitDim :: (HasDimensions dp) => KDTree dp -> DimensionIndex dp
nextSplitDim t 
    | splitdim t == maxBound = minBound
    | otherwise = splitdim t+1

insert :: (HasDimensions dp) => KDTree dp -> dp -> KDTree dp
insert Leaf dp = Node dp 0 Leaf Leaf
insert t dp
    | dp ! splitdim t > val t ! splitdim t = t { right = insert (right t) dp }
    | dp ! splitdim t < val t ! splitdim t = t { left = insert (left t) dp }

isValid :: (HasDimensions dp) => KDTree dp -> Bool
isValid Leaf = True
isValid t    = thisvalid && isValid (left t) && isValid (right t)
    where
        thisvalid = go (left t) (\dp -> dp ! splitdim t <= val t ! splitdim t)
                 && go (right t) (\dp -> dp ! splitdim t > val t ! splitdim t)
         
        go Leaf f = True
        go t f = f (val t) && go (left t) f && go (right t) f

size :: KDTree dp -> Int
size Leaf = 0
size t = 1 + (size $ left t) + (size $ right t)

depth :: KDTree dp -> Int
depth Leaf = 0
depth t = 1+max (depth $ left t) (depth $ right t)

nodesAtDepth :: Int -> KDTree dp -> Int
nodesAtDepth _ Leaf = 0
nodesAtDepth 0 _ = 1
nodesAtDepth i t = nodesAtDepth (i-1) (left t) + nodesAtDepth (i-1) (right t)

balanceScore :: KDTree dp -> Double
balanceScore t = (fromIntegral $ depthSum t) / (fromIntegral $ balanceDepthSum $ size t)

depthSum :: KDTree dp -> Int
depthSum t = go 0 t 
    where
        go i Leaf = 0
        go i t = i+(go (i+1) (left t)) + (go (i+1) (right t))

balanceDepthSum :: Int -> Int 
balanceDepthSum i = fromIntegral $ go 0 i
    where
        go exp i
            | 2^exp < i = (exp+1)*2^exp + go (exp+1) (i-2^exp)
            | otherwise = (exp+1)*i

-------------------------------------------------------------------------------
-- algebra

debug _ _ _ action = action
-- debug pos a b action = trace (pos
--     ++"\n  ra="++(show $ right a)
--     ++"\n  rb="++(show $ right b)
--     ++"\n  la="++(show $ left a)
--     ++"\n  lb="++(show $ left b)
--     ++"\n  va="++(show $ val a)
--     ++"\n  vb="++(show $ val b)
--     ++"\n  sa="++(show $ splitdim a)
--     ++"\n  sb="++(show $ splitdim b)
--     )
--     action

instance (HasDimensions dp) => Monoid (KDTree dp) where
    mempty = Leaf
    mappend Leaf Leaf = Leaf
    mappend Leaf a = a
    mappend a Leaf = a
    mappend a@(Node _ _ Leaf Leaf) b@(Node _ _ Leaf Leaf)
        | val a ! splitdim a >  val b ! splitdim a = debug "a1" a b $ 
            Node (val a) (splitdim a) (Node (val b) (nextSplitDim a) Leaf Leaf) Leaf
        | val a ! splitdim a <= val b ! splitdim a = debug "a2" a b $ 
            Node (val a) (splitdim a) Leaf (Node (val b) (nextSplitDim a) Leaf Leaf)
    mappend a b@(Node _ _ Leaf Leaf) = insert a $ val b
    mappend a@(Node _ _ Leaf Leaf) b = insert b $ val a
    mappend a b = if  splitdim a == splitdim b 
        then if
            | val a ! splitdim a <  val b ! splitdim a -> debug "b1" a b $ right a <> b 
                { left = left a <> left b <> train1dp (val a) 
                } 
            | val a ! splitdim a >  val b ! splitdim a -> debug "b2" a b $ left a <> b 
                { right = right a <> right b <> train1dp (val a)
                } 
            | val a ! splitdim a == val b ! splitdim a -> debug "b3" a b $ b 
                { left = left a <> left b <> train1dp (val a)
                , right = right a <> right b 
                }
        else debug "c1" a b $ ((train1dp $ val a) {splitdim = splitdim a} ) <> (right a <> (left a <> b))
       
       
        
---------------------------------------

-- instance Functor KDTree where
--     fmap f Leaf = Leaf
--     fmap f t = Node 
--         { val = f $ val t
--         , splitdim = fromInteger $ toInteger $ splitdim t
--         , left = fmap f $ left t
--         , right = fmap f $ right t
--         }

instance F.Foldable KDTree where
    foldr f i Leaf = i
    foldr f i t = F.foldr f (F.foldr f (f (val t) i) (right t)) (left t)

-- instance T.Traversable KDTree where
    

-------------------------------------------------------------------------------
-- model

instance (HasDimensions dp) => HomTrainer (KDTree dp) where
    type Datapoint (KDTree dp) = dp
    train1dp dp = Node dp 0 Leaf Leaf

-------------------------------------------------------------------------------
-- testing

instance Arbitrary (V.Vector Double) where
    arbitrary = do
        a <- choose (-100,100)
        b <- choose (-100,100)
        return $ V.fromList $ [a,b]

instance (HasDimensions dp, Arbitrary dp) => Arbitrary (KDTree dp) where
    arbitrary = do
        xs <- arbitrary<commit_msg>'Tree-Independent Dual-Tree Algorithms' is an amazing paper.<commit_after>module HLearn.DataStructures.KDTree
    where
    
import Control.Monad
import Control.Monad.Random
import qualified Data.Vector as V
import Test.QuickCheck
import Debug.Trace

import qualified Data.Foldable as F
import qualified Data.Traversable as T

import HLearn.Algebra

-------------------------------------------------------------------------------
-- data types

-- type DP = V.Vector Double
-- type family (Dimension dp)
-- 
-- (!) :: DP -> Dimension -> Double
-- (!) = (V.!)

class (Num (DimensionIndex x), Bounded (DimensionIndex x), Ord (DimensionIndex x), Ord (DimensionBase x)) => HasDimensions x where
    type DimensionIndex x 
    type DimensionBase x
    (!) :: x -> DimensionIndex x -> DimensionBase x

data KDTree dp
    = Leaf 
    | Node { val :: dp, splitdim :: DimensionIndex dp, left :: KDTree dp, right :: KDTree dp}
--     deriving (Read,Show,Eq,Ord)

deriving instance (Show dp, Show (DimensionIndex dp)) => Show (KDTree dp)

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

ppshow :: (Show dp) => KDTree dp -> String
ppshow = go [] 
    where
        go i Leaf = "" -- i++"*\n"
        go i t = i ++ " " ++ (show $ val t) ++ "\n" ++ go (i++"l") (left t) ++ go (i++"r") (right t)

basicshow :: (Show dp, Show (DimensionIndex dp)) => KDTree dp -> String
basicshow Leaf = "Leaf"
basicshow t = "Node { val="++show (val t)++", splitdim="++show (splitdim t)++", left="++go (left t)++", right="++go (right t)++" }"
    where
        go Leaf = "Leaf"
        go _ = "Node"


nextSplitDim :: (HasDimensions dp) => KDTree dp -> DimensionIndex dp
nextSplitDim t 
    | splitdim t == maxBound = minBound
    | otherwise = splitdim t+1

insert :: (HasDimensions dp) => KDTree dp -> dp -> KDTree dp
insert Leaf dp = Node dp 0 Leaf Leaf
insert t dp
    | dp ! splitdim t > val t ! splitdim t = t { right = insert (right t) dp }
    | dp ! splitdim t < val t ! splitdim t = t { left = insert (left t) dp }

isValid :: (HasDimensions dp) => KDTree dp -> Bool
isValid Leaf = True
isValid t    = thisvalid && isValid (left t) && isValid (right t)
    where
        thisvalid = go (left t) (\dp -> dp ! splitdim t <= val t ! splitdim t)
                 && go (right t) (\dp -> dp ! splitdim t > val t ! splitdim t)
         
        go Leaf f = True
        go t f = f (val t) && go (left t) f && go (right t) f

size :: KDTree dp -> Int
size Leaf = 0
size t = 1 + (size $ left t) + (size $ right t)

depth :: KDTree dp -> Int
depth Leaf = 0
depth t = 1+max (depth $ left t) (depth $ right t)

nodesAtDepth :: Int -> KDTree dp -> Int
nodesAtDepth _ Leaf = 0
nodesAtDepth 0 _ = 1
nodesAtDepth i t = nodesAtDepth (i-1) (left t) + nodesAtDepth (i-1) (right t)

balanceScore :: KDTree dp -> Double
balanceScore t = (fromIntegral $ depthSum t) / (fromIntegral $ balanceDepthSum $ size t)

depthSum :: KDTree dp -> Int
depthSum t = go 0 t 
    where
        go i Leaf = 0
        go i t = i+(go (i+1) (left t)) + (go (i+1) (right t))

balanceDepthSum :: Int -> Int 
balanceDepthSum i = fromIntegral $ go 0 i
    where
        go exp i
            | 2^exp < i = (exp+1)*2^exp + go (exp+1) (i-2^exp)
            | otherwise = (exp+1)*i

-------------------------------------------------------------------------------
-- algebra

debug _ _ _ action = action
-- debug pos a b action = trace (pos
--     ++"\n  ra="++(show $ right a)
--     ++"\n  rb="++(show $ right b)
--     ++"\n  la="++(show $ left a)
--     ++"\n  lb="++(show $ left b)
--     ++"\n  va="++(show $ val a)
--     ++"\n  vb="++(show $ val b)
--     ++"\n  sa="++(show $ splitdim a)
--     ++"\n  sb="++(show $ splitdim b)
--     )
--     action

instance (HasDimensions dp) => Monoid (KDTree dp) where
    mempty = Leaf
    mappend Leaf Leaf = Leaf
    mappend Leaf a = a
    mappend a Leaf = a
    mappend a@(Node _ _ Leaf Leaf) b@(Node _ _ Leaf Leaf)
        | val a ! splitdim a >  val b ! splitdim a = debug "a1" a b $ 
            Node (val a) (splitdim a) (Node (val b) (nextSplitDim a) Leaf Leaf) Leaf
        | val a ! splitdim a <= val b ! splitdim a = debug "a2" a b $ 
            Node (val a) (splitdim a) Leaf (Node (val b) (nextSplitDim a) Leaf Leaf)
    mappend a b@(Node _ _ Leaf Leaf) = insert a $ val b
    mappend a@(Node _ _ Leaf Leaf) b = insert b $ val a
    mappend a b = if  splitdim a == splitdim b 
        then if
            | val a ! splitdim a <  val b ! splitdim a -> debug "b1" a b $ right a <> b 
                { left = left a <> left b <> train1dp (val a) 
                } 
            | val a ! splitdim a >  val b ! splitdim a -> debug "b2" a b $ left a <> b 
                { right = right a <> right b <> train1dp (val a)
                } 
            | val a ! splitdim a == val b ! splitdim a -> debug "b3" a b $ b 
                { left = left a <> left b <> train1dp (val a)
                , right = right a <> right b 
                }
        else debug "c1" a b $ ((train1dp $ val a) {splitdim = splitdim a} ) <> (right a <> (left a <> b))
       
       
        
---------------------------------------

-- instance Functor KDTree where
--     fmap f Leaf = Leaf
--     fmap f t = Node 
--         { val = f $ val t
--         , splitdim = fromInteger $ toInteger $ splitdim t
--         , left = fmap f $ left t
--         , right = fmap f $ right t
--         }

instance F.Foldable KDTree where
    foldr f i Leaf = i
    foldr f i t = F.foldr f (F.foldr f (f (val t) i) (right t)) (left t)

class Prunable t where
    pfoldr :: (t a -> Bool) -> (a -> b -> b) -> b -> t a -> b

instance Prunable KDTree where
    pfoldr p f i t = if p t 
        then i
        else pfoldr p f (pfoldr p f (f (val t) i) (right t)) (left t) 

class (F.Foldable t) => DualFoldable t where
    dfoldr :: ((a,a) -> b -> b) -> b -> t a -> t a -> b
    dfoldr f i t1 t2 = foldr f i [(x,y) | x <- (F.toList t1), y <- (F.toList t2)]

-- instance T.Traversable KDTree where
    

-------------------------------------------------------------------------------
-- model

instance (HasDimensions dp) => HomTrainer (KDTree dp) where
    type Datapoint (KDTree dp) = dp
    train1dp dp = Node dp 0 Leaf Leaf

-------------------------------------------------------------------------------
-- testing

instance Arbitrary (V.Vector Double) where
    arbitrary = do
        a <- choose (-100,100)
        b <- choose (-100,100)
        return $ V.fromList $ [a,b]

instance (HasDimensions dp, Arbitrary dp) => Arbitrary (KDTree dp) where
    arbitrary = do
        xs <- arbitrary
{'fraction_non_alphanumeric': 0.1949286846275753, 'fraction_numerical': 0.007606973058637084, 'mean_word_length': 1.7312270071413114, 'pattern_counts': {'":': 0, '<': 45, '<?xml version=': 0, '>': 113, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '11412', 'n_tokens_mistral': 4303, 'n_tokens_neox': 4070, 'n_words': 1562}