{-# LANGUAGE CPP                       #-}
{-# LANGUAGE DeriveFoldable            #-}
{-# LANGUAGE DeriveFunctor             #-}
{-# LANGUAGE DeriveTraversable         #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TemplateHaskell           #-}
{-# LANGUAGE TypeFamilies              #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Diagrams.TwoD.Layout.Tree
-- Copyright   :  (c) 2011 Brent Yorgey
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  byorgey@cis.upenn.edu
--
-- A collection of methods for laying out various kinds of trees.
-- This module is still experimental, and more layout methods will
-- probably be added over time.
--
-- Laying out a rose tree using a symmetric layout:
--
-- > import Data.Tree
-- > import Diagrams.TwoD.Layout.Tree
-- >
-- > t1 = Node 'A' [Node 'B' (map lf "CDE"), Node 'F' [Node 'G' (map lf "HIJKLM"), Node 'N' (map lf "OPQR")]]
-- >   where lf x = Node x []
-- >
-- > exampleSymmTree =
-- >   renderTree ((<> circle 1 # fc white) . text . (:[]))
-- >              (~~)
-- >              (symmLayout' (with & slHSep .~ 4 & slVSep .~ 4) t1)
-- >   # centerXY # pad 1.1
--
-- <<diagrams/src_Diagrams_TwoD_Layout_Tree_exampleSymmTree.svg#diagram=exampleSymmTree&width=300>>
--
-- Laying out a rose tree of diagrams, with spacing automatically
-- adjusted for the size of the diagrams:
--
-- > import Data.Tree
-- > import Data.Maybe (fromMaybe)
-- > import Diagrams.TwoD.Layout.Tree
-- >
-- > tD = Node (rect 1 3)
-- >        [ Node (circle 0.2) []
-- >        , Node (hcat . replicate 3 $ circle 1) []
-- >        , Node (eqTriangle 5) []
-- >        ]
-- >
-- > exampleSymmTreeWithDs =
-- >   renderTree id (~~)
-- >   (symmLayout' (with & slWidth  .~ fromMaybe (0,0) . extentX
-- >                      & slHeight .~ fromMaybe (0,0) . extentY)
-- >      tD)
-- >   # centerXY # pad 1.1
--
-- <<diagrams/src_Diagrams_TwoD_Layout_Tree_exampleSymmTreeWithDs.svg#diagram=exampleSymmTreeWithDs&width=300>>
--
-- Using a variant symmetric layout algorithm specifically for binary trees:
--
-- > import Diagrams.TwoD.Layout.Tree
-- > import Diagrams.Prelude hiding (Empty)
-- >
-- > drawT = maybe mempty (renderTree (const (circle 0.05 # fc black)) (~~))
-- >       . symmLayoutBin' (with & slVSep .~ 0.5)
-- >
-- > tree500 = drawT t # centerXY # pad 1.1
-- >   where t = genTree 500 0.05
-- >         -- genTree 500 0.05 randomly generates trees of size 500 +/- 5%,
-- >         -- definition not shown
--
-- <<diagrams/src_Diagrams_TwoD_Layout_Tree_tree500.svg#diagram=tree500&width=400>>
--
-- Using force-based layout on a binary tree:
--
-- > {-# LANGUAGE NoMonomorphismRestriction #-}
-- > import Diagrams.Prelude hiding (Empty)
-- > import Diagrams.TwoD.Layout.Tree
-- >
-- > gent 0 = Empty
-- > gent n = BNode n (gent (n-1)) (gent (n-1))
-- >
-- > Just t' = uniqueXLayout 1 1 (gent 4)
-- >
-- > fblEx = renderTree (\n -> (text (show n) # fontSizeL 0.5
-- >                             <> circle 0.3 # fc white))
-- >             (~~)
-- >             (forceLayoutTree t')
-- >         # centerXY # pad 1.1
--
-- <<diagrams/src_Diagrams_TwoD_Layout_Tree_fblEx.svg#diagram=fblEx&width=300>>
--
-- Using a radial layout:
--
-- > import Diagrams.Prelude
-- > import Diagrams.TwoD.Layout.Tree
-- > import Data.Tree
-- >
-- > t = Node 'A' [Node 'B' (map lf "CDE"), Node 'F' [Node 'G' (map lf "HIJKLM"), Node 'N' (map lf "OPQRS")], Node 'T' (map lf "UVWXYZ")]
-- >   where lf x = Node x []
-- >
-- > radialEx =
-- >    renderTree (\n -> (text (show n) # fontSizeG 0.5
-- >                             <> circle 0.5 # fc white))
-- >              (~~) (radialLayout t)
-- >    # centerXY # pad 1.1
--
-- <<diagrams/src_Diagrams_TwoD_Layout_Tree_radialEx.svg#diagram=radialEx&width=300>>

module Diagrams.TwoD.Layout.Tree
       ( -- * Binary trees
         -- $BTree

         BTree(..)
       , leaf

         -- * Layout algorithms

         -- ** Unique-x layout

       , uniqueXLayout

         -- ** Radial layout

       , radialLayout

         -- ** Symmetric layout

         -- $symmetric
       , symmLayout
       , symmLayout'
       , symmLayoutBin
       , symmLayoutBin'
       , SymmLayoutOpts(..), slHSep, slVSep, slWidth, slHeight

         -- ** Force-directed layout
         -- $forcedirected

       , forceLayoutTree
       , forceLayoutTree'
       , ForceLayoutTreeOpts(..), forceLayoutOpts, edgeLen, springK, staticK

       , treeToEnsemble
       , label
       , reconstruct

         -- * Rendering

       , renderTree
       , renderTree'

       ) where

import           Physics.ForceLayout

import           Control.Arrow       (first, second, (&&&), (***))
import           Control.Monad.State

import           Data.Default
import qualified Data.Foldable       as F
import           Data.Function       (on)
import           Data.List           (mapAccumL)
import qualified Data.Map            as M
import           Data.Maybe
import qualified Data.Traversable    as T
import           Data.Tree

import           Control.Lens        (makeLenses, view, (+=), (-=), (^.))
import           Diagrams
import           Linear              ((*^))
import           Linear.Affine

#if !MIN_VERSION_base(4,11,0)
import           Data.Semigroup
#endif

#if __GLASGOW_HASKELL__ < 710
import           Control.Applicative
#endif

------------------------------------------------------------
--  Binary trees
------------------------------------------------------------

-- $BTree
-- There is a standard type of rose trees ('Tree') defined in the
-- @containers@ package, but there is no standard type for binary
-- trees, so we define one here.  Note, if you want to draw binary
-- trees with data of type @a@ at the leaves, you can use something
-- like @BTree (Maybe a)@ with @Nothing@ at internal nodes;
-- 'renderTree' lets you specify how to draw each node.

-- | Binary trees with data at internal nodes.
data BTree a = Empty | BNode a (BTree a) (BTree a)
  deriving (BTree a -> BTree a -> Bool
forall a. Eq a => BTree a -> BTree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: BTree a -> BTree a -> Bool
$c/= :: forall a. Eq a => BTree a -> BTree a -> Bool
== :: BTree a -> BTree a -> Bool
$c== :: forall a. Eq a => BTree a -> BTree a -> Bool
Eq, BTree a -> BTree a -> Bool
BTree a -> BTree a -> Ordering
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {a}. Ord a => Eq (BTree a)
forall a. Ord a => BTree a -> BTree a -> Bool
forall a. Ord a => BTree a -> BTree a -> Ordering
forall a. Ord a => BTree a -> BTree a -> BTree a
min :: BTree a -> BTree a -> BTree a
$cmin :: forall a. Ord a => BTree a -> BTree a -> BTree a
max :: BTree a -> BTree a -> BTree a
$cmax :: forall a. Ord a => BTree a -> BTree a -> BTree a
>= :: BTree a -> BTree a -> Bool
$c>= :: forall a. Ord a => BTree a -> BTree a -> Bool
> :: BTree a -> BTree a -> Bool
$c> :: forall a. Ord a => BTree a -> BTree a -> Bool
<= :: BTree a -> BTree a -> Bool
$c<= :: forall a. Ord a => BTree a -> BTree a -> Bool
< :: BTree a -> BTree a -> Bool
$c< :: forall a. Ord a => BTree a -> BTree a -> Bool
compare :: BTree a -> BTree a -> Ordering
$ccompare :: forall a. Ord a => BTree a -> BTree a -> Ordering
Ord, ReadPrec [BTree a]
ReadPrec (BTree a)
ReadS [BTree a]
forall a. Read a => ReadPrec [BTree a]
forall a. Read a => ReadPrec (BTree a)
forall a. Read a => Int -> ReadS (BTree a)
forall a. Read a => ReadS [BTree a]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [BTree a]
$creadListPrec :: forall a. Read a => ReadPrec [BTree a]
readPrec :: ReadPrec (BTree a)
$creadPrec :: forall a. Read a => ReadPrec (BTree a)
readList :: ReadS [BTree a]
$creadList :: forall a. Read a => ReadS [BTree a]
readsPrec :: Int -> ReadS (BTree a)
$creadsPrec :: forall a. Read a => Int -> ReadS (BTree a)
Read, Int -> BTree a -> ShowS
forall a. Show a => Int -> BTree a -> ShowS
forall a. Show a => [BTree a] -> ShowS
forall a. Show a => BTree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [BTree a] -> ShowS
$cshowList :: forall a. Show a => [BTree a] -> ShowS
show :: BTree a -> String
$cshow :: forall a. Show a => BTree a -> String
showsPrec :: Int -> BTree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> BTree a -> ShowS
Show, forall a b. a -> BTree b -> BTree a
forall a b. (a -> b) -> BTree a -> BTree b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> BTree b -> BTree a
$c<$ :: forall a b. a -> BTree b -> BTree a
fmap :: forall a b. (a -> b) -> BTree a -> BTree b
$cfmap :: forall a b. (a -> b) -> BTree a -> BTree b
Functor, forall a. Eq a => a -> BTree a -> Bool
forall a. Num a => BTree a -> a
forall a. Ord a => BTree a -> a
forall m. Monoid m => BTree m -> m
forall a. BTree a -> Bool
forall a. BTree a -> Int
forall a. BTree a -> [a]
forall a. (a -> a -> a) -> BTree a -> a
forall m a. Monoid m => (a -> m) -> BTree a -> m
forall b a. (b -> a -> b) -> b -> BTree a -> b
forall a b. (a -> b -> b) -> b -> BTree a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
product :: forall a. Num a => BTree a -> a
$cproduct :: forall a. Num a => BTree a -> a
sum :: forall a. Num a => BTree a -> a
$csum :: forall a. Num a => BTree a -> a
minimum :: forall a. Ord a => BTree a -> a
$cminimum :: forall a. Ord a => BTree a -> a
maximum :: forall a. Ord a => BTree a -> a
$cmaximum :: forall a. Ord a => BTree a -> a
elem :: forall a. Eq a => a -> BTree a -> Bool
$celem :: forall a. Eq a => a -> BTree a -> Bool
length :: forall a. BTree a -> Int
$clength :: forall a. BTree a -> Int
null :: forall a. BTree a -> Bool
$cnull :: forall a. BTree a -> Bool
toList :: forall a. BTree a -> [a]
$ctoList :: forall a. BTree a -> [a]
foldl1 :: forall a. (a -> a -> a) -> BTree a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> BTree a -> a
foldr1 :: forall a. (a -> a -> a) -> BTree a -> a
$cfoldr1 :: forall a. (a -> a -> a) -> BTree a -> a
foldl' :: forall b a. (b -> a -> b) -> b -> BTree a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> BTree a -> b
foldl :: forall b a. (b -> a -> b) -> b -> BTree a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> BTree a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> BTree a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> BTree a -> b
foldr :: forall a b. (a -> b -> b) -> b -> BTree a -> b
$cfoldr :: forall a b. (a -> b -> b) -> b -> BTree a -> b
foldMap' :: forall m a. Monoid m => (a -> m) -> BTree a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> BTree a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> BTree a -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> BTree a -> m
fold :: forall m. Monoid m => BTree m -> m
$cfold :: forall m. Monoid m => BTree m -> m
F.Foldable, Functor BTree
Foldable BTree
forall (t :: * -> *).
Functor t
-> Foldable t
-> (forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => BTree (m a) -> m (BTree a)
forall (f :: * -> *) a. Applicative f => BTree (f a) -> f (BTree a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> BTree a -> m (BTree b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> BTree a -> f (BTree b)
sequence :: forall (m :: * -> *) a. Monad m => BTree (m a) -> m (BTree a)
$csequence :: forall (m :: * -> *) a. Monad m => BTree (m a) -> m (BTree a)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> BTree a -> m (BTree b)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> BTree a -> m (BTree b)
sequenceA :: forall (f :: * -> *) a. Applicative f => BTree (f a) -> f (BTree a)
$csequenceA :: forall (f :: * -> *) a. Applicative f => BTree (f a) -> f (BTree a)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> BTree a -> f (BTree b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> BTree a -> f (BTree b)
T.Traversable)

-- | Convenient constructor for leaves.
leaf :: a -> BTree a
leaf :: forall a. a -> BTree a
leaf a
a = forall a. a -> BTree a -> BTree a -> BTree a
BNode a
a forall a. BTree a
Empty forall a. BTree a
Empty

------------------------------------------------------------
--  Layout algorithms
------------------------------------------------------------

--------------------------------------------------
-- Unique X layout for binary trees.  No
-- two nodes share the same X coordinate.

data Pos = Pos { Pos -> Int
_level :: Int
               , Pos -> Int
_horiz :: Int
               }
  deriving (Pos -> Pos -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Pos -> Pos -> Bool
$c/= :: Pos -> Pos -> Bool
== :: Pos -> Pos -> Bool
$c== :: Pos -> Pos -> Bool
Eq, Int -> Pos -> ShowS
[Pos] -> ShowS
Pos -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Pos] -> ShowS
$cshowList :: [Pos] -> ShowS
show :: Pos -> String
$cshow :: Pos -> String
showsPrec :: Int -> Pos -> ShowS
$cshowsPrec :: Int -> Pos -> ShowS
Show)

makeLenses ''Pos

pos2Point :: Num n => n -> n -> Pos -> P2 n
pos2Point :: forall n. Num n => n -> n -> Pos -> P2 n
pos2Point n
cSep n
lSep (Pos Int
l Int
h) = forall n. (n, n) -> P2 n
p2 (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
h forall a. Num a => a -> a -> a
* n
cSep, -forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
l forall a. Num a => a -> a -> a
* n
lSep)

-- | @uniqueXLayout xSep ySep t@ lays out the binary tree @t@ using a
--   simple recursive algorithm with the following properties:
--
--   * Every left subtree is completely to the left of its parent, and
--     similarly for right subtrees.
--
--   * All the nodes at a given depth in the tree have the same
--     y-coordinate. The separation distance between levels is given by
--     @ySep@.
--
--   * Every node has a unique x-coordinate. The separation between
--     successive nodes from left to right is given by @xSep@.

uniqueXLayout :: Num n => n -> n -> BTree a -> Maybe (Tree (a, P2 n))
uniqueXLayout :: forall n a. Num n => n -> n -> BTree a -> Maybe (Tree (a, P2 n))
uniqueXLayout n
cSep n
lSep BTree a
t = (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second) (forall n. Num n => n -> n -> Pos -> P2 n
pos2Point n
cSep n
lSep)
                forall a b. (a -> b) -> a -> b
$ forall s a. State s a -> s -> a
evalState (forall {m :: * -> *} {a}.
MonadState Pos m =>
BTree a -> m (Maybe (Tree (a, Pos)))
uniqueXLayout' BTree a
t) (Int -> Int -> Pos
Pos Int
0 Int
0)
  where uniqueXLayout' :: BTree a -> m (Maybe (Tree (a, Pos)))
uniqueXLayout' BTree a
Empty         = forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
        uniqueXLayout' (BNode a
a BTree a
l BTree a
r) = do
          forall {m :: * -> *}. MonadState Pos m => m ()
down
          Maybe (Tree (a, Pos))
l' <- BTree a -> m (Maybe (Tree (a, Pos)))
uniqueXLayout' BTree a
l
          forall {m :: * -> *}. MonadState Pos m => m ()
up
          Pos
p  <- forall {f :: * -> *}. MonadState Pos f => f Pos
mkNode
          forall {m :: * -> *}. MonadState Pos m => m ()
down
          Maybe (Tree (a, Pos))
r' <- BTree a -> m (Maybe (Tree (a, Pos)))
uniqueXLayout' BTree a
r
          forall {m :: * -> *}. MonadState Pos m => m ()
up
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (forall a. a -> [Tree a] -> Tree a
Node (a
a,Pos
p) (forall a. [Maybe a] -> [a]
catMaybes [Maybe (Tree (a, Pos))
l', Maybe (Tree (a, Pos))
r']))
        mkNode :: f Pos
mkNode = forall s (m :: * -> *). MonadState s m => m s
get forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* (Lens' Pos Int
horiz forall s (m :: * -> *) a.
(MonadState s m, Num a) =>
ASetter' s a -> a -> m ()
+= Int
1)

        down :: m ()
down = Lens' Pos Int
level forall s (m :: * -> *) a.
(MonadState s m, Num a) =>
ASetter' s a -> a -> m ()
+= Int
1
        up :: m ()
up   = Lens' Pos Int
level forall s (m :: * -> *) a.
(MonadState s m, Num a) =>
ASetter' s a -> a -> m ()
-= Int
1

--------------------------------------------------
-- "Symmetric" layout of rose trees.

-- $symmetric
-- \"Symmetric\" layout of rose trees, based on the algorithm described in:
--
-- Andrew J. Kennedy. /Drawing Trees/, J Func. Prog. 6 (3): 527-534,
-- May 1996.
--
-- Trees laid out using this algorithm satisfy:
--
--   1. Nodes at a given level are always separated by at least a
--   given minimum distance.
--
--   2. Parent nodes are centered with respect to their immediate
--   offspring (though /not/ necessarily with respect to the entire
--   subtrees under them).
--
--   3. Layout commutes with mirroring: that is, the layout of a given
--   tree is the mirror image of the layout of the tree's mirror
--   image.  Put another way, there is no inherent left or right bias.
--
--   4. Identical subtrees are always rendered identically.  Put
--   another way, the layout of any subtree is independent of the rest
--   of the tree.
--
--   5. The layouts are as narrow as possible while satisfying all the
--   above constraints.

-- | A tree with /relative/ positioning information.  The @n@
--   at each node is the horizontal /offset/ from its parent.
type Rel t n a = t (a, n)

-- | Shift a RelTree horizontally.
moveTree :: Num n => n -> Rel Tree n a -> Rel Tree n a
moveTree :: forall n a. Num n => n -> Rel Tree n a -> Rel Tree n a
moveTree n
x' (Node (a
a, n
x) [Tree (a, n)]
ts) = forall a. a -> [Tree a] -> Tree a
Node (a
a, n
xforall a. Num a => a -> a -> a
+n
x') [Tree (a, n)]
ts

-- | An /extent/ is a list of pairs, recording the leftmost and
--   rightmost (absolute) horizontal positions of a tree at each
--   depth.
newtype Extent n = Extent { forall n. Extent n -> [(n, n)]
getExtent :: [(n, n)] }

extent :: ([(n, n)] -> [(n, n)]) -> Extent n -> Extent n
extent :: forall n. ([(n, n)] -> [(n, n)]) -> Extent n -> Extent n
extent [(n, n)] -> [(n, n)]
f = forall n. [(n, n)] -> Extent n
Extent forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(n, n)] -> [(n, n)]
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall n. Extent n -> [(n, n)]
getExtent

consExtent :: (n, n) -> Extent n -> Extent n
consExtent :: forall n. (n, n) -> Extent n -> Extent n
consExtent = forall n. ([(n, n)] -> [(n, n)]) -> Extent n -> Extent n
extent forall b c a. (b -> c) -> (a -> b) -> a -> c
. (:)

-- | Shift an extent horizontally.
moveExtent :: Num n => n -> Extent n -> Extent n
moveExtent :: forall n. Num n => n -> Extent n -> Extent n
moveExtent n
x = (forall n. ([(n, n)] -> [(n, n)]) -> Extent n -> Extent n
extent forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map) ((forall a. Num a => a -> a -> a
+n
x) forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** (forall a. Num a => a -> a -> a
+n
x))

-- | Reflect an extent about the vertical axis.
flipExtent :: Num n => Extent n -> Extent n
flipExtent :: forall n. Num n => Extent n -> Extent n
flipExtent = (forall n. ([(n, n)] -> [(n, n)]) -> Extent n -> Extent n
extent forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map) (\(n
p,n
q) -> (-n
q, -n
p))

-- | Merge two non-overlapping extents.
mergeExtents :: Extent n -> Extent n -> Extent n
mergeExtents :: forall n. Extent n -> Extent n -> Extent n
mergeExtents (Extent [(n, n)]
e1) (Extent [(n, n)]
e2) = forall n. [(n, n)] -> Extent n
Extent forall a b. (a -> b) -> a -> b
$ forall {a} {b}. [(a, b)] -> [(a, b)] -> [(a, b)]
mergeExtents' [(n, n)]
e1 [(n, n)]
e2
  where

    mergeExtents' :: [(a, b)] -> [(a, b)] -> [(a, b)]
mergeExtents' [] [(a, b)]
qs = [(a, b)]
qs
    mergeExtents' [(a, b)]
ps [] = [(a, b)]
ps
    mergeExtents' ((a
p,b
_) : [(a, b)]
ps) ((a
_,b
q) : [(a, b)]
qs) = (a
p,b
q) forall a. a -> [a] -> [a]
: [(a, b)] -> [(a, b)] -> [(a, b)]
mergeExtents' [(a, b)]
ps [(a, b)]
qs

instance Semigroup (Extent n) where
  <> :: Extent n -> Extent n -> Extent n
(<>) = forall n. Extent n -> Extent n -> Extent n
mergeExtents

instance Monoid (Extent n) where
  mempty :: Extent n
mempty  = forall n. [(n, n)] -> Extent n
Extent []
  mappend :: Extent n -> Extent n -> Extent n
mappend = forall a. Semigroup a => a -> a -> a
(<>)

-- | Determine the amount to shift in order to \"fit\" two extents
--   next to one another.  The first argument is the separation to
--   leave between them.
fit :: (Num n, Ord n) => n -> Extent n -> Extent n -> n
fit :: forall n. (Num n, Ord n) => n -> Extent n -> Extent n -> n
fit n
hSep (Extent [(n, n)]
ps) (Extent [(n, n)]
qs) = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (n
0 forall a. a -> [a] -> [a]
: forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\(n
_,n
p) (n
q,n
_) -> n
p forall a. Num a => a -> a -> a
- n
q forall a. Num a => a -> a -> a
+ n
hSep) [(n, n)]
ps [(n, n)]
qs)

-- | Fit a list of subtree extents together using a left-biased
--   algorithm.  Compute a list of positions (relative to the leftmost
--   subtree which is considered to have position 0).
fitListL :: (Num n, Ord n) => n -> [Extent n] -> [n]
fitListL :: forall n. (Num n, Ord n) => n -> [Extent n] -> [n]
fitListL n
hSep = forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
mapAccumL Extent n -> Extent n -> (Extent n, n)
fitOne forall a. Monoid a => a
mempty
  where
    fitOne :: Extent n -> Extent n -> (Extent n, n)
fitOne Extent n
acc Extent n
e =
      let x :: n
x = forall n. (Num n, Ord n) => n -> Extent n -> Extent n -> n
fit n
hSep Extent n
acc Extent n
e
      in  (Extent n
acc forall a. Semigroup a => a -> a -> a
<> forall n. Num n => n -> Extent n -> Extent n
moveExtent n
x Extent n
e, n
x)

-- | Fit a list of subtree extents together with a right bias.
fitListR :: (Num n, Ord n) => n -> [Extent n] -> [n]
fitListR :: forall n. (Num n, Ord n) => n -> [Extent n] -> [n]
fitListR n
hSep = forall a. [a] -> [a]
reverse forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. Num a => a -> a
negate forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall n. (Num n, Ord n) => n -> [Extent n] -> [n]
fitListL n
hSep forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall n. Num n => Extent n -> Extent n
flipExtent forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [a]
reverse

-- | Compute a symmetric fitting by averaging the results of left- and
--   right-biased fitting.
fitList :: (Fractional n, Ord n) => n -> [Extent n] -> [n]
fitList :: forall n. (Fractional n, Ord n) => n -> [Extent n] -> [n]
fitList n
hSep = forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {a}. Fractional a => a -> a -> a
mean) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall n. (Num n, Ord n) => n -> [Extent n] -> [n]
fitListL n
hSep forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& forall n. (Num n, Ord n) => n -> [Extent n] -> [n]
fitListR n
hSep)
  where mean :: a -> a -> a
mean a
x a
y = (a
xforall a. Num a => a -> a -> a
+a
y)forall {a}. Fractional a => a -> a -> a
/a
2

-- | Options for controlling the symmetric tree layout algorithm.
data SymmLayoutOpts n a =
  SLOpts { forall n a. SymmLayoutOpts n a -> n
_slHSep   :: n -- ^ Minimum horizontal
                                         --   separation between sibling
                                         --   nodes.  The default is 1.
         , forall n a. SymmLayoutOpts n a -> n
_slVSep   :: n -- ^ Vertical separation
                                         --   between adjacent levels of
                                         --   the tree.  The default is 1.
         , forall n a. SymmLayoutOpts n a -> a -> (n, n)
_slWidth  :: a -> (n, n)
           -- ^ A function for measuring the horizontal extent (a pair
           --   of x-coordinates) of an item in the tree.  The default
           --   is @const (0,0)@, that is, the nodes are considered as
           --   taking up no space, so the centers of the nodes will
           --   be separated according to the @slHSep@ and @slVSep@.
           --   However, this can be useful, /e.g./ if you have a tree
           --   of diagrams of irregular size and want to make sure no
           --   diagrams overlap.  In that case you could use
           --   @fromMaybe (0,0) . extentX@.
         , forall n a. SymmLayoutOpts n a -> a -> (n, n)
_slHeight :: a -> (n, n)
           -- ^ A function for measuring the vertical extent of an
           --   item in the tree.  The default is @const (0,0)@.  See
           --   the documentation for 'slWidth' for more information.
         }

makeLenses ''SymmLayoutOpts

instance Num n => Default (SymmLayoutOpts n a) where
  def :: SymmLayoutOpts n a
def = SLOpts
          { _slHSep :: n
_slHSep   = n
1
          , _slVSep :: n
_slVSep   = n
1
          , _slWidth :: a -> (n, n)
_slWidth  = forall a b. a -> b -> a
const (n
0,n
0)
          , _slHeight :: a -> (n, n)
_slHeight = forall a b. a -> b -> a
const (n
0,n
0)
          }

-- | Actual recursive tree layout algorithm, which returns a tree
--   layout as well as an extent.
symmLayoutR :: (Fractional n, Ord n) => SymmLayoutOpts n a -> Tree a -> (Rel Tree n a, Extent n)
symmLayoutR :: forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> Tree a -> (Rel Tree n a, Extent n)
symmLayoutR SymmLayoutOpts n a
opts (Node a
a [Tree a]
ts) = (Tree (a, n)
rt, Extent n
ext)
  where ([Tree (a, n)]
trees, [Extent n]
extents) = forall a b. [(a, b)] -> ([a], [b])
unzip (forall a b. (a -> b) -> [a] -> [b]
map (forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> Tree a -> (Rel Tree n a, Extent n)
symmLayoutR SymmLayoutOpts n a
opts) [Tree a]
ts)
        positions :: [n]
positions        = forall n. (Fractional n, Ord n) => n -> [Extent n] -> [n]
fitList (SymmLayoutOpts n a
opts forall s a. s -> Getting a s a -> a
^. forall n a. Lens' (SymmLayoutOpts n a) n
slHSep) [Extent n]
extents
        pTrees :: [Tree (a, n)]
pTrees           = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall n a. Num n => n -> Rel Tree n a -> Rel Tree n a
moveTree [n]
positions [Tree (a, n)]
trees
        pExtents :: [Extent n]
pExtents         = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall n. Num n => n -> Extent n -> Extent n
moveExtent [n]
positions [Extent n]
extents
        ext :: Extent n
ext              = (SymmLayoutOpts n a
optsforall s a. s -> Getting a s a -> a
^.forall n a. Lens' (SymmLayoutOpts n a) (a -> (n, n))
slWidth) a
a forall n. (n, n) -> Extent n -> Extent n
`consExtent` forall a. Monoid a => [a] -> a
mconcat [Extent n]
pExtents
        rt :: Tree (a, n)
rt               = forall a. a -> [Tree a] -> Tree a
Node (a
a, n
0) [Tree (a, n)]
pTrees

-- | Symmetric tree layout algorithm specialized to binary trees.
--   Returns a tree layout as well as an extent.
symmLayoutBinR :: (Fractional n, Ord n) =>
                  SymmLayoutOpts n a -> BTree a -> (Maybe (Rel Tree n a), Extent n)
symmLayoutBinR :: forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> BTree a -> (Maybe (Rel Tree n a), Extent n)
symmLayoutBinR SymmLayoutOpts n a
_    BTree a
Empty         = (forall a. Maybe a
Nothing, forall a. Monoid a => a
mempty)
symmLayoutBinR SymmLayoutOpts n a
opts (BNode a
a BTree a
l BTree a
r) = (forall a. a -> Maybe a
Just Tree (a, n)
rt, Extent n
ext)
  where (Maybe (Tree (a, n))
l', Extent n
extL) = forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> BTree a -> (Maybe (Rel Tree n a), Extent n)
symmLayoutBinR SymmLayoutOpts n a
opts BTree a
l
        (Maybe (Tree (a, n))
r', Extent n
extR) = forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> BTree a -> (Maybe (Rel Tree n a), Extent n)
symmLayoutBinR SymmLayoutOpts n a
opts BTree a
r
        positions :: [n]
positions  = case (Maybe (Tree (a, n))
l', Maybe (Tree (a, n))
r') of
                       (Maybe (Tree (a, n))
Nothing, Maybe (Tree (a, n))
_) -> [n
0, SymmLayoutOpts n a
opts forall s a. s -> Getting a s a -> a
^. forall n a. Lens' (SymmLayoutOpts n a) n
slHSep forall {a}. Fractional a => a -> a -> a
/ n
2]
                       (Maybe (Tree (a, n))
_, Maybe (Tree (a, n))
Nothing) -> [-(SymmLayoutOpts n a
opts forall s a. s -> Getting a s a -> a
^. forall n a. Lens' (SymmLayoutOpts n a) n
slHSep) forall {a}. Fractional a => a -> a -> a
/ n
2, n
0]
                       (Maybe (Tree (a, n)), Maybe (Tree (a, n)))
_          -> forall n. (Fractional n, Ord n) => n -> [Extent n] -> [n]
fitList (SymmLayoutOpts n a
opts forall s a. s -> Getting a s a -> a
^. forall n a. Lens' (SymmLayoutOpts n a) n
slHSep) [Extent n
extL, Extent n
extR]
        pTrees :: [Tree (a, n)]
pTrees   = forall a. [Maybe a] -> [a]
catMaybes forall a b. (a -> b) -> a -> b
$ forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall n a. Num n => n -> Rel Tree n a -> Rel Tree n a
moveTree) [n]
positions [Maybe (Tree (a, n))
l',Maybe (Tree (a, n))
r']
        pExtents :: [Extent n]
pExtents = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall n. Num n => n -> Extent n -> Extent n
moveExtent [n]
positions [Extent n
extL, Extent n
extR]
        ext :: Extent n
ext = (SymmLayoutOpts n a
optsforall s a. s -> Getting a s a -> a
^.forall n a. Lens' (SymmLayoutOpts n a) (a -> (n, n))
slWidth) a
a forall n. (n, n) -> Extent n -> Extent n
`consExtent` forall a. Monoid a => [a] -> a
mconcat [Extent n]
pExtents
        rt :: Tree (a, n)
rt  = forall a. a -> [Tree a] -> Tree a
Node (a
a, n
0) [Tree (a, n)]
pTrees

-- | Run the symmetric rose tree layout algorithm on a given tree,
--   resulting in the same tree annotated with node positions.
symmLayout' :: (Fractional n, Ord n) => SymmLayoutOpts n a -> Tree a -> Tree (a, P2 n)
symmLayout' :: forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> Tree a -> Tree (a, P2 n)
symmLayout' SymmLayoutOpts n a
opts = forall n a.
(Num n, Ord n) =>
SymmLayoutOpts n a -> P2 n -> Rel Tree n a -> Tree (a, P2 n)
unRelativize SymmLayoutOpts n a
opts forall (f :: * -> *) a. (Additive f, Num a) => Point f a
origin forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> Tree a -> (Rel Tree n a, Extent n)
symmLayoutR SymmLayoutOpts n a
opts

-- | Run the symmetric rose tree layout algorithm on a given tree
--   using default options, resulting in the same tree annotated with
--   node positions.
symmLayout :: (Fractional n, Ord n) => Tree a -> Tree (a, P2 n)
symmLayout :: forall n a. (Fractional n, Ord n) => Tree a -> Tree (a, P2 n)
symmLayout = forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> Tree a -> Tree (a, P2 n)
symmLayout' forall a. Default a => a
def

-- | Lay out a binary tree using a slight variant of the symmetric
--   layout algorithm.  In particular, if a node has only a left child
--   but no right child (or vice versa), the child will be offset from
--   the parent horizontally by half the horizontal separation
--   parameter. Note that the result will be @Nothing@ if and only if
--   the input tree is @Empty@.
symmLayoutBin' :: (Fractional n, Ord n) => SymmLayoutOpts n a -> BTree a -> Maybe (Tree (a,P2 n))
symmLayoutBin' :: forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> BTree a -> Maybe (Tree (a, P2 n))
symmLayoutBin' SymmLayoutOpts n a
opts = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall n a.
(Num n, Ord n) =>
SymmLayoutOpts n a -> P2 n -> Rel Tree n a -> Tree (a, P2 n)
unRelativize SymmLayoutOpts n a
opts forall (f :: * -> *) a. (Additive f, Num a) => Point f a
origin) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> BTree a -> (Maybe (Rel Tree n a), Extent n)
symmLayoutBinR SymmLayoutOpts n a
opts

-- | Lay out a binary tree using a slight variant of the symmetric
--   layout algorithm, using default options.  In particular, if a
--   node has only a left child but no right child (or vice versa),
--   the child will be offset from the parent horizontally by half the
--   horizontal separation parameter. Note that the result will be
--   @Nothing@ if and only if the input tree is @Empty@.
symmLayoutBin :: (Fractional n, Ord n) => BTree a -> Maybe (Tree (a,P2 n))
symmLayoutBin :: forall n a.
(Fractional n, Ord n) =>
BTree a -> Maybe (Tree (a, P2 n))
symmLayoutBin = forall n a.
(Fractional n, Ord n) =>
SymmLayoutOpts n a -> BTree a -> Maybe (Tree (a, P2 n))
symmLayoutBin' forall a. Default a => a
def

-- | Given a fixed location for the root, turn a tree with
--   \"relative\" positioning into one with absolute locations
--   associated to all the nodes.
unRelativize :: (Num n, Ord n) =>
                SymmLayoutOpts n a -> P2 n -> Rel Tree n a -> Tree (a, P2 n)
unRelativize :: forall n a.
(Num n, Ord n) =>
SymmLayoutOpts n a -> P2 n -> Rel Tree n a -> Tree (a, P2 n)
unRelativize SymmLayoutOpts n a
opts P2 n
curPt (Node (a
a,n
hOffs) [Tree (a, n)]
ts)
    = forall a. a -> [Tree a] -> Tree a
Node (a
a, P2 n
rootPt) (forall a b. (a -> b) -> [a] -> [b]
map (forall n a.
(Num n, Ord n) =>
SymmLayoutOpts n a -> P2 n -> Rel Tree n a -> Tree (a, P2 n)
unRelativize SymmLayoutOpts n a
opts (P2 n
rootPt forall (p :: * -> *) a. (Affine p, Num a) => p a -> Diff p a -> p a
.+^ (n
vOffs forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ forall (v :: * -> *) n. (R2 v, Additive v, Num n) => v n
unit_Y))) [Tree (a, n)]
ts)
  where rootPt :: P2 n
rootPt = P2 n
curPt forall (p :: * -> *) a. (Affine p, Num a) => p a -> Diff p a -> p a
.+^ (n
hOffs forall (f :: * -> *) a. (Functor f, Num a) => a -> f a -> f a
*^ forall (v :: * -> *) n. (R1 v, Additive v, Num n) => v n
unitX)
        vOffs :: n
vOffs  = - forall a b. (a, b) -> a
fst ((SymmLayoutOpts n a
optsforall s a. s -> Getting a s a -> a
^.forall n a. Lens' (SymmLayoutOpts n a) (a -> (n, n))
slHeight) a
a)
               forall a. Num a => a -> a -> a
+ (forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. (SymmLayoutOpts n a
optsforall s a. s -> Getting a s a -> a
^.forall n a. Lens' (SymmLayoutOpts n a) (a -> (n, n))
slHeight) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> a
rootLabel) forall a b. (a -> b) -> a -> b
$ [Tree (a, n)]
ts)
               forall a. Num a => a -> a -> a
+ (SymmLayoutOpts n a
opts forall s a. s -> Getting a s a -> a
^. forall n a. Lens' (SymmLayoutOpts n a) n
slVSep)

--------------------------------------------------
--  Force-directed layout of rose trees

-- $forcedirected
-- Force-directed layout of rose trees.

data ForceLayoutTreeOpts n =
  FLTOpts
  { forall n. ForceLayoutTreeOpts n -> ForceLayoutOpts n
_forceLayoutOpts :: ForceLayoutOpts n -- ^ Options to the force layout simulator, including damping.
  , forall n. ForceLayoutTreeOpts n -> n
_edgeLen         :: n -- ^ How long edges should be, ideally.
                                           --   This will be the resting length for
                                           --   the springs.
  , forall n. ForceLayoutTreeOpts n -> n
_springK         :: n -- ^ Spring constant.  The
                                           --   bigger the constant,
                                           --   the more the edges
                                           --   push/pull towards their
                                           --   resting length.
  , forall n. ForceLayoutTreeOpts n -> n
_staticK         :: n -- ^ Coulomb constant.  The
                                           --   bigger the constant, the
                                           --   more sibling nodes repel
                                           --   each other.
  }

makeLenses ''ForceLayoutTreeOpts

instance Floating n => Default (ForceLayoutTreeOpts n) where
  def :: ForceLayoutTreeOpts n
def = FLTOpts
    { _forceLayoutOpts :: ForceLayoutOpts n
_forceLayoutOpts = forall a. Default a => a
def
    , _edgeLen :: n
_edgeLen = forall a. Floating a => a -> a
sqrt n
2
    , _springK :: n
_springK = n
0.05
    , _staticK :: n
_staticK = n
0.1
    }

-- | Assign unique ID numbers to the nodes of a tree, and generate an
--   'Ensemble' suitable for simulating in order to do force-directed
--   layout of the tree.  In particular,
--
--   * edges are modeled as springs
--
--   * nodes are modeled as point charges
--
--   * nodes are constrained to keep the same y-coordinate.
--
--   The input to @treeToEnsemble@ could be a tree already laid out by
--   some other method, such as 'uniqueXLayout'.
treeToEnsemble :: forall a n. Floating n => ForceLayoutTreeOpts n
               -> Tree (a, P2 n) -> (Tree (a, PID), Ensemble V2 n)
treeToEnsemble :: forall a n.
Floating n =>
ForceLayoutTreeOpts n
-> Tree (a, P2 n) -> (Tree (a, Int), Ensemble V2 n)
treeToEnsemble ForceLayoutTreeOpts n
opts Tree (a, P2 n)
t =
  ( forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first forall a b. (a, b) -> a
fst) Tree ((a, P2 n), Int)
lt
  , forall (v :: * -> *) n.
[([Edge], Point v n -> Point v n -> v n)]
-> Map Int (Particle v n) -> Ensemble v n
Ensemble
      [ ([Edge]
edges, \P2 n
pt1 P2 n
pt2 -> forall (v :: * -> *) a.
(Metric v, Fractional a) =>
v a -> v a -> v a
project forall (v :: * -> *) n. (R1 v, Additive v, Num n) => v n
unitX (forall (v :: * -> *) n.
(Metric v, Floating n) =>
n -> n -> Point v n -> Point v n -> v n
hookeForce (ForceLayoutTreeOpts n
opts forall s a. s -> Getting a s a -> a
^. forall n. Lens' (ForceLayoutTreeOpts n) n
springK) (ForceLayoutTreeOpts n
opts forall s a. s -> Getting a s a -> a
^. forall n. Lens' (ForceLayoutTreeOpts n) n
edgeLen) P2 n
pt1 P2 n
pt2))
      , ([Edge]
sibs,  \P2 n
pt1 P2 n
pt2 -> forall (v :: * -> *) a.
(Metric v, Fractional a) =>
v a -> v a -> v a
project forall (v :: * -> *) n. (R1 v, Additive v, Num n) => v n
unitX (forall (v :: * -> *) n.
(Metric v, Floating n) =>
n -> Point v n -> Point v n -> v n
coulombForce (ForceLayoutTreeOpts n
opts forall s a. s -> Getting a s a -> a
^. forall n. Lens' (ForceLayoutTreeOpts n) n
staticK) P2 n
pt1 P2 n
pt2))
      ]
      Map Int (Particle V2 n)
particleMap
  )

  where lt :: Tree ((a,P2 n), PID)
        lt :: Tree ((a, P2 n), Int)
lt = forall (t :: * -> *) a. Traversable t => t a -> t (a, Int)
label Tree (a, P2 n)
t

        particleMap :: M.Map PID (Particle V2 n)
        particleMap :: Map Int (Particle V2 n)
particleMap = forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
                    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second forall (v :: * -> *) n.
(Additive v, Num n) =>
Point v n -> Particle v n
initParticle)
                    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList
                    forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall {b} {a}. (b, a) -> (a, b)
swap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first forall a b. (a, b) -> b
snd)
                    forall a b. (a -> b) -> a -> b
$ Tree ((a, P2 n), Int)
lt
        swap :: (b, a) -> (a, b)
swap (b
x,a
y) = (a
y,b
x)

        edges, sibs :: [Edge]
        edges :: [Edge]
edges       = Tree Int -> [Edge]
extractEdges (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd Tree ((a, P2 n), Int)
lt)
        sibs :: [Edge]
sibs        = Forest Int -> [Edge]
extractSibs [forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd Tree ((a, P2 n), Int)
lt]

        extractEdges :: Tree PID -> [Edge]
        extractEdges :: Tree Int -> [Edge]
extractEdges (Node Int
i Forest Int
cs) = forall a b. (a -> b) -> [a] -> [b]
map (((,) Int
i) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> a
rootLabel) Forest Int
cs
                                    forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree Int -> [Edge]
extractEdges Forest Int
cs

        extractSibs :: Forest PID -> [Edge]
        extractSibs :: Forest Int -> [Edge]
extractSibs [] = []
        extractSibs Forest Int
ts = (\[Int]
is -> forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
is (forall a. [a] -> [a]
tail [Int]
is)) (forall a b. (a -> b) -> [a] -> [b]
map forall a. Tree a -> a
rootLabel Forest Int
ts)
                      forall a. [a] -> [a] -> [a]
++ Forest Int -> [Edge]
extractSibs (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. Tree a -> [Tree a]
subForest Forest Int
ts)

--        sz = ala Sum foldMap . fmap (const 1) $ t
--        sibs = [(x,y) | x <- [0..sz-2], y <- [x+1 .. sz-1]]

-- | Assign unique IDs to every node in a tree (or other traversable structure).
label :: (T.Traversable t) => t a -> t (a, PID)
label :: forall (t :: * -> *) a. Traversable t => t a -> t (a, Int)
label = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall s a. State s a -> s -> a
evalState Int
0 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
T.mapM (\a
a -> forall s (m :: * -> *). MonadState s m => m s
get forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Int
i -> forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (forall a. Num a => a -> a -> a
+Int
1) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return (a
a,Int
i))

-- | Reconstruct a tree (or any traversable structure) from an
--   'Ensemble', given unique identifier annotations matching the
--   identifiers used in the 'Ensemble'.
reconstruct :: (Functor t, Num n) => Ensemble V2 n -> t (a, PID) -> t (a, P2 n)
reconstruct :: forall (t :: * -> *) n a.
(Functor t, Num n) =>
Ensemble V2 n -> t (a, Int) -> t (a, P2 n)
reconstruct Ensemble V2 n
e = (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second)
                  (forall a. a -> Maybe a -> a
fromMaybe forall (f :: * -> *) a. (Additive f, Num a) => Point f a
origin forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view forall (v :: * -> *) n. Lens' (Particle v n) (Point v n)
pos) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b c. (a -> b -> c) -> b -> a -> c
flip forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup (Ensemble V2 n
eforall s a. s -> Getting a s a -> a
^.forall (v :: * -> *) n.
Lens' (Ensemble v n) (Map Int (Particle v n))
particles))

-- | Force-directed layout of rose trees, with default parameters (for
--   more options, see 'forceLayoutTree'').  In particular,
--
--   * edges are modeled as springs
--
--   * nodes are modeled as point charges
--
--   * nodes are constrained to keep the same y-coordinate.
--
--   The input could be a tree already laid out by some other method,
--   such as 'uniqueXLayout'.
forceLayoutTree :: (Floating n, Ord n) => Tree (a, P2 n) -> Tree (a, P2 n)
forceLayoutTree :: forall n a. (Floating n, Ord n) => Tree (a, P2 n) -> Tree (a, P2 n)
forceLayoutTree = forall n a.
(Floating n, Ord n) =>
ForceLayoutTreeOpts n -> Tree (a, P2 n) -> Tree (a, P2 n)
forceLayoutTree' forall a. Default a => a
def

-- | Force-directed layout of rose trees, with configurable parameters.
forceLayoutTree' :: (Floating n, Ord n) =>
                    ForceLayoutTreeOpts n -> Tree (a, P2 n) -> Tree (a, P2 n)
forceLayoutTree' :: forall n a.
(Floating n, Ord n) =>
ForceLayoutTreeOpts n -> Tree (a, P2 n) -> Tree (a, P2 n)
forceLayoutTree' ForceLayoutTreeOpts n
opts Tree (a, P2 n)
t = forall (t :: * -> *) n a.
(Functor t, Num n) =>
Ensemble V2 n -> t (a, Int) -> t (a, P2 n)
reconstruct (forall (v :: * -> *) n.
(Metric v, Num n, Ord n) =>
ForceLayoutOpts n -> Ensemble v n -> Ensemble v n
forceLayout (ForceLayoutTreeOpts n
optsforall s a. s -> Getting a s a -> a
^.forall n. Lens' (ForceLayoutTreeOpts n) (ForceLayoutOpts n)
forceLayoutOpts) Ensemble V2 n
e) Tree (a, Int)
ti
  where (Tree (a, Int)
ti, Ensemble V2 n
e) = forall a n.
Floating n =>
ForceLayoutTreeOpts n
-> Tree (a, P2 n) -> (Tree (a, Int), Ensemble V2 n)
treeToEnsemble ForceLayoutTreeOpts n
opts Tree (a, P2 n)
t

-- | Radial layout of rose trees, adapted from Andy Pavlo,
--   "Interactive, Tree-Based Graph Visualization", p. 18
--   (<http://www.cs.cmu.edu/~pavlo/static/papers/APavloThesis032006.pdf>)
radialLayout :: Tree a -> Tree (a, P2 Double)
radialLayout :: forall a. Tree a -> Tree (a, P2 Double)
radialLayout t :: Tree a
t@(Node a
a [Tree a]
_)
  = forall a. a -> [Tree a] -> Tree a
Node (a
a, forall (f :: * -> *) a. (Additive f, Num a) => Point f a
origin) (forall a.
Double
-> Double
-> Double
-> Int
-> Double
-> [Tree (a, NodeInfo)]
-> [Tree (a, P2 Double)]
assignPos Double
0 forall a. Floating a => a
pi Double
0 (NodeInfo -> Int
nodeLeaves NodeInfo
info) (forall a. Tree a -> Double
weight Tree a
t) [Tree (a, NodeInfo)]
ts)
  where
    Node (a
_,NodeInfo
info) [Tree (a, NodeInfo)]
ts = forall a. Tree a -> Tree (a, NodeInfo)
decorate Tree a
t

-- | Implementation of radial layout: @assignPos alpha beta theta k w ts@
--
--   * @alpha@, @beta@ define the bounds of an annular wedge around the root
--   * @k@ is #leaves of root and lambda is #leaves of vertex
--   * @theta@ is ?
--   * @w@ is ?
--
--   The algorithm used is an extension of Algorithm 1, Page 18 of
--   <http://www.cs.cmu.edu/~pavlo/static/papers/APavloThesis032006.pdf>.
--   See
--   <https://drive.google.com/file/d/0B3el1oMKFsOIVGVRYzJzWGwzWDA/view>
--   for more examples.
assignPos :: Double -> Double -> Double -> Int -> Double  -> [Tree (a, NodeInfo)] -> [Tree (a, P2 Double)]
assignPos :: forall a.
Double
-> Double
-> Double
-> Int
-> Double
-> [Tree (a, NodeInfo)]
-> [Tree (a, P2 Double)]
assignPos Double
_ Double
_ Double
_ Int
_ Double
_ [] = []
assignPos Double
alpha Double
beta Double
theta Int
k Double
w (Node (a
a, NodeInfo
info) [Tree (a, NodeInfo)]
ts1 : [Tree (a, NodeInfo)]
ts2)
  = forall a. a -> [Tree a] -> Tree a
Node (a
a, P2 Double
pt) (forall a.
Double
-> Double
-> Double
-> Int
-> Double
-> [Tree (a, NodeInfo)]
-> [Tree (a, P2 Double)]
assignPos Double
theta Double
u Double
theta Int
lambda Double
w [Tree (a, NodeInfo)]
ts1) forall a. a -> [a] -> [a]
: forall a.
Double
-> Double
-> Double
-> Int
-> Double
-> [Tree (a, NodeInfo)]
-> [Tree (a, P2 Double)]
assignPos Double
alpha Double
beta Double
u Int
k Double
w [Tree (a, NodeInfo)]
ts2
    where
      lambda :: Int
lambda  = NodeInfo -> Int
nodeLeaves NodeInfo
info
      u :: Double
u       = Double
theta forall a. Num a => a -> a -> a
+ (Double
beta forall a. Num a => a -> a -> a
- Double
alpha) forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
lambda forall {a}. Fractional a => a -> a -> a
/ forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k
      pt :: P2 Double
pt      = (Double
1 forall c. Coordinates c => PrevDim c -> FinalCoord c -> c
^& Double
0)
              # rotate (theta + u @@ rad)
              # scale (w * fromIntegral (nodeDepth info) / 2)

-- | Compute the length of radius determined by the number of children to avoid
--   node overlapping
weight :: Tree a -> Double
weight :: forall a. Tree a -> Double
weight Tree a
t = forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum forall a b. (a -> b) -> a -> b
$
               forall a b. (a -> b) -> [a] -> [b]
map (((\ Int
x -> forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x forall {a}. Fractional a => a -> a -> a
/ Double
2) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Int
length) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. Tree a -> a
rootLabel)
                    (forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Bool -> Bool
not forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Bool
null) forall a b. (a -> b) -> a -> b
$ forall a. (a -> a) -> a -> [a]
iterate (forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. Tree a -> [Tree a]
subForest) [Tree a
t])

data NodeInfo = NodeInfo
  { NodeInfo -> Int
nodeLeaves :: Int
  , NodeInfo -> Int
nodeDepth  :: Int
  }

decorate :: Tree a -> Tree (a, NodeInfo)
decorate :: forall a. Tree a -> Tree (a, NodeInfo)
decorate = forall a. Int -> Tree a -> Tree (a, NodeInfo)
decorate' Int
0

decorate' :: Int -> Tree a -> Tree (a, NodeInfo)
decorate' :: forall a. Int -> Tree a -> Tree (a, NodeInfo)
decorate' Int
d (Node a
a [Tree a]
ts) = forall a. a -> [Tree a] -> Tree a
Node (a
a, NodeInfo
info) [Tree (a, NodeInfo)]
ts'
  where
    ts' :: [Tree (a, NodeInfo)]
ts'   = forall a b. (a -> b) -> [a] -> [b]
map (forall a. Int -> Tree a -> Tree (a, NodeInfo)
decorate' (Int
dforall a. Num a => a -> a -> a
+Int
1)) [Tree a]
ts
    infos :: [NodeInfo]
infos = forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> a
rootLabel) [Tree (a, NodeInfo)]
ts'
    leaves :: Int
leaves
      | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Tree a]
ts   = Int
1
      | Bool
otherwise = forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map NodeInfo -> Int
nodeLeaves forall a b. (a -> b) -> a -> b
$ [NodeInfo]
infos

    info :: NodeInfo
info  = Int -> Int -> NodeInfo
NodeInfo Int
leaves Int
d

------------------------------------------------------------
--  Rendering
------------------------------------------------------------

-- | Draw a tree annotated with node positions, given functions
--   specifying how to draw nodes and edges.
renderTree :: (Monoid' m, Floating n, Ord n)
           => (a -> QDiagram b V2 n m) -> (P2 n -> P2 n -> QDiagram b V2 n m)
           -> Tree (a, P2 n) -> QDiagram b V2 n m
renderTree :: forall m n a b.
(Monoid' m, Floating n, Ord n) =>
(a -> QDiagram b V2 n m)
-> (P2 n -> P2 n -> QDiagram b V2 n m)
-> Tree (a, P2 n)
-> QDiagram b V2 n m
renderTree a -> QDiagram b V2 n m
n P2 n -> P2 n -> QDiagram b V2 n m
e = forall m n a b.
(Monoid' m, Floating n, Ord n) =>
(a -> QDiagram b V2 n m)
-> ((a, P2 n) -> (a, P2 n) -> QDiagram b V2 n m)
-> Tree (a, P2 n)
-> QDiagram b V2 n m
renderTree' a -> QDiagram b V2 n m
n (P2 n -> P2 n -> QDiagram b V2 n m
e forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` forall a b. (a, b) -> b
snd)

-- | Draw a tree annotated with node positions, given functions
--   specifying how to draw nodes and edges.  Unlike 'renderTree',
--   this version gives the edge-drawing function access to the actual
--   values stored at the nodes rather than just their positions.
renderTree' :: (Monoid' m, Floating n, Ord n)
           => (a -> QDiagram b V2 n m) -> ((a,P2 n) -> (a,P2 n) -> QDiagram b V2 n m)
           -> Tree (a, P2 n) -> QDiagram b V2 n m
renderTree' :: forall m n a b.
(Monoid' m, Floating n, Ord n) =>
(a -> QDiagram b V2 n m)
-> ((a, P2 n) -> (a, P2 n) -> QDiagram b V2 n m)
-> Tree (a, P2 n)
-> QDiagram b V2 n m
renderTree' a -> QDiagram b V2 n m
renderNode (a, P2 n) -> (a, P2 n) -> QDiagram b V2 n m
renderEdge = forall n a.
(InSpace V2 n a, Fractional n, Alignable a, HasOrigin a) =>
a -> a
alignT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (v :: * -> *) n a.
(InSpace v n a, R1 v, Fractional n, Alignable a, HasOrigin a) =>
a -> a
centerX forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree (a, P2 n) -> QDiagram b V2 n m
renderTreeR
  where
    renderTreeR :: Tree (a, P2 n) -> QDiagram b V2 n m
renderTreeR (Node (a
a,P2 n
p) [Tree (a, P2 n)]
cs) =
         a -> QDiagram b V2 n m
renderNode a
a forall a b. a -> (a -> b) -> b
# forall (v :: * -> *) n t.
(InSpace v n t, HasOrigin t) =>
Point v n -> t -> t
moveTo P2 n
p
      forall a. Semigroup a => a -> a -> a
<> forall a. Monoid a => [a] -> a
mconcat (forall a b. (a -> b) -> [a] -> [b]
map Tree (a, P2 n) -> QDiagram b V2 n m
renderTreeR [Tree (a, P2 n)]
cs)
      forall a. Semigroup a => a -> a -> a
<> forall a. Monoid a => [a] -> a
mconcat (forall a b. (a -> b) -> [a] -> [b]
map ((a, P2 n) -> (a, P2 n) -> QDiagram b V2 n m
renderEdge (a
a,P2 n
p) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Tree a -> a
rootLabel) [Tree (a, P2 n)]
cs)


-- > -- Critical size-limited Boltzmann generator for binary trees (used in example)
-- >
-- > import           Control.Applicative
-- > import           Control.Lens                   hiding (( # ), Empty)
-- > import           Control.Monad.Random
-- > import           Control.Monad.Reader
-- > import           Control.Monad.State
-- > import           Control.Monad.Trans.Maybe
-- >
-- > genTreeCrit :: ReaderT Int (StateT Int (MaybeT (Rand StdGen))) (BTree ())
-- > genTreeCrit = do
-- >   r <- getRandom
-- >   if r <= (1/2 :: Double)
-- >     then return Empty
-- >     else atom >> (BNode () <$> genTreeCrit <*> genTreeCrit)
-- >
-- > atom :: ReaderT Int (StateT Int (MaybeT (Rand StdGen))) ()
-- > atom = do
-- >   targetSize <- ask
-- >   curSize <- get
-- >   when (curSize >= targetSize) mzero
-- >   put (curSize + 1)
-- >
-- > genOneTree :: Int -> Int -> Double -> Maybe (BTree ())
-- > genOneTree seed size eps =
-- >   case mt of
-- >     Nothing -> Nothing
-- >     Just (t,sz) -> if sz >= minSz then Just t else Nothing
-- >
-- >   where
-- >     g          = mkStdGen seed
-- >     sizeWiggle = floor $ fromIntegral size * eps
-- >     maxSz = size + sizeWiggle
-- >     minSz = size - sizeWiggle
-- >     mt = (evalRand ?? g) . runMaybeT . (runStateT ?? 0) . (runReaderT ?? maxSz)
-- >        $ genTreeCrit
-- >
-- > genTree' :: Int -> Int -> Double -> BTree ()
-- > genTree' seed size eps =
-- >   case (genOneTree seed size eps) of
-- >     Nothing -> genTree' (seed+1) size eps
-- >     Just t  -> t
-- >
-- > genTree :: Int -> Double -> BTree ()
-- > genTree = genTree' 0