content: {-# LANGUAGE ScopedTypeVariables #-}
  
  module Test.SmartCheck.DataToTree
    ( forestReplaceChildren
    , getAtIdx
    , replaceAtIdx
    , getIdxForest
    , breadthLevels
    , mkSubstForest
    , depth
    , tooDeep
    ) where
  
  import Test.SmartCheck.Types
  
  import Data.Tree
  import Data.List
  import Data.Maybe
  import Data.Typeable
  
  --------------------------------------------------------------------------------
  -- Operations on Trees and Forests.
  --------------------------------------------------------------------------------
  
  -- | Return the list of values at each level in a Forest Not like levels in
  -- Data.Tree (but what I imagined it should have done!).
  breadthLevels :: Forest a -> [[a]]
  breadthLevels forest =
    takeWhile (not . null) go
    where
    go = map (getLevel forest) [0..]
  
  --------------------------------------------------------------------------------
  
  -- | Return the elements at level i from a forest.  0-based indexing.
  getLevel :: Forest a -> Int -> [a]
  getLevel fs 0 = map rootLabel fs
  getLevel fs n = concatMap (\fs' -> getLevel (subForest fs') (n-1)) fs
  
  --------------------------------------------------------------------------------
  
  -- | Get the depth of a Forest.  0-based (an empty Forest has depth 0).
  depth :: Forest a -> Int
  depth forest = if null ls then 0 else maximum ls
    where
    ls = map depth' forest
    depth' (Node _ [])      = 1
    depth' (Node _ forest') = 1 + depth forest'
  
  --------------------------------------------------------------------------------
  
  -- | How many members are at level i in the Tree?
  levelLength :: Int -> Tree a -> Int
  levelLength 0 t = length (subForest t)
  levelLength n t = sum $ map (levelLength (n-1)) (subForest t)
  
  --------------------------------------------------------------------------------
  
  -- | Get the tree at idx in a forest.  Nothing if the index is out-of-bounds.
  getIdxForest :: Forest a -> Idx -> Maybe (Tree a)
  getIdxForest forest (Idx (0 :: Int) n) =
    if length forest > n then Just (forest !! n)
      else Nothing
  getIdxForest forest idx              =
    -- Should be a single Just x in the list, holding the value.
    listToMaybe . catMaybes . snd $ acc
  
    where
    acc = mapAccumL findTree (column idx) (map Just forest)
  
    l = level idx - 1
    -- Invariant: not at the right level yet.
    findTree :: Int -> Maybe (Tree a) -> (Int, Maybe (Tree a))
    findTree n Nothing  = (n, Nothing)
    findTree n (Just t) =
      let len = levelLength l t in
      if n < 0 -- Already found index
        then (n, Nothing)
        else if n < len -- Big enough to index, so we climb down this one.
               then let t' = getIdxForest (subForest t) (Idx l n) in
                    (n-len, t')
               else (n-len, Nothing)
  
  --------------------------------------------------------------------------------
  
  -- Morally, we should be using generic zippers and a nice, recursive breadth-first search function, e.g.
  
  {-
  
  data Tree = N Int Tree Tree
            | E
  
  index :: Int -> Tree -> Tree
  index = index' []
    where
    index' :: [Tree] -> Int -> Tree -> Tree
    index' _      0   t           = t
    index' []     idx (N i t0 t1) = index' [t1]             (idx-1) t0
    index' (k:ks) idx E           = index' ks               (idx-1) k
    index' (k:ks) idx (N i t0 t1) = index' (ks ++ [t0, t1]) (idx-1) k
  
  -}
  
  -- | Returns the value at index idx.  Returns nothing if the index is out of
  -- bounds.
  getAtIdx :: SubTypes a
           => a         -- ^ Value
           -> Idx       -- ^ Index of hole
           -> Maybe Int -- ^ Maximum depth we want to extract
           -> Maybe SubT
  getAtIdx d Idx { level = l, column = c } maxDepth
    | tooDeep l maxDepth = Nothing
    | length lev > c     = Just (lev !! c)
    | otherwise          = Nothing
    where
    lev = getLevel (subTypes d) l
  
  --------------------------------------------------------------------------------
  
  tooDeep :: Int -> Maybe Int -> Bool
  tooDeep l = maybe False (l >)
  
  --------------------------------------------------------------------------------
  
  data SubStrat = Parent   -- ^ Replace everything in the path from the root to
                           -- here.  Used as breadcrumbs to the value.  Chop the
                           -- subforest.
                | Children -- ^ Replace a value and all of its subchildren.
    deriving  (Show, Read, Eq)
  
  --------------------------------------------------------------------------------
  
  forestReplaceParent, forestReplaceChildren :: Forest a -> Idx -> a -> Forest a
  forestReplaceParent   = sub Parent
  forestReplaceChildren = sub Children
  
  --------------------------------------------------------------------------------
  
  sub :: SubStrat -> Forest a -> Idx -> a -> Forest a
  -- on right level, and we'll assume correct subtree.
  sub strat forest (Idx (0 :: Int) n) a =
    snd $ mapAccumL f 0 forest
    where
    f i node | i == n    = ( i+1, news )
             | otherwise = ( i+1, node )
  
      where
      news = case strat of
               Parent   -> Node a []
               Children -> fmap (const a) (forest !! n)
  
  sub strat forest idx a =
    snd $ mapAccumL findTree (column idx) forest
    where
    l = level idx - 1
    -- Invariant: not at the right level yet.
    findTree n t
      -- Already found index
      | n < 0     = (n, t)
      -- Big enough to index, so we climb down this one.
      | n < len   = (n-len, newTree)
      | otherwise = (n-len, t)
      where
      len = levelLength l t
      newTree = Node newRootLabel (sub strat (subForest t) (Idx l n) a)
      newRootLabel = case strat of
                       Parent   -> a
                       Children -> rootLabel t
  
  --------------------------------------------------------------------------------
  -- Operations on SubTypes.
  --------------------------------------------------------------------------------
  
  -- | Make a substitution Forest (all proper children).  Initially we don't
  -- replace anything.
  mkSubstForest :: SubTypes a => a -> b -> Forest b
  mkSubstForest a b = map tMap (subTypes a)
    where tMap = fmap (const b)
  
  --------------------------------------------------------------------------------
  
  -- | Replace a value at index idx generically in a Tree/Forest generically.
  replaceAtIdx :: (SubTypes a, Typeable b)
               => a     -- ^ Parent value
               -> Idx   -- ^ Index of hole to replace
               -> b     -- ^ Value to replace with
               -> Maybe a
  replaceAtIdx m idx = replaceChild m (forestReplaceParent subF idx Subst)
    where
    subF = mkSubstForest m Keep
  
  --------------------------------------------------------------------------------
  