{-# LANGUAGE MultiParamTypeClasses #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Graphics.Rendering.Diagrams.Types
-- Copyright   :  (c) Brent Yorgey 2008
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  byorgey@gmail.com
-- Stability   :  experimental
-- Portability :  portable
--
-- Layout definitions for "Graphics.Rendering.Diagrams", an embedded
-- domain-specific language (EDSL) for creating simple diagrams.
--
-----------------------------------------------------------------------------

module Graphics.Rendering.Diagrams.Layouts
  ( -- * Union

    (##)
  , union, unionA

    -- * Lists

  , (<>), (//)
  , hcat, vcat
  , hcatA, vcatA
  , hsep, vsep
  , hsepA, vsepA
  , hdistrib, vdistrib
  , hdistribA, vdistribA

  , position, positionA
  , positionAlong, positionAlongA

  , grid, gridA, gridAs

  , VAlignment
  , top, vcenter, bottom
  , HAlignment
  , left, hcenter, right

    -- * Tree

  , tree

    -- * Miscellaneous

  , pad, padA
  , showBBox, showBBoxes
  , withSize

  ) where

import Graphics.Rendering.Diagrams.Types
import Graphics.Rendering.Diagrams.Attributes
import Graphics.Rendering.Diagrams.Shapes
import Graphics.Rendering.Diagrams.Paths
import Graphics.Rendering.Diagrams.Engine (sizeAndPos)
import Control.Arrow (first)
import Control.Monad.Cont

import Data.List (transpose)
import Data.Tree

-- Union -----------------------------------------------

-- | The union layout, which lays out diagrams superimposed on one
--   another.  The diagrams can be aligned both vertically and
--   horizontally.  A union layout with centered alignment in both
--   axes is simply the identity layout which does no repositioning of
--   subdiagrams.
data UnionLayout = UnionL HAlignment VAlignment
instance LayoutClass UnionLayout [] where
  layoutSizeAndPos _ [] = ((0,0), [])
  layoutSizeAndPos (UnionL halign valign) pds = ((w,h), positionedDs)
    where (sizes, diagrams) = unzip pds
          (xs, ys)          = unzip sizes
          (w,h)             = (maximum xs, maximum ys)
          xs'               = aligns halign w xs
          ys'               = aligns valign h ys
          positionedDs      = zipWith3 translate xs' ys' diagrams

-- | Create a 'Diagram' as a union of subdiagrams which will not be
--   repositioned.  If the subdiagrams overlap, they will appear with
--   the first 'Diagram' on the bottom, and the last on top.
union :: [Diagram] -> Diagram
union = unionA hcenter vcenter

-- | Superimpose one diagram atop another.  @d1 \#\# d2@ results in a
--   diagram in which @d2@ is on top of @d1@ (i.e., @d1@ is drawn
--   first, then @d2@).
(##) :: Diagram -> Diagram -> Diagram
d1 ## d2 = union [d1, d2]

-- | Create a 'Diagram' as a union of subdiagrams superimposed on one
--   another, aligned vertically and/or horizontally.
unionA :: HAlignment -> VAlignment -> [Diagram] -> Diagram
unionA ha va = Compound . Layout (UnionL ha va)

-- Lists -----------------------------------------------

-- | The possible ways to arrange a list of diagrams.
data ListType = H    -- ^ in a horizontal row, left to right
              | V    -- ^ in a vertical column, top to bottom
  deriving (Eq, Show, Read)

-- | A list of 'Diagram's can be aligned in one of several ways.
data Alignment = TopLeft      -- ^ Align along top/left edges, i.e. in
                              --   the negative direction
               | Center       -- ^ Align centers
               | BottomRight  -- ^ Align bottom/right edges, i.e. in
                              --   the positive direction
  deriving (Eq, Show, Read)

-- | Vertical alignment.
type VAlignment = Alignment

-- | Horizontal alignment.
type HAlignment = Alignment

top, vcenter, bottom :: VAlignment
left, hcenter, right :: HAlignment
top = TopLeft
bottom = BottomRight
left = TopLeft
right = BottomRight
vcenter = Center
hcenter = Center

-- | Compute new coordinates for the centers of the subdiagrams of a
--   list, based on the alignment.
aligns :: Alignment    -- ^ the alignment to use
       -> Double       -- ^ the total height (resp. width) we're working with
       -> [Double]     -- ^ the individual heights (resp. widths) of
                       --   the diagrams to align
       -> [Double]     -- ^ the amount to translate each diagram in
                       --   order to be properly aligned
aligns a h hs = map (align a h) hs

-- | Given a requested alignment, the total width (resp. height) of
--   the enclosing bounding box, and the width (resp. height) of a
--   diagram, compute the offset needed to properly align the diagram.
align :: Alignment    -- ^ the alignment to use
      -> Double       -- ^ the total height (resp. width) we're working with
      -> Double       -- ^ the height (resp. widths) of the diagram to
                      --   align
      -> Double       -- ^ the offset needed to align the diagram
align a h x = alignOffset a (h - x)

-- | Compute an offset corresponding to a given alignment and width.
alignOffset :: Alignment -> Double -> Double
alignOffset TopLeft x     = -x/2
alignOffset Center  _     = 0
alignOffset BottomRight x = x/2


-- | A list of 'Diagram's can be distributed in one of several ways.
data Distribution = Sep Double  -- ^ Put a constant separation between
                                --   each pair of diagrams
                  | Distrib Alignment Double
                                -- ^ @Distrib align sep@ represents
                                --   evenly spaced diagrams,
                                --   distributed with the @align@ of a
                                --   diagram placed every @sep@ units.
                                --   For example, @Distrib TopLeft 50@
                                --   means that the top/left of a
                                --   diagram will occur every 50
                                --   units.
  deriving (Eq, Show, Read)

-- | A horizontal or vertical list of diagrams, with configurable
--   alignment and distribution.
data List = List ListType Alignment Distribution
instance LayoutClass List [] where
  layoutSizeAndPos _ [] = ((0,0), [])
  layoutSizeAndPos (List typ algn dist) dss =

        -- unzip sizes and diagrams, flipping sizes for a vertical layout
    let (sizes, diagrams) = first (map (mswap typ)) $ unzip dss

        -- calculate the total size, given the method of distribution
        size  = listSize  dist sizes

        -- calculate new positions for the subdiagrams, flipping
        --   positions back for a vertical layout
        pos   = map (mswap typ) $ listPosns algn dist size sizes

        -- apply appropriate translates, and flip the final size for a
        --   vertical layout
    in  (mswap typ size, zipWith (uncurry translate) pos diagrams)

-- | Conditionally swap points, so vertical lists can be laid out
--   using the same code as horizontal lists, with appropriate swaps
--   before and after layout processing.
mswap :: ListType -> Point -> Point
mswap H (x,y) = (x,y)
mswap V (x,y) = (y,x)

-- | Given a method of distribution and the sizes of the subdiagrams
--   for a horizontal list, compute the total size of the entire list
listSize :: Distribution -> [Point] -> Point
listSize (Sep sep) ss = (x,y)
  where (xs,ys) = unzip ss
        x = sum xs + (sep * fromIntegral (length ss - 1))
        y = maximum ys
listSize (Distrib _ sep) ss = (x,y)
  where x = sep * fromIntegral (length ss)
        y = maximum (map snd ss)

-- | Calculate the final positions for the centers of the subdiagrams
--   of a horizontal list, based on the alignment and distribution.
listPosns :: Alignment      -- ^ vertical alignment of list elements
          -> Distribution   -- ^ horizontal distribution
          -> Point          -- ^ pre-calculated total size of the list
          -> [Point]        -- ^ sizes of individual elements
          -> [Point]        -- ^ new coordinates for the center of
                            --   each list element
listPosns a (Sep sep) (w,h) sizes = zip xs ys
  where (ws,hs)    = unzip sizes
        offsets    = scanl (+) 0 ws
        offsWSeps  = zipWith (+) offsets (zipWith (*) [0..] (repeat sep))
        centers    = zipWith (+) (init offsWSeps) (map (/2) ws)
        xs         = map (subtract (w/2)) centers
        ys         = aligns a h hs

listPosns a (Distrib da sep) (w,h) sizes = zip xs ys
  where (ws,hs)    = unzip sizes
        offsets    = zipWith (*) [0..] (repeat sep)
        centers    = zipWith (+) offsets (map (centerOffs da) ws)
        centerOffs TopLeft width = width/2
        centerOffs Center  _     = sep/2
        centerOffs BottomRight width = sep - width/2
        xs         = map (subtract (w/2)) centers
        ys         = aligns a h hs

-- | Lay out a list of 'Diagram's horizontally from left to right,
--   aligned along their top edges.
hcat :: [Diagram] -> Diagram
hcat = hsep 0

-- | @d1 \<\> d2@ is a 'Diagram' with @d1@ to the left of @d2@, aligned
--   along their top edges.
(<>) :: Diagram -> Diagram -> Diagram
d1 <> d2 = hcat [d1, d2]

-- | Lay out a list of 'Diagram's horizontally from left to right,
--   with the given vertical alignment ('top', 'vcenter', or 'bottom').
hcatA :: VAlignment -> [Diagram] -> Diagram
hcatA = hsepA 0

-- | Lay out a list of 'Diagram's horizontally, aligned along their
--   top edges, with a given amount of separation in between each pair.
hsep :: Double        -- ^ amount of separation between each pair of diagrams
     -> [Diagram] -> Diagram
hsep sep = hsepA sep top

-- | Lay out a list of 'Diagram's horizontally, with the given
--   amount of separation in between each pair, using the given
--   vertical alignment ('top', 'vcenter', or 'bottom').
hsepA :: Double       -- ^ amount of separation between each pair of diagrams
      -> VAlignment   -- ^ alignment to use ('top', 'vcenter', or 'bottom')
      -> [Diagram] -> Diagram
hsepA = sepA H

-- | Distribute a list of 'Diagram's horizontally according to a
--   regular spacing, aligned along their top edges.
hdistrib :: Double      -- ^ How far from one diagram to the next?
         -> HAlignment  -- ^ Distribute according to which parts of
                        --   the diagrams ('left', 'hcenter', 'right')?
         -> [Diagram] -> Diagram
hdistrib sep alignD = hdistribA sep alignD top

-- | Distribute a list of 'Diagram's horizontally according to a
--   regular spacing, with the given alignment.
hdistribA :: Double     -- ^ How far from one diagram to the next?
          -> HAlignment -- ^ Distribute according to which parts of
                        --   the diagrams ('left', 'hcenter', 'right')?
          -> VAlignment -- ^ alignment to use ('top', 'vcenter', 'bottom')
          -> [Diagram] -> Diagram
hdistribA = distribA H

-- | Lay out a list of 'Diagram's vertically from top to bottom,
--   aligned along their left edges.
vcat :: [Diagram] -> Diagram
vcat = vsep 0

-- | @d1 \/\/ d2@ is a 'Diagram' with @d1@ above @d2@, aligned
--   along their left edges.
(//) :: Diagram -> Diagram -> Diagram
d1 // d2 = vcat [d1, d2]

-- | Lay out a list of 'Diagram's vertically from top to bottom,
--   with the given horizontal alignment ('left', 'hcenter', or 'right').
vcatA :: HAlignment -> [Diagram] -> Diagram
vcatA = vsepA 0

-- | Lay out a list of 'Diagram's vertically, aligned along their
--   left edges, with a given amount of separation in between each pair.
vsep :: Double       -- ^ amount of separation between each pair of diagrams
     -> [Diagram] -> Diagram
vsep sep = vsepA sep left

-- | Lay out a list of 'Diagram's vertically, with the given
--   amount of separation in between each pair, using the given
--   horizontal alignment ('left', 'hcenter', or 'right').

vsepA :: Double       -- ^ amount of separation between each pair of diagrams
      -> HAlignment   -- ^ alignment to use ('left', 'hcenter', or 'right')
      -> [Diagram] -> Diagram
vsepA = sepA V

-- | Distribute a list of 'Diagram's vertically according to a regular
--   spacing, aligned along their left edges.
vdistrib :: Double     -- ^ How far from one diagram to the next?
         -> VAlignment -- ^ Distribute according to which parts of
                       --   the diagrams ('top', 'vcenter', 'bottom')?
         -> [Diagram] -> Diagram
vdistrib sep alignD = vdistribA sep alignD left

-- | Distribute a list of 'Diagram's vertically according to a
--   regular spacing, with the given alignment.
vdistribA :: Double     -- ^ How far from one diagram to the next?
          -> VAlignment -- ^ Distribute according to which parts of
                        --   the diagrams ('top', 'vcenter', 'bottom')?
          -> HAlignment -- ^ alignment to use ('left', 'hcenter', 'right')
          -> [Diagram] -> Diagram
vdistribA = distribA V

-- | The mother-combinator for all the @sep@ functions, used
--   internally.
sepA :: ListType -> Double -> Alignment -> [Diagram] -> Diagram
sepA lt sep algn = Compound . Layout (List lt algn (Sep sep))

-- | The mother-combinator for all the @distrib@ functions, used
--   internally.
distribA :: ListType -> Double -> Alignment -> Alignment -> [Diagram] -> Diagram
distribA lt sep alignD algn =
  Compound . Layout (List lt algn (Distrib alignD sep))


-- Explicit subdiagram positioning ---------------------

data Positioned = Positioned [Point] HAlignment VAlignment
instance LayoutClass Positioned [] where
  layoutSizeAndPos _ [] = ((0,0), [])
  layoutSizeAndPos (Positioned ps ha va) dss = ((w,h), pds)
    where (sizes, diagrams) = unzip dss
          (ws, hs) = unzip sizes
          (pxs, pys) = unzip ps

          xs' = zipWith (\x w1 -> x - alignOffset ha w1) pxs ws
          ys' = zipWith (\y h1 -> y - alignOffset va h1) pys hs

          hws = map (/2) ws
          hhs = map (/2) hs
          xmax = maximum (zipWith (+) xs' hws)
          xmin = minimum (zipWith (-) xs' hws)
          w = xmax - xmin
          ymax = maximum (zipWith (+) ys' hhs)
          ymin = minimum (zipWith (-) ys' hhs)
          h = ymax - ymin

          xs'norm = map (subtract (xmin + w/2)) xs'
          ys'norm = map (subtract (ymin + h/2)) ys'

          pds = zipWith3 translate xs'norm ys'norm diagrams

-- | Create a diagram from a list of subdiagrams with explicit
--   positions in a local coordinate system.  Each subdiagram will be
--   positioned with its center at the corresponding
--   position. @position@ is equivalent to @positionA hcenter
--   vcenter@.
position :: [(Point, Diagram)] -> Diagram
position = positionA hcenter vcenter

-- | Create a diagram from a list of subdiagrams with explicit
--   positions in a local coordinate system.  The alignment options
--   specify what part of each subdiagram should be placed on the
--   corresponding position.  For example, @positionA left top@ will
--   position the top left corner of each subdiagram at the
--   corresponding point.
positionA :: HAlignment -> VAlignment -> [(Point,Diagram)] -> Diagram
positionA ha va pds = Compound (Layout (Positioned ps ha va) ds)
  where (ps, ds) = unzip pds

-- | Create a diagram from a list of subdiagrams and a given path, by
--   positioning the subdiagrams at successive vertices of the path.
--   If there are more diagrams than path vertices, the extra diagrams
--   will be discarded.
positionAlong :: [Diagram] -> Path -> Diagram
positionAlong = positionAlongA hcenter vcenter

-- | A version of 'positionAlong' with explicit alignment.
positionAlongA :: HAlignment -> VAlignment -> [Diagram] -> Path -> Diagram
positionAlongA ha va ds p = positionA ha va $ zip (pathToVertices (0,0) p) ds

-- Tree ------------------------------------------

-- | Lay out a 'Tree' (from "Data.Tree") of 'Diagram's in a top-down
--   fashion.  This layout is experimental; future releases of the
--   Diagrams library are planned which will be able to automatically
--   draw edges between nodes in the tree.
tree :: Double    -- ^ separation between layers
     -> Double    -- ^ separation between siblings
     -> Tree Diagram
     -> Diagram
tree _  _  (Node d []) = d
tree ls ss (Node d ts) = vsepA ls hcenter [ d, hsep ss (map (tree ls ss) ts) ]

-- Miscellaneous ---------------------------------

newtype Identity a = Identity a
instance Functor Identity where
  fmap f (Identity x) = Identity (f x)

data Padded = Padded Double Double HAlignment VAlignment
instance LayoutClass Padded Identity where
  layoutSizeAndPos (Padded dw dh ha va) (Identity ((w,h), d))
      = ((w', h'), [translate (align ha w' w) (align va h' h) d])
    where w' = w + dw
          h' = h + dh

-- | Add extra padding to a diagram.  @pad w h d@ is a diagram which
--   is the same as @d@, but with @w@ units added to the width and @h@
--   units added to the height, with @d@ centered in the available
--   space.  Thus @pad w h@ is equivalent to @padA w h hcenter vcenter@.
pad :: Double -> Double -> Diagram -> Diagram
pad pw ph d = padA pw ph hcenter vcenter d

-- | Add extra padding to a diagram, aligning the diagram as indicated
--   within the avilable space.
padA :: Double -> Double -> HAlignment -> VAlignment -> Diagram -> Diagram
padA pw ph aw ah d = Compound (Layout (Padded pw ph aw ah) (Identity d))



data ShowBBox = ShowBBox
instance LayoutClass ShowBBox Identity where
  layoutSizeAndPos _ (Identity ((w,h), d))
      = ((w,h), [defaultAttributes $ rect w h, d])

-- | Show a rectangle denoting a diagram's bounding box, in addition
--   to the diagram itself.
showBBox :: Diagram -> Diagram
showBBox Empty = Empty
showBBox d = Compound . Layout ShowBBox . Identity $ d

-- | Show the bounding boxes of a diagram and all its subdiagrams.
showBBoxes :: Diagram -> Diagram
showBBoxes Empty = Empty
showBBoxes d@(Prim _) = showBBox d
showBBoxes (Ann a d) = showBBox (Ann a (showBBoxes d))
showBBoxes (Compound (Layout l ds)) = showBBox (Compound (Layout l (fmap showBBoxes ds)))
showBBoxes (Union ds) = showBBox (Union (map showBBoxes ds))
showBBoxes (Sized p d) = showBBox (Sized p (showBBoxes d))

data FromSize = FromSize (Double -> Double -> Diagram)
instance LayoutClass FromSize Identity where
  layoutSizeAndPos (FromSize f) (Identity  ((w,h),_))
    = let d = f w h in (fst (sizeAndPos d), [d])

-- | Create one diagram using the current size of another.
--   The new diagram is returned, the old one is discarded.
withSize :: (Double -> Double -> Diagram) -- ^ Function for new diagram
         -> Diagram                       -- ^ Old diagram
         -> Diagram
withSize f dia = Compound (Layout (FromSize f) (Identity dia))

diagSize :: Diagram -> Cont Diagram (Double, Double) 
diagSize d = Cont { runCont = \f -> flip withSize d $ \w h -> f (w, h) } 

-- | Align diagrams into a grid, specifying individual alignments for each item.
--   Warning: there is currently an exponential performace blowup if you nest grids
--   (exponential in how deep the nesting is).
gridAs :: [[(HAlignment, VAlignment)]] -> [[Diagram]] -> Diagram

-- | Align diagrams into a grid with each item aligned as specified.
--   Warning: there is currently an exponential performace blowup if you nest grids
--   (exponential in how deep the nesting is).
gridA :: HAlignment -> VAlignment -> [[Diagram]] -> Diagram

-- | Align diagrams into a grid, with each item centered horizontally and vertically
--   Warning: there is currently an exponential performace blowup if you nest grids.
--   (exponential in how deep the nesting is).
grid :: [[Diagram]] -> Diagram

grid = gridA hcenter vcenter
gridA h v = gridAs (repeat (repeat (h, v)))
gridAs alignss diagss
  = flip runCont id $
     do sizess <- mapM (mapM diagSize) diagss
        let widths = map (maximum . map fst) (transpose sizess)
            heights = map (maximum . map snd) sizess
            newsizess = [[(w, h) | w <- widths] | h <- heights]
            adjss = zipWith (zipWith (\ (nw, nh) (ow, oh) -> (nw - ow, nh - oh))) newsizess sizess
            padTo (width, height) (halign, valign) = padA width height halign valign
        return $ vcat $ map hcat $ zipWith3 (zipWith3 padTo) adjss alignss diagss