{-# LANGUAGE DeriveDataTypeable     #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE UndecidableInstances   #-}
module Plots.Legend
 (
   -- * Legend
   Legend
 , HasLegend (..)

   -- * Drawing a legend
 , drawLegend

 ) where

import           Control.Lens         hiding (none, ( # ))
import           Data.Default
import           Data.Typeable
import           Diagrams.TwoD.Text

import           Diagrams.BoundingBox
import           Diagrams.Prelude

import           Plots.Types

-- | The data type to describe how to draw a legend. For legend entries
--   see 'Plots.Types.LegendEntry'.
data Legend b n = Legend
  { forall b n. Legend b n -> Placement
lPlacement   :: Placement
  , forall b n. Legend b n -> n
lGap         :: n
  , forall b n. Legend b n -> Style V2 n
lStyle       :: Style V2 n
  , forall b n. Legend b n -> n
lSpacing     :: n
  , forall b n. Legend b n -> n
lTextWidth   :: n
  , forall b n. Legend b n -> String -> QDiagram b V2 n Any
lTextF       :: String -> QDiagram b V2 n Any
  , forall b n. Legend b n -> Style V2 n
lTextStyle   :: Style V2 n
  , forall b n. Legend b n -> Orientation
lOrientation :: Orientation
  , forall b n. Legend b n -> Bool
lVisible     :: Bool
  } deriving Typeable

type instance V (Legend b n) = V2
type instance N (Legend b n) = n

class HasLegend a b | a -> b where
  -- | Lens onto the 'Legend' of something.
  legend :: Lens' a (Legend b (N a))

  -- | The 'Placement' of the legend relative to the 'Plots.Axis.Axis'.
  legendPlacement :: Lens' a Placement
  legendPlacement = (Legend b (N a) -> f (Legend b (N a))) -> a -> f a
forall a b. HasLegend a b => Lens' a (Legend b (N a))
legend ((Legend b (N a) -> f (Legend b (N a))) -> a -> f a)
-> ((Placement -> f Placement)
    -> Legend b (N a) -> f (Legend b (N a)))
-> (Placement -> f Placement)
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Legend b (N a) -> Placement)
-> (Legend b (N a) -> Placement -> Legend b (N a))
-> Lens (Legend b (N a)) (Legend b (N a)) Placement Placement
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Legend b (N a) -> Placement
forall b n. Legend b n -> Placement
lPlacement (\Legend b (N a)
l Placement
a -> Legend b (N a)
l {lPlacement :: Placement
lPlacement = Placement
a})

  -- | The gap between the legend and the axis.
  legendGap :: Lens' a (N a)
  legendGap = (Legend b (N a) -> f (Legend b (N a))) -> a -> f a
forall a b. HasLegend a b => Lens' a (Legend b (N a))
legend ((Legend b (N a) -> f (Legend b (N a))) -> a -> f a)
-> ((N a -> f (N a)) -> Legend b (N a) -> f (Legend b (N a)))
-> (N a -> f (N a))
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Legend b (N a) -> N a)
-> (Legend b (N a) -> N a -> Legend b (N a))
-> Lens (Legend b (N a)) (Legend b (N a)) (N a) (N a)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Legend b (N a) -> N a
forall b n. Legend b n -> n
lGap (\Legend b (N a)
l N a
a -> Legend b (N a)
l {lGap :: N a
lGap = N a
a})

  -- | The style applied to the surronding box of the legend.
  legendStyle :: Lens' a (Style V2 (N a))
  legendStyle = (Legend b (N a) -> f (Legend b (N a))) -> a -> f a
forall a b. HasLegend a b => Lens' a (Legend b (N a))
legend ((Legend b (N a) -> f (Legend b (N a))) -> a -> f a)
-> ((Style V2 (N a) -> f (Style V2 (N a)))
    -> Legend b (N a) -> f (Legend b (N a)))
-> (Style V2 (N a) -> f (Style V2 (N a)))
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Legend b (N a) -> Style V2 (N a))
-> (Legend b (N a) -> Style V2 (N a) -> Legend b (N a))
-> Lens
     (Legend b (N a)) (Legend b (N a)) (Style V2 (N a)) (Style V2 (N a))
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Legend b (N a) -> Style V2 (N a)
forall b n. Legend b n -> Style V2 n
lStyle (\Legend b (N a)
l Style V2 (N a)
a -> Legend b (N a)
l {lStyle :: Style V2 (N a)
lStyle = Style V2 (N a)
a})

  -- | The spacing between entries in the legend.
  legendSpacing :: Lens' a (N a)
  legendSpacing = (Legend b (N a) -> f (Legend b (N a))) -> a -> f a
