{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE GADTs #-}
module ArbitraryEnzymeTree
  ( EnzymeTree
  , getTreeSize
  , getElement
  , getShape
  , getSubTrees
  , treeFromList
  , generateTreeFromList
  , generateDotFromTree
  , generateDotFromPoisonTree
  , poisonTree
  , makeHead
  ) where

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

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

getElement :: EnzymeTree s a -> a
getElement (EnzymeTree _ e _) = e

getSubTrees :: EnzymeTree s a -> [EnzymeTree TreeShape a]
getSubTrees (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

generateTreeFromList :: Int -> [a] -> IO (EnzymeTree s a)
generateTreeFromList s a = flip treeFromList a . getShape <$> generate (resize s arbitrary :: Gen (EnzymeTree s ()))


generateDotFromTree :: (Show a, IsString b, Monoid b) => b -> EnzymeTree s a -> b
generateDotFromTree name t = "digraph " <> name <> " {\n"
                             <> generateDotFromTree' t
                             <> "}\n"
  where
    generateDotFromTree' :: (Show a, IsString b, Monoid b) => EnzymeTree s a -> b
    generateDotFromTree' (EnzymeTree _ num ns) = 
        "    " <> wrap (ts num) <> " -> { "
                   <> mconcat (intersperse " "  (wrap . ts . getElement <$> ns))
               <> " };\n"
               <> mconcat (generateDotFromTree' <$> ns)
      where 
        ts = fromString . show
        wrap x = "\"" <> x <> "\""

generateDotFromPoisonTree :: (Show a, IsString b, Monoid b) => b -> Double -> EnzymeTree s (Double,a) -> b
generateDotFromPoisonTree name pl t = "digraph " <> name <> " {\n"
                             <> generateDotFromTree' t
                             <> "}\n"
  where
    generateDotFromTree' :: (Show a, IsString b, Monoid b) => EnzymeTree s (Double,a) -> b
    generateDotFromTree' (EnzymeTree _ (d,c) ns) = 
        "    " <> wrap (ts c) <> " -> { "
                   <> mconcat (intersperse " "  (wrap . ts . snd . getElement <$> ns))
               <> " };\n"
               <> (if d > pl then "    " <> wrap (ts c) <> " [style=filled, fillcolor=\"0," <> ts d <> ",0.9\"];\n" else "")
               <> mconcat (generateDotFromTree' <$> ns)
      where 
        ts :: (Show a, IsString b) => a -> b
        ts = fromString . show
        wrap x = "\"" <> x <> "\""

poisonTree :: [Double] -> EnzymeTree s a -> EnzymeTree s (Double, a)
poisonTree ds t@(EnzymeTree s _ _) = go Nothing 0 annotatedTree
  where
      annotatedTree = (,) <$> treeFromList s ds <*> t
      go :: Maybe Double -> Int -> EnzymeTree t (Double, a) -> EnzymeTree t (Double, a)
      go Nothing   i parent@(EnzymeTree s' (p,a) childs) = EnzymeTree s' (p/5, a) $ (\(_,ts) -> go (Just         $ p/5 ) (i+1) ts) <$> zip [1..] childs
      go (Just pe) i   this@(EnzymeTree s' (p,a) childs) = EnzymeTree s' (p',a)   $ (\(j,ts) -> go (Just . min 1 $ j*p') (i+1) ts) <$> zip [1..] childs
          where
              p' = max pe (p / i')
              i' = fromIntegral $ 6 - min i 5 -- 100% effective poision only at level xx or deeper

makeHead :: a -> EnzymeTree s a -> EnzymeTree s a
makeHead c (EnzymeTree s a ts) = EnzymeTree s c ts
