-- | Tests for tries

module TestSuite.Trie where

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

import Test.Tasty
import Test.Tasty.QuickCheck
import Test.QuickCheck

import Prelude hiding ( mapM , mapM_ , concat , concatMap , lookup )

import Control.Applicative
import Control.Monad hiding ( mapM , mapM_ , forM , forM_ )
import Data.Foldable hiding ( toList )
import Data.Traversable
import Data.List ( sort , group , nubBy , nub , (\\) , foldl' )

import Data.Map (Map)
import qualified Data.Map as Map 

import Data.Generics.Fixplate.Base
-- import Data.Generics.Fixplate.Misc
import Data.Generics.Fixplate.Attributes
import Data.Generics.Fixplate.Morphisms
import Data.Generics.Fixplate.Traversals
import Data.Generics.Fixplate.Trie as Trie

import TestSuite.Tools
import TestSuite.Misc

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

testgroup_Trie :: TestTree
testgroup_Trie = testGroup "Trie"
  [ testProperty "difference"        prop_difference
  , testProperty "differenceWith"    prop_differenceWith
  , testProperty "union"             prop_union
  , testProperty "intersection"      prop_intersection

  , testProperty "universeBag"       prop_unibag_naive
  , testProperty "universeBag /2"    prop_unibag_naive_2
  , testProperty "christmasTree"     prop_christmasTree
  , testProperty "christmasTree /2"  prop_christmasTree_2
  , testProperty "christmasTree /3"  prop_christmasTree_3

  , testProperty "fromList"          prop_fromList_naive
  , testProperty "bag"               prop_bag
  , testProperty "bag /b"            prop_bag_b
  , testProperty "toList . fromList" prop_fromList_toList
  , testProperty "multiSetToList"    prop_multiSetToList_b

  , testProperty "insert"            prop_insert
  , testProperty "delete"            prop_delete
  , testProperty "update"            prop_update
  , testProperty "delete . insert"   prop_insert_delete
  , testProperty "insert . delete"   prop_delete_insert
  , testProperty "lookup"            prop_lookup
  , testProperty "lookup /notfound"  prop_lookup_notfound
  , testProperty "singleton"         prop_singleton
  ]

--------------------------------------------------------------------------------
-- Tests

{-
runtests_Trie :: IO ()
runtests_Trie = do
  quickCheck prop_difference
  quickCheck prop_differenceWith
  quickCheck prop_union
  quickCheck prop_intersection

  quickCheck prop_unibag_naive
  quickCheck prop_unibag_naive_2
  quickCheck prop_christmasTree
  quickCheck prop_christmasTree_2
  quickCheck prop_christmasTree_3

  quickCheck prop_fromList_naive
  quickCheck prop_bag
  quickCheck prop_bag_b
  quickCheck prop_fromList_toList
  quickCheck prop_multiSetToList_b
  quickCheck prop_insert
  quickCheck prop_delete
  quickCheck prop_update
  quickCheck prop_insert_delete
  quickCheck prop_delete_insert
  quickCheck prop_lookup
  quickCheck prop_lookup_notfound
  quickCheck prop_singleton
-}

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

newtype Multiplicity = Multiplicity { unMultiplicity :: Int } deriving (Eq,Ord,Show)

instance Arbitrary Multiplicity where
  arbitrary = do
    n <- choose (1, 7)
    return (Multiplicity n)

newtype MultiSet = MultiSet { unMultiSet :: [(Multiplicity, FixT Label)] } deriving (Eq,Ord,Show)

instance Arbitrary MultiSet where arbitrary = MultiSet <$> arbitrary

multiSetToList :: MultiSet -> [FixT Label]
multiSetToList (MultiSet mxs) = go mxs where
  go [] = []
  go ((Multiplicity n, x):rest) = replicate n x ++ go rest

multiSetToList_b :: MultiSet -> [FixT Label]
multiSetToList_b (MultiSet mxs) = go mxs [] where
  go [] [] = []
  go [] ys = go ys []
  go ((Multiplicity n, x):rest) ys = if n>0 
    then x : go rest ( (Multiplicity (n-1), x) : ys )
    else go rest ys

newtype FiniteMap = FiniteMap { unFiniteMap :: [(FixT Label,Char)] } deriving (Eq,Ord,Show)

instance Arbitrary FiniteMap where arbitrary = (FiniteMap . nubBy (equating fst)) <$> arbitrary

type TrieT = Trie (TreeF Label) Char

finiteMap :: FiniteMap -> TrieT
finiteMap (FiniteMap fmap) = fromList fmap

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

fromListNaive :: (Traversable f, OrdF f) => [(Mu f, a)] -> Trie f a
fromListNaive ts = Prelude.foldl worker Trie.empty ts where
  worker trie (tree,value) = Trie.insertWith id const tree value trie 

universeBagNaive :: (Functor f, Foldable f, OrdF f) => Mu f -> Trie f Int
universeBagNaive = bag . universe

mapBag :: Ord a => [a] -> Map a Int
mapBag xs = Data.List.foldl' f Map.empty xs where
  f old x = Map.insertWith (+) x 1 old

--------------------
    
prop_unibag_naive :: FixT Label -> Bool
prop_unibag_naive tree = toList (universeBag tree) == toList (universeBagNaive tree)

prop_unibag_naive_2 :: FixT Bool -> Bool
prop_unibag_naive_2 tree = toList (universeBag tree) == toList (universeBagNaive tree)

prop_fromList_naive :: FiniteMap -> Bool
prop_fromList_naive (FiniteMap list) = toList (fromList list) == toList (fromListNaive list)

prop_bag :: MultiSet -> Bool
prop_bag mset = (sort $ toList $ bag $ multiSetToList mset) == sort (map f $ unMultiSet mset) where
  f (Multiplicity k, x) = (x,k)

prop_bag_b :: MultiSet -> Bool
prop_bag_b mset = (sort $ toList $ bag $ multiSetToList_b mset) == sort (map f $ unMultiSet mset) where
  f (Multiplicity k, x) = (x,k)

prop_fromList_toList :: FiniteMap -> Bool
prop_fromList_toList (FiniteMap list) = sort (toList (fromList list)) == sort list

prop_multiSetToList_b :: MultiSet -> Bool
prop_multiSetToList_b mset = toList (bag (multiSetToList mset)) == toList (bag (multiSetToList_b mset)) 

prop_insert :: FixT Label -> Char -> FiniteMap -> Bool
prop_insert key ch (FiniteMap list) = sort (toList (insert key ch trie)) == sort ((key,ch) : toList trie) where
  trie = fromList list

prop_delete :: Int -> FiniteMap -> Bool
prop_delete i (FiniteMap list) = (n==0) || (toList (delete key trie) == toList trie \\ [(key,value)]) where
  trie = fromList list
  n = length list
  k = mod i n
  (key,value) = list!!k

prop_update :: Char -> Int -> FiniteMap -> Bool
prop_update new i (FiniteMap list) = (n==0) || (toList (update f key trie) == replace (toList trie)) where
  trie = fromList list
  n = length list
  k = mod i n
  (key,value) = list!!k
  replace [] = []
  replace (this@(k,x):rest) = if k==key 
    then case f x of 
      Nothing -> rest
      Just y  -> (k,y) : replace rest
    else this : replace rest    
  f old = if old < 'A' then Nothing else Just new

prop_insert_delete :: FixT Label -> Char -> FiniteMap -> Bool
prop_insert_delete key ch (FiniteMap list) = toList (delete key (insert key ch trie)) == toList trie where
  trie = delete key (fromList list)   -- !

prop_delete_insert :: Int -> FiniteMap -> Bool
prop_delete_insert i (FiniteMap list) = (n==0) || (toList (insert key value (delete key trie)) == toList trie) where
  trie = fromList list
  n = length list
  k = mod i n
  (key,value) = list!!k

prop_lookup :: Int -> FiniteMap -> Bool
prop_lookup i (FiniteMap list) = (n==0) || (Just value == lookup key trie) where
  trie = fromList list
  n = length list
  k = mod i n
  (key,value) = list!!k

prop_lookup_notfound :: FixT Label -> FiniteMap -> Bool
prop_lookup_notfound key (FiniteMap list) = lookup key trie == Nothing where
  trie = delete key (fromList list)   -- !#endif

prop_singleton :: FixT Label -> Char -> Bool
prop_singleton tree ch = toList (singleton tree ch) == [(tree,ch)]

prop_intersection :: MultiSet -> Bool
prop_intersection mset = {- trace ("--"++show n++"--") -} (itrie == imap) where

  list = multiSetToList_b mset
  n = length list
  k = div n 3
  l = div (2*n) 3
  xs = take l list
  ys = drop k list

  itrie = sort $     toList $     intersectionWith (+) (   bag xs) (   bag ys)
  imap  = sort $ Map.toList $ Map.intersectionWith (+) (mapBag xs) (mapBag ys) 

prop_union :: MultiSet -> Bool
prop_union mset = {- trace ("--"++show n++"--") -} (utrie == umap) where

  list = multiSetToList_b mset
  n = length list
  k = div n 3
  l = div (2*n) 3
  xs = take l list
  ys = drop k list

  utrie = sort $     toList $     unionWith (+) (   bag xs) (   bag ys)
  umap  = sort $ Map.toList $ Map.unionWith (+) (mapBag xs) (mapBag ys) 

prop_difference :: MultiSet -> Bool
prop_difference mset = {- trace ("--"++show [length xs , length ys, length dtrie]++"--") -} (dtrie == dmap) where

  list = multiSetToList_b mset
  n = length list
  k = div n 3
  l = div (2*n) 3
  xs = take l list
  ys = drop k list

  dtrie = sort $     toList $     difference (   bag xs) (   bag ys)
  dmap  = sort $ Map.toList $ Map.difference (mapBag xs) (mapBag ys) 

prop_differenceWith :: MultiSet -> Bool
prop_differenceWith mset = {- trace ("--"++show [length xs , length ys, length dtrie]++"--") -} (dtrie == dmap) where

  list = multiSetToList_b mset
  n = length list
  k = div n 3
  l = div (2*n) 3
  xs = take l list
  ys = drop k list

  f x y = if y<=2 then Just (x+1) else Nothing

  dtrie = sort $     toList $     differenceWith f (   bag xs) (   bag ys)
  dmap  = sort $ Map.toList $ Map.differenceWith f (mapBag xs) (mapBag ys) 

prop_christmasTree :: FixT Label -> Bool
prop_christmasTree tree = toList (attribute (christmasTree tree)) == toList (universeBag tree)

prop_christmasTree_3 :: FixT Bool -> Bool
prop_christmasTree_3 tree = toList (attribute (christmasTree tree)) == toList (universeBag tree)

-- we reduce the labels so that there is more chance for collisions
prop_christmasTree_2 :: Bool -> FixT Label -> Bool
prop_christmasTree_2 b tree0 = toList (attribute (christmasTree tree)) == toList (universeBag tree) where
  tree = transform f tree0
  f = if b 
    then \(Fix (TreeF (Label label) ts)) -> Fix $ TreeF (Label (take 1 label)) ts
    else \(Fix (TreeF (Label label) ts)) -> Fix $ TreeF (Label ""            ) ts

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