-- | A class for tree types and representations of selections on tree types, as well as functions for converting between text and tree selections.
module Language.GroteTrap.Trees (

  -- * Paths and navigation
  Path, root,
  Nav, up, into, down, left, right, sibling,
  
  -- * Tree types
  Tree(..), followM, follow, depth, selectDepth, flatten,
  
  -- * Tree selections
  Selectable(..), TreeSelection,
  select, allSelections, selectionToRange, rangeToSelection, posToPath, isValidRange,
  
  -- * Suggesting and fixing
  suggest, repair

  ) where

import Language.GroteTrap.Range

import Data.List (sortBy, findIndex)
import Data.Maybe (isJust)
import Control.Monad.Error ()


------------------------------------
-- Paths and navigation
------------------------------------

-- | A path in a tree. Each integer denotes the selection of a child; these indices are 0-relative.
type Path  =  [Int]

-- | @root@ is the empty path.
root :: Path
root = []

-- | Navigation transforms one path to another.
type Nav = Path -> Path

-- | Move up to parent node. Moving up from root has no effect.
up :: Nav
up [] = []
up path = init path

-- | Move down into the nth child node.
into    ::  Int -> Nav
into i  =   (++[i])

-- | Move down into first child node.
down :: Nav
down = into 0

-- | Move left one sibling.
left :: Nav
left = sibling (-1)

-- | Move right one sibling.
right :: Nav
right = sibling 1

-- | Move @n@ siblings (@n@ can be negative).
sibling      ::  Int -> Nav
sibling 0 p  =   p  -- because sibling 0 [] == []
sibling d p  =   if newindex < 0 then p else into newindex parent
  where  index     = last p
         newindex  = index + d
         parent    = up p


------------------------------------
-- Parents and children
------------------------------------

-- | Tree types.
class Tree p where
  -- | Yields this tree's subtrees.
  children :: p -> [p]

-- | Breadth-first, pre-order traversal.
flatten :: Tree t => t -> [t]
flatten t = t : concatMap flatten (children t)

-- | Follows a path in a tree, returning the result in a monad.
followM :: (Monad m, Tree t) => t -> Path -> m t
followM parent [] = return parent
followM parent (t:ts) = do
  c <- childM parent t
  followM c ts

-- | Moves down into a child.
childM :: (Monad m, Tree p) => p -> Int -> m p
childM t i = if i >= 0 && i < length cs
              then return (cs !! i)
              else fail ("child " ++ show i ++ " does not exist")
  where cs = children t

-- | Follows a path in a tree.
follow :: Tree t => t -> Path -> t
follow t = fromError . followM t

fromError :: Either String a -> a
fromError = either error id

{-
indexIn :: (Eq p, Parent p) => p -> p -> Maybe Int
indexIn child = elemIndex child . children
-}


-- | Yields the depth of the tree.
depth :: Tree t => t -> Int
depth t
  | null depths = 1
  | otherwise   = 1 + maximum (map depth $ children t)
  where depths = map depth $ children t


-- | Yields all ancestors at the specified depth.
selectDepth :: Tree t => Int -> t -> [t]
selectDepth 0 t = [t]
selectDepth d t = concatMap (selectDepth (d - 1)) (children t)



------------------------------------
-- Tree selections
------------------------------------


-- | Selection in a tree. The path indicates the left side of the selection; the int tells how many siblings to the right are included in the selection.
type TreeSelection = (Path, Int)


-- | Selectable trees.
class Tree t => Selectable t where
  -- | Tells whether complete subranges of children may be selected in this tree. If not, valid TreeSelections in this tree always have a second element @0@.
  allowSubranges :: t -> Bool


-- | Enumerates all possible selections of a tree.
allSelections :: Selectable a => a -> [TreeSelection]
allSelections p = ([], 0) : subranges ++ recurse where
  subranges
    | allowSubranges p  = [ ([from], to - from)
                          | from <- [0 .. length cs - 2]
                          , to <- [from + 1 .. length cs - 1]
                          , from > 0 || to < length cs - 1
                          ]
    | otherwise         = []
  cs = children p
  recurse = concat $ zipWith label cs [0 ..]
  label c i = map (rt i) (allSelections c)
  rt i (path, offset) = (i : path, offset)

-- | Selects part of a tree.
select :: (Monad m, Tree t) => t -> TreeSelection -> m [t]
select t (path, offset) = (sequence . map (followM t) . take (offset + 1) . iterate right) path

-- | Computes the range of a valid selection.
selectionToRange :: (Tree a, KnowsPosition a) => a -> TreeSelection -> Range
selectionToRange parent (path, offset) = (from, to) where
  from = begin $ follow parent path
  to   = end   $ follow parent (sibling offset path)


-- | Converts a specified range to a corresponding selection and returns it in a monad.
rangeToSelection :: (Tree a, KnowsPosition a, Monad m) => a -> Range -> m TreeSelection
rangeToSelection p ran@(b, e)
  -- If the range matches that of the root, we're done.
  | range p == ran  = return ([], 0)

  | otherwise =
      -- Find the children whose ranges contain b and e.
      case ( findIndex (\c -> b `inRange` range c) cs
           , findIndex (\c -> e `inRange` range c) cs) of

               (Just l, Just r) ->
                   if l == r
                   -- b and e are contained by the same child!
                   -- Recurse into child.
                   then rangeToSelection (cs !! l) ran
                          -- ... and prepend child index, of course.
                          >>= (\(path, offset) -> return (l : path, offset))

                   else if begin (cs !! l) == b && end (cs !! r) == e
                   -- b is the beginning of l, and e is the end
                   -- of r: a selection of a range of children.
                   -- Note that r - l > 0; else it would've been
                   -- caught by the previous test.
                   -- This also means that there are many ways
                   -- to select a single node: either select it
                   -- directly, or select all its children.
                   then return ([l], r - l)

                   -- All other cases are bad.
                   else fail "text selection does not have corresponding tree selection"

               -- Either position is not contained
               -- within any child. Can't be valid.
               _ -> fail "text selection does not have corresponding tree selection"

      where cs = children p


-- | Returns the path to the deepest descendant whose range contains the specified position.
posToPath :: (Tree a, KnowsPosition a) => a -> Pos -> Path
posToPath p pos = case break (\c -> pos `inRange` range c) (children p) of
  (_, []) -> []
  (no, c:_) -> length no : posToPath c pos


-- | Tells whether the text selection corresponds to a tree selection.
isValidRange :: (KnowsPosition a, Selectable a) => a -> Range -> Bool
isValidRange p = isJust . rangeToSelection p


------------------------------------
-- Suggesting and fixing
------------------------------------


-- | Yields all possible selections, ordered by distance to the specified range, closest first.
suggest :: (Selectable a, KnowsPosition a) => a -> Range -> [TreeSelection]
suggest p r = sortBy distance $ allSelections p where
  distance s1 s2 = (selectionToRange p s1 `distRange` r) `compare` (selectionToRange p s2 `distRange` r)


-- | Takes @suggest@'s first suggestion and yields its range.
repair :: (KnowsPosition a, Selectable a) => a -> Range -> Range
repair p = selectionToRange p . head . suggest p