{-# LANGUAGE BangPatterns, FlexibleContexts #-}
-- | This module contain functions for drawing diagrams of
-- dendrograms.
module Diagrams.Dendrogram
    ( -- * High-level interface
      -- $runnableExample
      dendrogram
    , Width(..)

      -- * Low-level interface
    , dendrogramPath
    , fixedWidth
    , variableWidth
    , X
    , hcatB
    ) where

-- from base
import Control.Arrow (first, second)

-- from hierarchical-clustering
import Data.Clustering.Hierarchical (Dendrogram(..), elements)

-- from diagrams-lib
import Diagrams.Prelude


-- $runnableExample
--
-- Given a dendrogram @dendro :: 'Dendrogram' a@ and a function
-- @drawItem :: a -> Diagram b R2@ for drawing the items on the
-- leaves of @dendro@, just use @'dendrogram' 'Variable' drawItem
-- dendro :: Diagram b R2@ to draw a diagram of @dendro@.
--
-- Runnable example which produces something like
-- <https://patch-tag.com/r/felipe/hierarchical-clustering-diagrams/snapshot/current/content/pretty/example.png>:
--
-- @
--import Data.Clustering.Hierarchical (Dendrogram(..))
--import Diagrams.Prelude (Diagram, R2, atop, lw, pad, roundedRect, text, (\#))
--import Diagrams.Backend.Cairo.CmdLine (Cairo, defaultMain)
--import qualified Diagrams.Dendrogram as D
--
--main :: IO ()
--main = defaultMain diagram
--
--diagram :: Diagram Cairo R2
--diagram = D.'dendrogram' 'D.Fixed' char test \# lw 0.1 \# pad 1.1
--
--char :: Char -> Diagram Cairo R2
--char c = pad 1.3 $ roundedRect (1,1) 0.1 \`atop\` text [c]
--
--test :: Dendrogram Char
--test = Branch 5
--         (Branch 2
--           (Branch 1
--             (Leaf \'A\')
--             (Leaf \'B\'))
--           (Leaf \'C\'))
--         (Leaf \'D\')
-- @


-- | @dendrogram width drawItem dendro@ is a drawing of the
-- dendrogram @dendro@ using @drawItem@ to draw its leafs.  The
-- @width@ parameter controls how whether all items have the same
-- width or not ('Fixed' or 'Variable', respectively, see
-- 'Width').
--
-- Note: you should probably use 'alignT' to align your items.
dendrogram :: (Monoid m, Renderable (Path R2) b) =>
              Width
           -> (a -> AnnDiagram b R2 m)
           -> Dendrogram a
           -> AnnDiagram b R2 m
dendrogram width_ drawItem dendro = (stroke path_ # value mempty)
                                               ===
                                         (items # alignL)
  where
    (path_, items) =
        case width_ of
          Fixed -> let drawnItems = map drawItem (elements dendro)
                       w = width (head drawnItems)
                       (dendro', _) = fixedWidth w dendro
                   in (dendrogramPath dendro', hcatB drawnItems)
          Variable -> first dendrogramPath $ variableWidth drawItem dendro


-- | The width of the items on the leafs of a dendrogram.
data Width =
      Fixed
      -- ^ @Fixed@ assumes that all items have a fixed width
      -- (which is automatically calculated).  This mode is
      -- faster than @Variable@, especially when you have many
      -- items.
    | Variable
      -- ^ @Variable@ does not assume that all items have a fixed
      -- width, so each item may have a different width.  This
      -- mode is slower since it has to calculate the width of
      -- each item separately.


-- | A dendrogram path that can be 'stroke'@d@ later.  This function
-- assumes that the 'Leaf'@s@ of your 'Dendrogram' are already in
-- the right position.
dendrogramPath :: Dendrogram X -> Path R2
dendrogramPath = mconcat . fst . go []
    where
      go acc (Leaf x)       = (acc, (x, 0))
      go acc (Branch d l r) = (path : acc'', pos)
        where
          (acc',  (!xL, !yL)) = go acc  l
          (acc'', (!xR, !yR)) = go acc' r

          path = fromVertices [ P (xL, yL)
                              , P (xL, d)
                              , P (xR, d)
                              , P (xR, yR)]
          pos  = ((xL + xR) / 2, d)


-- | The horizontal position of a dendrogram Leaf.
type X = Double


-- | @fixedWidth w@ positions the 'Leaf'@s@ of a 'Dendrogram'
-- assuming that they have the same width @w@.  Also returns the
-- total width.
fixedWidth :: Double -> Dendrogram a -> (Dendrogram X, Double)
fixedWidth w = second (subtract half_w) . go half_w
    where
      half_w = w/2
      go !y (Leaf _)       = (Leaf y, y + w)
      go !y (Branch d l r) = (Branch d l' r', y'')
          where
            (l', !y')  = go y  l
            (r', !y'') = go y' r


-- | @variableWidth draw@ positions the 'Leaf'@s@ of a
-- 'Dendrogram' according to the diagram generated by 'draw'.
-- Each 'Leaf' may have a different width.  Also returns the
-- resulting diagram having all 'Leaf'@s@ drawn side-by-side.
--
-- Note: you should probably use 'alignT' to align your items.
variableWidth :: (Monoid m) =>
                 (a -> AnnDiagram b R2 m)
              -> Dendrogram a
              -> (Dendrogram X, AnnDiagram b R2 m)
variableWidth draw = finish . go 0 []
    where
      go !y acc (Leaf a) = (Leaf y', y'', dia : acc)
          where
            dia  = draw a
            !w   = width dia
            !y'  = y + w/2
            !y'' = y + w
      go !y acc (Branch d l r) = (Branch d l' r', y'', acc'')
          where
            (l', !y',  acc'') = go y  acc' l -- yes, this is acc'
            (r', !y'', acc')  = go y' acc r
      finish (dendro, _, dias) = (dendro, hcatB dias)
      -- We used to concatenate diagrams inside 'go' using (|||).
      -- However, pathological dendrograms (such as those created
      -- using single linkage) may be highly unbalanced, creating
      -- a performance problem for 'variableWidth'.


-- | Like 'hcat', but balanced.  Much better performance.  Use it
-- for concatenating the items of your dendrogram.
hcatB :: Monoid m => [AnnDiagram b R2 m] -> AnnDiagram b R2 m
hcatB [y] = y
hcatB ys  = hcatB $ dubs ys
  where dubs (x1:x2:xs) = x1 ||| x2 : dubs xs
        dubs [x]        = [x]
        dubs []         = []