{-# LANGUAGE NoMonomorphismRestriction, DeriveFunctor, OverlappingInstances, DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}

{- |
  Module     : Data.RTree.Base
  Copyright  : Copyright (c) 2014, Birte Wagner, Sebastian Philipp
  License    : MIT

  Maintainer : Birte Wagner, Sebastian Philipp (sebastian@spawnhost.de)
  Stability  : experimental
  Portability: not portable

  Internal implementations. Use Data.RTree instead

-}


module Data.RTree.Base
(
    RTree,
    empty,
    singleton,
    insert,
    union,
    lookup,
    lookupRange,
    fromList,
    toList,
    delete,
    length,
    null,
    keys,
    values,
    mapMaybe,
    foldWithMBB,
    getMBB,

    -- | testing

    pp,
    isValid,
    unionDistinct,
    getC1,
    getC2,
    getC3,
    getC4
)
where

import           Prelude hiding (lookup, length, null)

import           Data.Binary
import           Data.Function
import           Data.List (maximumBy, minimumBy, partition)
import qualified Data.List as L (length)
import           Data.Maybe (catMaybes, isJust)
import qualified Data.Maybe as Maybe (mapMaybe)
import           Data.Typeable (Typeable)

import           Control.Applicative ((<$>))
import           Control.DeepSeq (NFData, rnf)

import           GHC.Generics (Generic)

import           Data.RTree.MBB hiding (mbb)

data RTree a = 
      Node4 {getMBB :: {-# UNPACK #-} ! MBB, getC1 :: ! (RTree a), getC2 :: ! (RTree a), getC3 :: ! (RTree a), getC4 :: ! (RTree a) }
    | Node3 {getMBB :: {-# UNPACK #-} ! MBB, getC1 :: ! (RTree a), getC2 :: ! (RTree a), getC3 :: ! (RTree a) }
    | Node2 {getMBB :: {-# UNPACK #-} ! MBB, getC1 :: ! (RTree a), getC2 :: ! (RTree a) }
    | Node  {getMBB ::                  MBB, getChildren' :: [RTree a] }
    | Leaf  {getMBB :: {-# UNPACK #-} ! MBB, getElem :: a}
    | Empty
    deriving (Show, Eq, Functor, Typeable, Generic)

m, n :: Int
m = 2
n = 4


unionMBB' :: RTree a -> RTree a -> MBB
unionMBB' x y = unionMBB [getMBB x, getMBB y]

-- ---------------
-- smart constuctors

-- | creates an empty tree
empty :: RTree a
empty = Empty

-- | returns True, if empty
--
-- prop> null empty = True
null :: RTree a -> Bool
null Empty = True
null _     = False

-- | creates a single element tree
singleton :: MBB -> a -> RTree a
singleton mbb x = Leaf mbb x

node :: MBB -> [RTree a] -> RTree a
node mbb [x,y]     = Node2 mbb x y
node mbb [x,y,z]   = Node3 mbb x y z
node mbb [x,y,z,w] = Node4 mbb x y z w
node _   []        = error "node: empty"
node mbb xs        = Node mbb xs

createNodeWithChildren :: [RTree a] -> RTree a
createNodeWithChildren c = node (unionMBB $ getMBB <$> c) c

norm :: RTree a -> RTree a
norm (Node4 mbb x y z w) = Node mbb [x,y,z,w]
norm (Node3 mbb x y z)   = Node mbb [x,y,z]
norm (Node2 mbb x y)     = Node mbb [x,y]
norm x = x

getChildren :: RTree a -> [RTree a]
getChildren Empty = error "getChildren: Empty"
getChildren Leaf{} = error "getChildren: Leaf"
getChildren t = getChildren' $ norm t

-- ----------------------------------
-- Lists

-- | creates a tree out of pairs
fromList :: [(MBB, a)] -> RTree a
fromList l = fromList' $ (uncurry singleton) <$> l

fromList' :: [RTree a] -> RTree a
fromList' [] = error "fromList' empty"
fromList' ts = foldr1 unionDistinct ts

-- | creates a list of pairs out of a  tree
--
-- prop> toList t = zip (keys t) (values t)
toList :: RTree a -> [(MBB, a)]
toList Empty = []
toList (Leaf mbb x) = [(mbb, x)]
toList t = concatMap toList $ getChildren t

-- | returns all keys in this tree
--
-- prop> toList t = zip (keys t) (values t)
keys :: RTree a -> [MBB]
keys = foldWithMBB handleLeaf handleNode []
    where
    handleLeaf mbb _ = [mbb]
    handleNode _ xs  = concat xs

-- | returns all values in this tree
--
-- prop> toList t = zip (keys t) (values t)
values :: RTree a -> [a]
values = foldWithMBB handleLeaf handleNode []
    where
    handleLeaf _ x = [x]
    handleNode _ xs  = concat xs


-- ----------------------------------
-- insert 

-- | inserts an element whith the given 'MBB' and a value in a tree
insert :: MBB -> a -> RTree a -> RTree a
insert mbb e oldRoot = unionDistinct (singleton mbb e) oldRoot

-- | unifies left and right RTeee. Works only, if they don't contain common keys. Much faster than union, though. 
unionDistinct :: RTree a -> RTree a -> RTree a
unionDistinct Empty{} t           = t
unionDistinct t       Empty{}     = t
unionDistinct t1@Leaf{} t2@Leaf{}
    | on (==) getMBB t1 t2 = t1
    | otherwise = createNodeWithChildren [t1, t2] -- root case
unionDistinct left right
    | depth left > depth right              = unionDistinct right left
    | depth left == depth right             = fromList' $ (getChildren left) ++ [right]
    | (L.length $ getChildren newNode) > n = createNodeWithChildren $ splitNode newNode
    | otherwise                             = newNode
    where
    newNode = addLeaf left right

addLeaf :: RTree a -> RTree a -> RTree a
addLeaf left right 
    | depth left + 1 == depth right = node (left `unionMBB'` right) (left : nonEq)
    | otherwise                     = node (left `unionMBB'` right) newChildren
    where
    newChildren = findNodeWithMinimalAreaIncrease left (getChildren right)
    (eq, nonEq) = partition (on (==) getMBB left) $ getChildren right

findNodeWithMinimalAreaIncrease :: RTree a -> [RTree a] -> [RTree a]
findNodeWithMinimalAreaIncrease leaf children = splitMinimal xsAndIncrease
    where
--  xsAndIncrease :: [(RTree a, Double)]    
    xsAndIncrease = zip children ((areaIncreasesWith leaf) <$> children)
    minimalIncrease = minimum $ snd <$> xsAndIncrease
--  xsAndIncrease' :: [(RTree a, Double)]   
    splitMinimal [] = []
    splitMinimal ((t,mbb):xs)
        | mbb == minimalIncrease = unionDistinctSplit leaf t ++ (fst <$> xs)
        | otherwise            = t : splitMinimal xs

unionDistinctSplit :: RTree a -> RTree a -> [RTree a]
unionDistinctSplit leaf e
    | (L.length $ getChildren newLeaf) > n = splitNode newLeaf
    | otherwise = [newLeaf]
    where
    newLeaf = addLeaf leaf e

-- | /O(n²)/ solution
splitNode :: RTree a -> [RTree a]
splitNode Leaf{} = error "splitNode: Leaf"
splitNode e = [createNodeWithChildren x1, createNodeWithChildren x2]
    where
    (l, r) = findGreatestArea $ getChildren e
    (x1, x2) = quadSplit [l] [r] unfinished
    unfinished = filter (on (/=) getMBB l) $ filter (on (/=) getMBB r) $ getChildren e

findGreatestArea :: [RTree a] -> (RTree a, RTree a)
findGreatestArea xs = (x', y')
    where
    xs' = zip xs [(1::Int)..]
    listOfTripels = [(fst x, fst y, on unionMBB' fst x y) | x <- xs', y <- xs', ((<) `on` snd) x y]
    (x', y', _) = maximumBy (compare `on` (\(_,_,x) -> area x)) listOfTripels


quadSplit :: [RTree a] -> [RTree a] -> [RTree a] -> ([RTree a], [RTree a])
quadSplit left right [] = (left, right)
quadSplit left right unfinished
    | (L.length left) + (L.length unfinished)  <= m = (left ++ unfinished, right)
    | (L.length right) + (L.length unfinished) <= m = (left, right ++ unfinished)
    | isLeft''                                      = quadSplit (minimumElem : left) right  newRest
    | otherwise                                     = quadSplit left  (minimumElem : right) newRest
        where
--      makeTripel :: RTree a -> (RTree a, Bool, Double)
        makeTripel x = (x, isLeft, growth)
            where
            isLeft = (areaIncreasesWithLeft) < (areaIncreasesWithRight)
            growth = case isLeft of
                True -> areaIncreasesWithLeft
                False -> areaIncreasesWithRight
            areaIncreasesWithLeft  = (areaIncreasesWith x (createNodeWithChildren left))
            areaIncreasesWithRight = (areaIncreasesWith x (createNodeWithChildren right))
        (minimumElem, isLeft'', _) = minimumBy (compare `on` (\(_,_,g) -> g)) $ makeTripel <$> unfinished
        newRest = (filter (on (/=) getMBB minimumElem) unfinished)

--mergeNodes :: RTree a -> RTree a -> RTree a
--mergeNodes x@Node{} y@Node{} = node (unionMBB' x y) (on (++) getChildren x y)
--mergeNodes _ _               = error "no merge for Leafs"

-- ------------
-- helpers


areaIncreasesWith :: RTree a -> (RTree a) -> Double
areaIncreasesWith newElem current = newArea - currentArea
    where
    currentArea = area $ getMBB current
    newArea = area $ unionMBB' newElem current

-- -----------------
-- lookup

-- | returns the value if it exists in the tree
lookup :: MBB -> RTree a -> Maybe a
lookup _   Empty = Nothing
lookup mbb t@Leaf{}
    | mbb == getMBB t = Just $ getElem t
    | otherwise = Nothing
lookup mbb t = case founds of 
    [] -> Nothing
    x:_ -> Just x
    where
    matches = filter (\x -> (getMBB x) `containsMBB` mbb) $ getChildren t
    founds = catMaybes $ map (lookup mbb) matches

-- | returns all values, which are located in the given bounding box. 
lookupRange :: MBB -> RTree a -> [a]
lookupRange _ Empty = []
lookupRange mbb t@Leaf{}
    | mbb `containsMBB` (getMBB t) = [getElem t]
    | otherwise = []
lookupRange mbb t = founds
    where
    matches = filter intersectRTree $ getChildren t
    founds = concatMap (lookupRange mbb) matches
    intersectRTree x = isJust $ mbb `intersectMBB` (getMBB x)

-- -----------
-- delete

-- | Delete a key and its value from the RTree. When the key is not a member of the tree, the original tree is returned.
delete :: MBB -> RTree a -> RTree a
delete _   Empty = Empty
delete mbb  t@Leaf{} 
    | mbb == getMBB t = Empty
    | otherwise       = t
delete mbb root
    | L.length (getChildren newRoot) == 1 = head $ getChildren newRoot
    | otherwise                           = newRoot
    where
    newRoot = delete' mbb root


delete' :: MBB -> RTree a -> RTree a
delete' mbb  t@Leaf{} 
    | mbb == getMBB t = Empty
    | otherwise       = t
delete' mbb t = fromList' $ orphans ++ [newValidNode]
    where
    (matches, noMatches) = partition (\x -> (getMBB x) `containsMBB` mbb) $ getChildren t
    matches' = filter (not . null) $ map (delete' mbb) matches
    (orphans, validMatches) = foldr handleInvalid ([], []) matches'
--  handleInvalid :: RTree a -> ([RTree a], [RTree a]) -> ([RTree a], [RTree a])
    handleInvalid l@Leaf{} (orphans', validMatches') = (orphans', l:validMatches')
    handleInvalid invalidNode (orphans', validMatches')
        | L.length children < m = (children ++ orphans', validMatches')
        | otherwise = (orphans', invalidNode:validMatches')
        where
        children = getChildren invalidNode
    newValidNode = createNodeWithChildren $ validMatches ++ noMatches

-- ---------------
foldWithMBB :: (MBB -> a -> b) -> (MBB -> [b] -> b) -> b -> RTree a -> b
foldWithMBB _ _ n' Empty    = n'
foldWithMBB f _ _ t@Leaf{} = f (getMBB t) (getElem t)
foldWithMBB f g n' t        = g (getMBB t) $ foldWithMBB f g n' <$> (getChildren t)

-- | unifies the first and the second tree into one.
union :: RTree a -> RTree a -> RTree a
union Empty Empty = Empty
union t1 t2
    | depth t1 <= depth t2 = foldr (uncurry insert) t2 (toList t1)
    | otherwise            = union t2 t1

-- | map, which also filters Nothing values
mapMaybe :: (a -> Maybe b) -> RTree a -> RTree b
mapMaybe f t = fromList $ Maybe.mapMaybe func $ toList t
    where
    func (mbb,x) = case f x of
            Nothing -> Nothing
            Just x' -> Just (mbb, x')

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

isValid :: Show b => b -> RTree a -> Bool
isValid _ Leaf{} = True
isValid context x = case L.length c >= m && L.length c <= n && (and $ (isValid context) <$> c) && (isBalanced x) of
    True -> True
    False -> error ( "invalid " ++ show (L.length c) ++ " " ++ show context )
    where
    isBalanced :: RTree a -> Bool 
    isBalanced (Leaf _ _ ) = True
    isBalanced x' = (and $ isBalanced <$> c') && (and $ (== depth (head c')) <$> (depth <$> c'))
        where
        c' = getChildren x'
    c = getChildren x

-- ----------------------
i_ :: String
i_ = "    "

pp :: (Show a) => RTree a -> IO ()
pp = pp' ""

pp' :: (Show a) => String -> RTree a -> IO ()
pp' i Empty                 = putStrLn $ i ++ "Empty"
pp' i (Leaf mbb x)          = putStrLn $ i ++ "Leaf " ++ (show mbb) ++ " " ++ (show x)
pp' i (Node mbb cs)         = do
    putStrLn $ i ++ "Node " ++ (show mbb)
    mapM_ (pp' (i ++ i_)) cs
pp' i (Node2 mbb c1 c2)       = do
    putStrLn $ i ++ "Node2 " ++ (show mbb)
    mapM_ (pp' (i ++ i_)) [c1, c2]
pp' i (Node3 mbb c1 c2 c3)    = do
    putStrLn $ i ++ "Node3 " ++ (show mbb)
    mapM_ (pp' (i ++ i_)) [c1, c2, c3]
pp' i (Node4 mbb c1 c2 c3 c4) = do
    putStrLn $ i ++ "Node4 " ++ (show mbb)
    mapM_ (pp' (i ++ i_)) [c1, c2, c3, c4]

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

depth :: RTree a -> Int
depth (Leaf _ _ ) = 0
depth t = 1 + (depth $ head $ getChildren t)

-- | returns the number of elements in a tree
length :: RTree a -> Int 
length Empty = 0
length (Leaf {}) = 1
length t = sum $ length <$> (getChildren t)

--delete' :: MBB -> RTree a -> Either (RTree a) [(MBB, a)]

instance NFData a => NFData (RTree a) where
    rnf (Empty)               = ()
    rnf (Leaf _ e)            = {-rnf m `seq`-} rnf e
    rnf (Node _ cs)           = {-rnf m `seq`-} rnf cs
    rnf (Node2 _ c1 c2)       = {-rnf m `seq`-} rnf c1 `seq` rnf c2
    rnf (Node3 _ c1 c2 c3)    = {-rnf m `seq`-} rnf c1 `seq` rnf c2 `seq` rnf c3
    rnf (Node4 _ c1 c2 c3 c4) = {-rnf m `seq`-} rnf c1 `seq` rnf c2 `seq` rnf c3 `seq` rnf c4


instance  (Binary a) => Binary (RTree a)