-----------------------------------------------------------------------------
--
-- Module      :  Data.EdgeTree
-- Copyright   :  (c) 2012-16 Brian W Bush
-- License     :  MIT
--
-- Maintainer  :  Brian W Bush <consult@brianwbush.info>
-- Stability   :  Stable
-- Portability :  Portable
--
-- | Trees with data on vertices and edges.
--
-----------------------------------------------------------------------------


{-# LANGUAGE Safe #-}


module Data.EdgeTree (
-- * Types
  EdgeTree(..)
, TreeEdge(..)
-- * Functions
, createEdgeTree
, filterEdgeTree
, filterEdgeTree'
, mapEdgeTree
, mapEdgeTree'
, mapTop
, truncateEdgeTree
-- * Input/output
, putEdgeForest
, hPutEdgeForest
, putEdgeTree
, putEdgeTree'
, hPutEdgeTree
, hPutEdgeTree'
, putTreeEdge'
, hPutTreeEdge'
) where


import System.IO (Handle, hPutStrLn, stdout)


-- | A tree with data at the vertex and edges radiating from it.
data EdgeTree e v =
  EdgeTree {
    vertex :: v
  , edges :: [TreeEdge e v]
  }
    deriving (Read, Show)

instance Functor (EdgeTree e) where
  fmap f (EdgeTree v ee) = EdgeTree (f v) (map (fmap f) ee)

instance Foldable (EdgeTree e) where
  foldMap f (EdgeTree v []) = f v
  foldMap f (EdgeTree v ee) = mconcat (f v : map (foldMap f) ee)

-- TODO: The edgle labels are foldable, too.

instance Traversable (EdgeTree e) where
  traverse f (EdgeTree v ee) =
        EdgeTree
    <$> f v
    <*> sequenceA (map (traverse f) ee)


-- | An edge with data and connecting to a tree.
data TreeEdge e v =
  TreeEdge {
    edge :: e
  , target :: EdgeTree e v
  }
    deriving (Read, Show)

instance Functor (TreeEdge e) where
  fmap f (TreeEdge e t) = TreeEdge e (fmap f t)

instance Foldable (TreeEdge e) where
  foldMap f (TreeEdge _ t) = foldMap f t

instance Traversable (TreeEdge e) where
  traverse f (TreeEdge e t) =
        TreeEdge e
    <$> traverse f t


-- | Create a tree.
createEdgeTree ::
     (a -> v)     -- ^ Function for labelling vertices.
  -> (a -> e)     -- ^ Function for labelling edges.
  -> (a -> [a])   -- ^ Function for generating objects radiating from the starting object.
  -> a            -- ^ The starting objects.
  -> EdgeTree e v -- ^ The tree.
createEdgeTree vertexLabeller edgeLabeller generator start =
  EdgeTree {
    vertex = vertexLabeller start
  , edges = map (createTreeEdge vertexLabeller edgeLabeller generator) $ generator start
  }


-- | Create an edge.
createTreeEdge ::
     (a -> v)     -- ^ Function for labelling vertices.
  -> (a -> e)     -- ^ Function for labelling edges.
  -> (a -> [a])   -- ^ Function for generating objects radiating from the starting object.
  -> a            -- ^ The starting objects.
  -> TreeEdge e v -- ^ The tree.
createTreeEdge vertexLabeller edgeLabeller generator start =
  TreeEdge {
    edge = edgeLabeller start
  , target = createEdgeTree vertexLabeller edgeLabeller generator start
  }


-- | Filter a tree.
filterEdgeTree ::
     ((v, e, v) -> Bool) -- ^ Function for filtering based on vertex-edge-vertex labelling.
  -> EdgeTree e v        -- ^ The tree.
  -> EdgeTree e v        -- ^ The filtered tree.
filterEdgeTree f edgeTree@(EdgeTree v ee) =
  edgeTree {
    edges = map filterTreeEdge $ filter filterVertex ee
  }
    where
      filterVertex (TreeEdge e (EdgeTree v' _)) = f (v, e, v')
      filterTreeEdge treeEdge@(TreeEdge _ t) = treeEdge {target = filterEdgeTree f t}


-- | Filter a tree.
filterEdgeTree' ::
     (v -> e -> Bool) -- ^ Function for filtering based on vextex-edge labelling.
  -> EdgeTree e v     -- ^ The tree.
  -> EdgeTree e v     -- ^ The filtered tree.
filterEdgeTree' = filterEdgeTree . uncurry2of3


-- | Evaluate a function on vertices of a tree.
mapEdgeTree ::
     ((v, e, v) -> w) -- ^ Function for evaluating vertex-edge-vertex triplets.
  -> w                -- ^ The new value for the root of the tree.
  -> EdgeTree e v     -- ^ The tree.
  -> EdgeTree e w     -- ^ The transformed tree.
mapEdgeTree f start (EdgeTree v ee) =
  EdgeTree {
    vertex = start
  , edges = map mapTreeEdge ee
  }
    where
      mapTreeEdge treeEdge@(TreeEdge e t@(EdgeTree v' _)) =
        treeEdge {
          target = mapEdgeTree f (f (v, e, v')) t
        }


-- | Evaluate a function on vertices of a tree.
mapEdgeTree' ::
     (v -> e -> w) -- ^ Function for evaluating vertex-edge-vertex triplets.
  -> w             -- ^ The new value for the root of the tree.
  -> EdgeTree e v  -- ^ The tree.
  -> EdgeTree e w  -- ^ The transformed tree.
mapEdgeTree' = mapEdgeTree . uncurry2of3


-- | Uncurry the first two elements of a triplet.
uncurry2of3 :: (a -> b -> d) -> (a, b, c) -> d
uncurry2of3 f (x, y, _) = f x y


-- | Apply a function to the first subtrees.
mapTop :: (EdgeTree e v -> a) -> EdgeTree e v -> [(e, a)]
mapTop f (EdgeTree _ edges') = map (\(TreeEdge edge' target') -> (edge', f target')) edges'


-- | Truncate a tree at a particular depth.
truncateEdgeTree :: Int -> EdgeTree e v -> EdgeTree e v
truncateEdgeTree 0 (EdgeTree vertex' _) = EdgeTree vertex' []
truncateEdgeTree n (EdgeTree vertex' edges') = EdgeTree vertex' $ map (truncateTreeEdge (n - 1)) edges'


-- | Truncate a tree at a particular depth.
truncateTreeEdge :: Int -> TreeEdge e v -> TreeEdge e v
truncateTreeEdge n (TreeEdge edge' target') = TreeEdge edge' $ truncateEdgeTree n target'


-- | Print a forest.
putEdgeForest ::
     Int                 -- ^ How many levels to print.
  -> (a -> String)       -- ^ Function for rendering the label for a tree.
  -> (v -> String)       -- ^ Function for rendering vertex labels.
  -> (e -> String)       -- ^ Function for rendering edge labels.
  -> [(a, EdgeTree e v)] -- ^ The forest.
  -> IO ()               -- ^ The action for printing the forest.
putEdgeForest = hPutEdgeForest stdout


-- | Print a forest.
hPutEdgeForest ::
     Handle              -- ^ Where to print the forest.
  -> Int                 -- ^ How many levels to print.
  -> (a -> String)       -- ^ Function for rendering the label for a tree.
  -> (v -> String)       -- ^ Function for rendering vertex labels.
  -> (e -> String)       -- ^ Function for rendering edge labels.
  -> [(a, EdgeTree e v)] -- ^ The forest.
  -> IO ()               -- ^ The action for printing the forest.
hPutEdgeForest handle depth showGroup showVertex showEdge =
  do
    let
      putGroup (group, edgeTree) =
        do
          hPutStrLn handle $ showGroup group
          hPutEdgeTree' handle depth ". " showVertex showEdge edgeTree
    mapM_ putGroup


-- | Print a tree.
putEdgeTree ::
     Int           -- ^ How many levels to print.
  -> (v -> String) -- ^ Function for rendering vertex labels.
  -> (e -> String) -- ^ Function for rendering edge labels.
  -> EdgeTree e v  -- ^ The tree.
  -> IO ()         -- ^ The action for printing the tree.
putEdgeTree = hPutEdgeTree stdout


-- | Print a tree.
putEdgeTree' ::
     Int           -- ^ How many levels to print.
  -> String        -- ^ The string for indentation, which will be prefixed to each line output.
  -> (v -> String) -- ^ Function for rendering vertex labels.
  -> (e -> String) -- ^ Function for rendering edge labels.
  -> EdgeTree e v  -- ^ The tree.
  -> IO ()         -- ^ The action for printing the tree.
putEdgeTree' = hPutEdgeTree' stdout


-- | Print a tree.
hPutEdgeTree ::
     Handle        -- ^ Where to print the tree.
  -> Int           -- ^ How many levels to print.
  -> (v -> String) -- ^ Function for rendering vertex labels.
  -> (e -> String) -- ^ Function for rendering edge labels.
  -> EdgeTree e v  -- ^ The tree.
  -> IO ()         -- ^ The action for printing the tree.
hPutEdgeTree = flip flip "" . hPutEdgeTree'


-- | Print a tree.
hPutEdgeTree' ::
     Handle        -- ^ Where to print the tree.
  -> Int           -- ^ How many levels to print.
  -> String        -- ^ The string for indentation, which will be prefixed to each line output.
  -> (v -> String) -- ^ Function for rendering vertex labels.
  -> (e -> String) -- ^ Function for rendering edge labels.
  -> EdgeTree e v  -- ^ The tree.
  -> IO ()         -- ^ The action for printing the tree.
hPutEdgeTree' handle depth indent showVertex showEdge (EdgeTree v ee) =
  do
    hPutStrLn handle $ indent ++ showVertex v
    mapM_ (hPutTreeEdge' handle depth (indent ++ ". ") showVertex showEdge) ee


-- | Print an edge.
putTreeEdge' ::
     Int           -- ^ How many levels to print.
  -> String        -- ^ The string for indentation, which will be prefixed to each line output.
  -> (v -> String) -- ^ Function for rendering vertex labels.
  -> (e -> String) -- ^ Function for rendering edge labels.
  -> TreeEdge e v  -- ^ The edge.
  -> IO ()         -- ^ THe action for printing the edge.
putTreeEdge' = hPutTreeEdge' stdout


-- | Print an edge.
hPutTreeEdge' ::
     Handle        -- ^ Where to print the tree.
  -> Int           -- ^ How many levels to print.
  -> String        -- ^ The string for indentation, which will be prefixed to each line output.
  -> (v -> String) -- ^ Function for rendering vertex labels.
  -> (e -> String) -- ^ Function for rendering edge labels.
  -> TreeEdge e v  -- ^ The edge.
  -> IO ()         -- ^ THe action for printing the edge.
hPutTreeEdge' handle depth indent showVertex showEdge (TreeEdge e t) =
  do
    hPutStrLn handle $ indent ++ showEdge e
    if depth > 0
      then hPutEdgeTree' handle (depth - 1) (indent ++ ". ") showVertex showEdge t
      else hPutStrLn handle $ indent ++ ". <<TRUNCATED>>"