forall a b. HasLegend a b => Lens' a (Legend b (N a))
legend ((Legend b (N a) -> f (Legend b (N a))) -> a -> f a)
-> ((N a -> f (N a)) -> Legend b (N a) -> f (Legend b (N a)))
-> (N a -> f (N a))
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Legend b (N a) -> N a)
-> (Legend b (N a) -> N a -> Legend b (N a))
-> Lens (Legend b (N a)) (Legend b (N a)) (N a) (N a)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Legend b (N a) -> N a
forall b n. Legend b n -> n
lSpacing (\Legend b (N a)
l N a
a -> Legend b (N a)
l {lSpacing :: N a
lSpacing = N a
a})

  -- | The space given for the text in the legend.
  legendTextWidth :: Lens' a (N a)
  legendTextWidth = (Legend b (N a) -> f (Legend b (N a))) -> a -> f a
forall a b. HasLegend a b => Lens' a (Legend b (N a))
legend ((Legend b (N a) -> f (Legend b (N a))) -> a -> f a)
-> ((N a -> f (N a)) -> Legend b (N a) -> f (Legend b (N a)))
-> (N a -> f (N a))
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Legend b (N a) -> N a)
-> (Legend b (N a) -> N a -> Legend b (N a))
-> Lens (Legend b (N a)) (Legend b (N a)) (N a) (N a)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Legend b (N a) -> N a
forall b n. Legend b n -> n
lTextWidth (\Legend b (N a)
l N a
a -> Legend b (N a)
l {lTextWidth :: N a
lTextWidth = N a
a})

  -- | The function to generate the legend text.
  legendTextFunction :: Lens' a (String -> QDiagram b V2 (N a) Any)
  legendTextFunction = (Legend b (N a) -> f (Legend b (N a))) -> a -> f a
forall a b. HasLegend a b => Lens' a (Legend b (N a))
legend ((Legend b (N a) -> f (Legend b (N a))) -> a -> f a)
-> (((String -> QDiagram b V2 (N a) Any)
     -> f (String -> QDiagram b V2 (N a) Any))
    -> Legend b (N a) -> f (Legend b (N a)))
-> ((String -> QDiagram b V2 (N a) Any)
    -> f (String -> QDiagram b V2 (N a) Any))
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Legend b (N a) -> String -> QDiagram b V2 (N a) Any)
-> (Legend b (N a)
    -> (String -> QDiagram b V2 (N a) Any) -> Legend b (N a))
-> Lens
     (Legend b (N a))
     (Legend b (N a))
     (String -> QDiagram b V2 (N a) Any)
     (String -> QDiagram b V2 (N a) Any)
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Legend b (N a) -> String -> QDiagram b V2 (N a) Any
forall b n. Legend b n -> String -> QDiagram b V2 n Any
lTextF (\Legend b (N a)
l String -> QDiagram b V2 (N a) Any
a -> Legend b (N a)
l {lTextF :: String -> QDiagram b V2 (N a) Any
lTextF = String -> QDiagram b V2 (N a) Any
a})

  -- | The style applied to the legend text.
  legendTextStyle :: Lens' a (Style V2 (N a))
  legendTextStyle = (Legend b (N a) -> f (Legend b (N a))) -> a -> f a
forall a b. HasLegend a b => Lens' a (Legend b (N a))
legend ((Legend b (N a) -> f (Legend b (N a))) -> a -> f a)
-> ((Style V2 (N a) -> f (Style V2 (N a)))
    -> Legend b (N a) -> f (Legend b (N a)))
-> (Style V2 (N a) -> f (Style V2 (N a)))
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Legend b (N a) -> Style V2 (N a))
-> (Legend b (N a) -> Style V2 (N a) -> Legend b (N a))
-> Lens
     (Legend b (N a)) (Legend b (N a)) (Style V2 (N a)) (Style V2 (N a))
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Legend b (N a) -> Style V2 (N a)
forall b n. Legend b n -> Style V2 n
lTextStyle (\Legend b (N a)
l Style V2 (N a)
a -> Legend b (N a)
l {lTextStyle :: Style V2 (N a)
lTextStyle = Style V2 (N a)
a})

  -- | The way the legend entries are listed. (This will likely be
  --   replaced by a grid-like system)
  legendOrientation :: Lens' a Orientation
  legendOrientation = (Legend b (N a) -> f (Legend b (N a))) -> a -> f a
