-- | Functions for creating rose trees (from "Data.Tree") of a specified
-- size.
module Data.Tree.Replicate where

import           Control.Monad
import           Data.Tree

-- | @'replicateA' n x@ replicates the action @x@ @n@ times.
replicateA :: Applicative f => Int -> f a -> f (Tree a)
replicateA t x = go t
  where
    go n
      | n <= 1 = flip Node [] <$> x
    go n =
        let m =
                head
                    [ y
                    | y <- [1 ..]
                    , y * y >= (n - 1) ]
            lm = (n - 1) `div` m
        in if m * lm == (n - 1)
               then Node <$> x <*> replicateM lm (go m)
               else Node <$> x <*>
                    ((:) <$> go ((n - 1) - m * lm) <*> replicateM lm (go m))

-- | @'replicateTree' n a@ creates a tree of size @n@ filled with @a@.
--
-- >>> putStr (drawTree (replicateTree 4 "."))
-- .
-- |
-- +- .
-- |
-- `- .
--    |
--    `- .
--
-- prop> n > 0 ==> length (replicateTree n x) == n
replicateTree :: Int -> a -> Tree a
replicateTree t x = go t where
 go n | n <= 1 = Node x []
 go n =
   let m = head [ y | y <- [1..], y * y >= (n-1) ]
       lm = (n-1) `div` m
   in if m * lm == (n-1)
      then Node x (replicate lm (go m))
      else Node x (go ((n-1) - m * lm) : replicate lm (go m))