{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveGeneric #-}
module Arbitrary where

import Test.QuickCheck
import GHC.Generics (Generic(..))
import Text.Pretty.Simple
import Data.List

data TreeShape = Root
               { treeSize :: Int
               , treeSubSizes :: [TreeShape]
               } deriving (Show, Eq)

data EnzymeTree s a = EnzymeTree TreeShape a [EnzymeTree TreeShape a] deriving (Show, Eq, Generic, Functor, Foldable, Traversable)

instance Applicative (EnzymeTree s) where
    pure a = EnzymeTree (Root 1 []) a []
    EnzymeTree s f fs <*> EnzymeTree _ a as = EnzymeTree s (f a) (zipWith (<*>) fs as)

instance Arbitrary a => Arbitrary (EnzymeTree s a) where
    arbitrary = sized arbitrarySizedEnzymeTree

arbitrarySizedEnzymeTree :: Arbitrary a => Int -> Gen (EnzymeTree s a)
arbitrarySizedEnzymeTree m = do
  t <- arbitrary
  n <- choose (m `div` 4, m)
  ts <- if n == 0 then return []
                  else sequenceA $ replicate (2*n) (arbitrarySizedEnzymeTree (n `div` 2))
  let (_,ts') = foldr (go m) (0,[]) ts
      go :: Int -> EnzymeTree s a -> (Int, [EnzymeTree s a]) -> (Int, [EnzymeTree s a])
      go m x (s,ts)
        | m == 0             = (s,ts)
        | s + getTreeSize x > m = (s,ts)
        | otherwise          = (s + getTreeSize x, x:ts)

      sz = sum $ getTreeSize <$> ts'
      ss = (\(EnzymeTree a _ _) -> a) <$> ts'
  return (EnzymeTree (Root (sz+1) ss) t ts')

getTreeSize :: EnzymeTree s a -> Int
getTreeSize (EnzymeTree (Root a _) _ _) = a

getShape :: EnzymeTree s a -> TreeShape
getShape (EnzymeTree s _ _) = s

treeFromList :: TreeShape -> [a] -> EnzymeTree s a
treeFromList (Root n ns) (a:as) = EnzymeTree (Root n ns) a (unfoldr go (ns,as))
  where
      go :: ([TreeShape],[a]) -> Maybe (EnzymeTree s a,([TreeShape],[a]))
      go (n:ns,as) = Just (treeFromList n (take (treeSize n) as), (ns, drop (treeSize n) as))
      go ([],_)    = Nothing