forall a b. HasLegend a b => Lens' a (Legend b (N a))
legend ((Legend b (N a) -> f (Legend b (N a))) -> a -> f a)
-> ((Orientation -> f Orientation)
    -> Legend b (N a) -> f (Legend b (N a)))
-> (Orientation -> f Orientation)
-> a
-> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Legend b (N a) -> Orientation)
-> (Legend b (N a) -> Orientation -> Legend b (N a))
-> Lens (Legend b (N a)) (Legend b (N a)) Orientation Orientation
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Legend b (N a) -> Orientation
forall b n. Legend b n -> Orientation
lOrientation (\Legend b (N a)
l Orientation
a -> Legend b (N a)
l {lOrientation :: Orientation
lOrientation = Orientation
a})

instance HasLegend (Legend b n) b where
  legend :: Lens' (Legend b n) (Legend b (N (Legend b n)))
legend = (Legend b (N (Legend b n)) -> f (Legend b (N (Legend b n))))
-> Legend b n -> f (Legend b n)
forall a. a -> a
id

instance HasGap (Legend b n) where
  gap :: Lens' (Legend b n) (N (Legend b n))
gap = (N (Legend b n) -> f (N (Legend b n)))
-> Legend b n -> f (Legend b n)
forall a b. HasLegend a b => Lens' a (N a)
legendGap

instance HasPlacement (Legend b n) where
  placement :: Lens' (Legend b n) Placement
placement = (Placement -> f Placement) -> Legend b n -> f (Legend b n)
forall a b. HasLegend a b => Lens' a Placement
legendPlacement

instance (TypeableFloat n, Renderable (Text n) b) => Default (Legend b n) where
  def :: Legend b n
def = Legend
    { lPlacement :: Placement
lPlacement   = Placement
rightTop
    , lGap :: n
lGap         = n
20
    , lSpacing :: n
lSpacing     = n
20
    , lTextWidth :: n
lTextWidth   = n
60
    , lStyle :: Style V2 n
lStyle       = Style V2 n
forall a. Monoid a => a
mempty
    , lTextF :: String -> QDiagram b V2 n Any
lTextF       = TextAlignment n -> String -> QDiagram b V2 n Any
forall n b.
(TypeableFloat n, Renderable (Text n) b) =>
TextAlignment n -> String -> QDiagram b V2 n Any
mkText (n -> n -> TextAlignment n
forall n. n -> n -> TextAlignment n
BoxAlignedText n
0 n
0.5)
    , lTextStyle :: Style V2 n
lTextStyle   = Style V2 n
forall a. Monoid a => a
mempty Style V2 n -> (Style V2 n -> Style V2 n) -> Style V2 n
forall a b. a -> (a -> b) -> b
& Measure n -> Style V2 n -> Style V2 n
forall a n.
(N a ~ n, Typeable n, HasStyle a) =>
Measure n -> a -> a
fontSize (n -> Measure n
forall n. n -> Measure n
output n
11)
    , lOrientation :: Orientation
lOrientation = Orientation
Vertical
    , lVisible :: Bool
lVisible     = Bool
True
    }

instance HasVisibility (Legend b n) where
  visible :: Lens' (Legend b n) Bool
visible = (Legend b n -> Bool)
-> (Legend b n -> Bool -> Legend b n) -> Lens' (Legend b n) Bool
forall s a b t. (s -> a) -> (s -> b -> t) -> Lens s t a b
lens Legend b n -> Bool
forall b n. Legend b n -> Bool
lVisible (\Legend b n
l Bool
a -> Legend b n
l {lVisible :: Bool
lVisible = Bool
a})

instance TypeableFloat n => HasStyle (Legend b n) where
  applyStyle :: Style (V (Legend b n)) (N (Legend b n)) -> Legend b n -> Legend b n
applyStyle Style (V (Legend b n)) (N (Legend b n))
sty = ASetter (Legend b n) (Legend b n) (Style V2 n) (Style V2 n)
-> (Style V2 n -> Style V2 n) -> Legend b n -> Legend b n
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter (Legend b n) (Legend b n) (Style V2 n) (Style V2 n)
forall a b. HasLegend a b => Lens' a (Style V2 (N a))
legendStyle (Style (V (Style V2 n)) (N (Style V2 n)) -> Style V2 n -> Style V2 n
forall a. HasStyle a => Style (V a) (N a) -> a -> a
applyStyle Style (V (Style V2 n)) (N (Style V2 n))
Style (V (Legend b n)) (N (Legend b n))
sty)

