-- | Auxillary functions useful for testing

{-# LANGUAGE CPP, 
             DeriveFunctor, DeriveFoldable, DeriveTraversable, StandaloneDeriving,
             FlexibleInstances, TypeSynonymInstances
  #-}
module TestSuite.Tools where

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

import Control.Applicative
import Control.Monad hiding (mapM, mapM_, forM, forM_)
import Data.List (sort)
import Data.Foldable
import Data.Traversable
import Prelude hiding (foldl,foldr,mapM,mapM_,concat,concatMap)

import Text.Show
import Text.Read

import Data.Generics.Fixplate.Base
-- import Data.Generics.Fixplate.Misc

import Test.QuickCheck
import TestSuite.Misc

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

maxChildren :: Int
maxChildren = 7

data Tree label
  = Tree label [Tree label] 
  deriving (Eq,Ord,Show,Read,Functor,Foldable,Traversable)

data TreeF label t 
  = TreeF label [t]
  deriving (Eq,Ord,Show,Read,Functor,Foldable,Traversable)
  
type FixT label = Mu (TreeF label)  
  
instance Eq   label => EqF   (TreeF label) where equalF     = (==)
instance Ord  label => OrdF  (TreeF label) where compareF   = compare
instance Show label => ShowF (TreeF label) where showsPrecF = showsPrec
#ifdef __GLASGOW_HASKELL__
instance Read label => ReadF (TreeF label) where readPrecF  = readPrec
#else
instance Read label => ReadF (TreeF label) where readsPrecF = readsPrec
#endif
  
treeF :: l -> [Mu (TreeF l)] -> Mu (TreeF l)
treeF s = Fix . TreeF s

attrTreeF :: a -> l -> [Attr (TreeF l) a] -> Attr (TreeF l) a
attrTreeF x s = Fix . Ann x . TreeF s

--------------------------------------------------------------------------------
-- * draw trees

printTree :: Tree Label -> IO ()
printTree = printTree' (\(Label s) -> s)

printTreeF :: FixT Label -> IO ()
printTreeF = printTreeF' (\(Label s) -> s)

printTree' :: (a -> String) -> Tree a -> IO ()
printTree' h = go 0 where
  go i (Tree label children) = do
    putStrLn $ if i>0 
      then concat (replicate (i-1) "| " ++ ["|-", h label])
      else h label
    mapM_ (go (i+1)) children

printTreeF' :: (a -> String) -> Mu (TreeF a) -> IO ()
printTreeF' h = go 0 where
  go i (Fix (TreeF label children)) = do
    putStrLn $ if i>0 
      then concat (replicate (i-1) "| " ++ ["|-", h label])
      else h label
    mapM_ (go (i+1)) children
    
--------------------------------------------------------------------------------
-- * random trees

rndTree :: IO (Tree Label)
rndTree = liftM (!!7) $ sample' arbitrary

rndFixT :: IO (FixT Label)
rndFixT = liftM (!!7) $ sample' arbitrary

--------------------------------------------------------------------------------
-- * conversion

toFixT :: Tree l -> Mu (TreeF l)
toFixT (Tree s ts) = treeF s (map toFixT ts)

fromFixT :: FixT l -> Tree l
fromFixT (Fix (TreeF s ts)) = Tree s (map fromFixT ts)

fromAttr :: Attr (TreeF l) a -> Tree (l,a)
fromAttr (Fix (Ann x (TreeF s ts))) = Tree (s,x) (map fromAttr ts)

toAttr :: Tree (l,a) -> Attr (TreeF l) a 
toAttr (Tree (s,x) ts) = Fix (Ann x (TreeF s (map toAttr ts)))

--------------------------------------------------------------------------------
-- * arbitrary

pairs :: [a] -> [(a,a)]
pairs (x:xs@(y:_)) = (x,y):(pairs xs)
pairs [_]          = []
pairs []           = error "pairs: empty list"

-- | @genPartition n k@ partitions n elements into k groups randomly,
-- and gives back the sizes (which can be zero, too)
genPartition :: Int -> Int -> Gen [Int]
genPartition n k = do
  sep <- replicateM (k-1) $ choose (0,n)
  let ps = pairs (0 : sort sep ++ [n]) 
  return (map (\(x,y) -> (y-x)) ps)

newtype Label = Label String deriving (Eq,Ord,Show,Read)

unLabel :: Label -> String
unLabel (Label s) = s

instance Arbitrary Label where
  arbitrary = do
    n <- choose (2, 8)
    liftM Label $ vectorOf n $ oneof [ choose ('a','z') , choose ('A','Z') ]
  
instance Arbitrary l => Arbitrary (Tree l) where
  shrink (Tree s sub) = [ Tree s sub' | sub' <- shrink sub ] 
  arbitrary = sized mkTree where
    mkTree n = do
      s <- arbitrary
      case n of
        0 -> return (Tree s [])
        1 -> mkTree 0 >>= \t -> return (Tree s [t])
        _ -> do
          k <- choose (1, min maxChildren n)
          ls <- genPartition (n-1) k
          subtrees <- forM ls $ \l -> mkTree l
          return (Tree s subtrees)

instance Arbitrary l => Arbitrary (Mu (TreeF l)) where
  shrink (Fix (TreeF s sub)) = [ Fix (TreeF s sub') | sub' <- shrink sub ] 
  arbitrary = sized mkTree  where
    mkTree n = do
      s <- arbitrary
      case n of
        0 -> return (treeF s [])
        1 -> mkTree 0 >>= \t -> return (treeF s [t])
        _ -> do
          k <- choose (1, min maxChildren n)
          ls <- genPartition (n-1) k
          subtrees <- forM ls $ \l -> mkTree l
          return (treeF s subtrees)

{-          
instance (Arbitrary a, Arbitrary x) => Arbitrary (Ann TreeF a x) where
  shrink (Ann a x) = [ Ann a y | y <- shrink x ]
  arbitrary = do
    a <- arbitrary
    x <- arbitrary
-}

instance (Arbitrary a, Arbitrary l) => Arbitrary (Attr (TreeF l) a) where
  shrink (Fix (Ann a (TreeF s sub))) = [ Fix (Ann a (TreeF s sub')) | sub' <- shrink sub ] 
  arbitrary = sized mkTree  where
    mkTree n = do
      s <- arbitrary
      a <- arbitrary
      case n of
        0 -> return (attrTreeF a s [])
        1 -> mkTree 0 >>= \t -> return (attrTreeF a s [t])
        _ -> do
          k <- choose (1, min maxChildren n)
          ls <- genPartition (n-1) k
          subtrees <- forM ls $ \l -> mkTree l
          return (attrTreeF a s subtrees)
              
--------------------------------------------------------------------------------