instance HasOrientation (Legend b n) where
  orientation :: Lens' (Legend b n) Orientation
orientation = (Orientation -> f Orientation) -> Legend b n -> f (Legend b n)
forall a b. HasLegend a b => Lens' a Orientation
legendOrientation

-- | Draw a legend to the bounding box using the legend entries and
--   legend options.
drawLegend
  :: (TypeableFloat n,
      Renderable (Path V2 n) b)
  => BoundingBox V2 n                -- ^ bounding box to place legend against
  -> [(QDiagram b V2 n Any, String)] -- ^ diagram pictures along with their key
  -> Legend b n                      -- ^ options for drawing the legend
  -> QDiagram b V2 n Any             -- ^ rendered legend
drawLegend :: forall n b.
(TypeableFloat n, Renderable (Path V2 n) b) =>
BoundingBox V2 n
-> [(QDiagram b V2 n Any, String)]
-> Legend b n
-> QDiagram b V2 n Any
drawLegend BoundingBox V2 n
bb [(QDiagram b V2 n Any, String)]
entries Legend b n
l
  | Legend b n
l Legend b n -> Getting Bool (Legend b n) Bool -> Bool
forall s a. s -> Getting a s a -> a
^. Getting Bool (Legend b n) Bool
forall a. HasVisibility a => Lens' a Bool
hidden Bool -> Bool -> Bool
|| [(QDiagram b V2 n Any, String)] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(QDiagram b V2 n Any, String)]
entries = QDiagram b V2 n Any
forall a. Monoid a => a
mempty
  | Bool
otherwise   = BoundingBox V2 n
-> Placement -> n -> QDiagram b V2 n Any -> QDiagram b V2 n Any
forall n a b.
(InSpace V2 n a, SameSpace a b, Enveloped a, HasOrigin b,
 Alignable b) =>
a -> Placement -> n -> b -> b
placeAgainst
                    BoundingBox V2 n
bb
                    (Legend b n
l Legend b n -> Getting Placement (Legend b n) Placement -> Placement
forall s a. s -> Getting a s a -> a
^. Getting Placement (Legend b n) Placement
forall a b. HasLegend a b => Lens' a Placement
legendPlacement)
                    (Legend b n
l Legend b n -> Getting n (Legend b n) n -> n
forall s a. s -> Getting a s a -> a
^. Getting n (Legend b n) n
forall a b. HasLegend a b => Lens' a (N a)
legendGap)
                    (QDiagram b V2 n Any
ledge QDiagram b V2 n Any -> QDiagram b V2 n Any -> QDiagram b V2 n Any
forall a. Semigroup a => a -> a -> a
<> QDiagram b V2 n Any
back)
  where
    w :: n
w = Legend b n
l Legend b n -> Getting n (Legend b n) n -> n
forall s a. s -> Getting a s a -> a
^. Getting n (Legend b n) n
forall a b. HasLegend a b => Lens' a (N a)
legendTextWidth
    h :: n
h = Legend b n
l Legend b n -> Getting n (Legend b n) n -> n
forall s a. s -> Getting a s a -> a
^. Getting n (Legend b n) n
forall a b. HasLegend a b => Lens' a (N a)
legendSpacing
    --
    ledge :: QDiagram b V2 n Any
ledge = ((QDiagram b V2 n Any, String) -> QDiagram b V2 n Any)
-> [(QDiagram b V2 n Any, String)] -> [QDiagram b V2 n Any]
forall a b. (a -> b) -> [a] -> [b]
map (QDiagram b V2 n Any, String) -> QDiagram b V2 n Any
mkLabels [(QDiagram b V2 n Any, String)]
entries
              # orient (l ^. legendOrientation) hcat vcat
              # alignTL

    back :: QDiagram b V2 n Any
back = QDiagram b V2 n Any
backRect
             # applyStyle (l ^. legendStyle)
             # alignTL
    backRect :: QDiagram b V2 n Any
backRect = Orientation
-> QDiagram b V2 n Any
-> QDiagram b V2 n Any
-> QDiagram b V2 n Any
forall o a. HasOrientation o => o -> a -> a -> a
orient (Legend b n
l Legend b n
-> Getting Orientation (Legend b n) Orientation -> Orientation
forall s a. s -> Getting a s a -> a
^. Getting Orientation (Legend b n) Orientation
forall a b. HasLegend a b => Lens' a Orientation
legendOrientation)
      (n -> n -> QDiagram b V2 n Any
forall n t. (InSpace V2 n t, TrailLike t) => n -> n -> t
rect (n
nEntries n -> n -> n
forall a. Num a => a -> a -> a
* n
entryWidth) n
h             )
      (n -> n -> QDiagram b V2 n Any
forall n t. (InSpace V2 n t, TrailLike t) => n -> n -> t
rect n
entryWidth              (n
h n -> n -> n
forall a. Num a => a -> a -> a
* n
nEntries))
    nEntries :: n
nEntries = Int -> n
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([(QDiagram b V2 n Any, String)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(QDiagram b V2 n Any, String)]
entries)

    -- Each legend picture has a width equal to the height of each
    -- legend entry. The picture also has a 5 unit buffer either side of
    -- it.
    entryWidth :: n
entryWidth = n
w n -> n -> n
forall a. Num a => a -> a -> a
+ n
10 n -> n -> n
forall a. Num a => a -> a -> a
+ n
h

    mkLabels :: (QDiagram b V2 n Any, String) -> QDiagram b V2 n Any
mkLabels (QDiagram b V2 n Any
pic, String
txt) = n -> QDiagram b V2 n Any
forall (v :: * -> *) n b m.
(Metric v, R1 v, OrderedField n) =>
n -> QDiagram b v n m
strutX n
5 QDiagram b V2 n Any -> QDiagram b V2 n Any -> QDiagram b V2 n Any
forall n a.
(InSpace V2 n a, Juxtaposable a, Semigroup a) =>
a -> a -> a
||| QDiagram b V2 n Any
pic' QDiagram b V2 n Any -> QDiagram b V2 n Any -> QDiagram b V2 n Any
forall n a.
(InSpace V2 n a, Juxtaposable a, Semigroup a) =>
a -> a -> a
||| n -> QDiagram b V2 n Any
forall (v :: * -> *) n b m.
(Metric v, R1 v, OrderedField n) =>
n -> QDiagram b v n m
strutX n
5 QDiagram b V2 n Any -> QDiagram b V2 n Any -> QDiagram b V2 n Any
forall n a.
(InSpace V2 n a, Juxtaposable a, Semigroup a) =>
a -> a -> a
||| QDiagram b V2 n Any
label where
      pic' :: QDiagram b V2 n Any
pic'  = QDiagram b V2 n Any
pic QDiagram b V2 n Any
-> (QDiagram b V2 n Any -> QDiagram b V2 n Any)
-> QDiagram b V2 n Any
forall a b. a -> (a -> b) -> b
# BoundingBox V2 n -> QDiagram b V2 n Any -> QDiagram b V2 n Any
forall (v :: * -> *) n a m b.
(InSpace v n a, Monoid' m, Enveloped a) =>
a -> QDiagram b v n m -> QDiagram b v n m
withEnvelope (Point V2 n -> Point V2 n -> BoundingBox V2 n
forall (v :: * -> *) n.
(Additive v, Foldable v, Ord n) =>
Point v n -> Point v n -> BoundingBox v n
fromCorners (n -> Point V2 n
forall (f :: * -> *) a. Applicative f => a -> f a
pure (-n
hn -> n -> n
forall a. Fractional a => a -> a -> a
/n
2)) (n -> Point V2 n
forall (f :: * -> *) a. Applicative f => a -> f a
pure (n
hn -> n -> n
forall a. Fractional a => a -> a -> a
/n
2)))
      label :: QDiagram b V2 n Any
label = Getting
  (String -> QDiagram b V2 n Any)
  (Legend b n)
  (String -> QDiagram b V2 n Any)
-> Legend b n -> String -> QDiagram b V2 n Any
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting
  (String -> QDiagram b V2 n Any)
  (Legend b n)
  (String -> QDiagram b V2 n Any)
forall a b.
HasLegend a b =>
Lens' a (String -> QDiagram b V2 (N a) Any)
legendTextFunction Legend b n
l String
txt
                # applyStyle (l ^. legendTextStyle)
                # withEnvelope (fromCorners origin (mkP2 w h) # moveTo (mkP2 0 (-h/2)))

-- wrapPic :: RealFloat n => V2 n -> QDiagram b V2 n Any -> QDiagram b V2 n Any
-- wrapPic ((^/ 2) -> v) d
--   = d # sizedAs (fromCorners (origin .-^ v) (origin .+^ v))