--  This Source Code Form is subject to the terms of the Mozilla Public
--  License, v. 2.0. If a copy of the MPL was not distributed with this
--  file, You can obtain one at http://mozilla.org/MPL/2.0/.

{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE NamedFieldPuns #-}

-- |
-- Copyright: © 2018-2020 IOHK
-- License: MPL-2.0
-- Stability: experimental
--
-- Plot 'Histogram' as unicode strings; for terminal plotting.
-- For example, one could plot answers and their frequencies to the question:
--
-- /"Who's your favorite metal band?"/
--
-- > plot $ Histogram
-- >     { width  = 80
-- >     , height = 24
-- >     , bins   =
-- >         [ ( "Metallica", 78 )
-- >         , ( "Iron Maiden", 61 )
-- >         , ( "Slayer", 16 )
-- >         , ( "Dimmu Borgir", 3 )
-- >         , ( "Ghost", 48 )
-- >         ]
-- >     }
--
-- > 78 ┤┌─────────────┐
-- >    ││             │
-- >    ││             │
-- >    ││             │
-- >    ││             │
-- > 62 ┤│             │
-- >    ││             ├─────────────┐
-- >    ││             │             │
-- >    ││             │             │
-- >    ││             │             │
-- > 47 ┤│             │             │                            ┌─────────────┐
-- >    ││             │             │                            │             │
-- >    ││             │             │                            │             │
-- >    ││             │             │                            │             │
-- >    ││             │             │                            │             │
-- > 31 ┤│             │             │                            │             │
-- >    ││             │             │                            │             │
-- >    ││             │             │                            │             │
-- >    ││             │             │                            │             │
-- >    ││             │             │                            │             │
-- > 16 ┤│             │             ├─────────────┐              │             │
-- >    ││             │             │             │              │             │
-- >    ││             │             │             │              │             │
-- >    ││             │             │             │              │             │
-- >    ││             │             │             ├──────────────┤             │
-- >        Metallica    Iron Maiden     Slayer      Dimmu Borgir      Ghost
module Plot.Text.Histogram
    ( Histogram (..)
    , plot
    ) where

import Control.Applicative
    ( (<|>) )
import Data.List
    ( nub )
import Data.Maybe
    ( fromJust, fromMaybe )

-- | Model a distribution as an [Histogram](https://en.wikipedia.org/wiki/Histogram).
-- This representation works for continuous non-overlapping distributions. Axes
-- are labelled automatically based on the bins.
data Histogram = Histogram
    { Histogram -> Int
width  :: !Int
    , Histogram -> Int
height :: !Int
    , Histogram -> [(String, Int)]
bins   :: ![(String, Int)]
    } deriving (Int -> Histogram -> ShowS
[Histogram] -> ShowS
Histogram -> String
(Int -> Histogram -> ShowS)
-> (Histogram -> String)
-> ([Histogram] -> ShowS)
-> Show Histogram
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Histogram] -> ShowS
$cshowList :: [Histogram] -> ShowS
show :: Histogram -> String
$cshow :: Histogram -> String
showsPrec :: Int -> Histogram -> ShowS
$cshowsPrec :: Int -> Histogram -> ShowS
Show, ReadPrec [Histogram]
ReadPrec Histogram
Int -> ReadS Histogram
ReadS [Histogram]
(Int -> ReadS Histogram)
-> ReadS [Histogram]
-> ReadPrec Histogram
-> ReadPrec [Histogram]
-> Read Histogram
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Histogram]
$creadListPrec :: ReadPrec [Histogram]
readPrec :: ReadPrec Histogram
$creadPrec :: ReadPrec Histogram
readList :: ReadS [Histogram]
$creadList :: ReadS [Histogram]
readsPrec :: Int -> ReadS Histogram
$creadsPrec :: Int -> ReadS Histogram
Read)

-- | Convert an 'Histogram' to 'String'
plot :: Histogram -> String
plot :: Histogram -> String
plot Histogram{Int
width :: Int
$sel:width:Histogram :: Histogram -> Int
width,Int
height :: Int
$sel:height:Histogram :: Histogram -> Int
height,[(String, Int)]
bins :: [(String, Int)]
$sel:bins:Histogram :: Histogram -> [(String, Int)]
bins}
    | Int
height Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 0 Bool -> Bool -> Bool
|| Int
width Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 0 =
        ""
    | Bool
otherwise =
        [String] -> String
unlines ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ Plot -> [String]
getPlot (Plot -> [String]) -> Plot -> [String]
forall a b. (a -> b) -> a -> b
$ [Plot] -> Plot
forall a. Monoid a => [a] -> a
mconcat ([Plot] -> Plot) -> [Plot] -> Plot
forall a b. (a -> b) -> a -> b
$ Int -> [Axis] -> Plot
axis Int
widthYAxis [Axis]
ys Plot -> [Plot] -> [Plot]
forall a. a -> [a] -> [a]
: (Bar -> Plot
bar (Bar -> Plot) -> ((String, Int) -> Bar) -> (String, Int) -> Plot
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, Int) -> Bar
forall i. Real i => (String, i) -> Bar
mkBar ((String, Int) -> Plot) -> [(String, Int)] -> [Plot]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(String, Int)]
bins)
  where
    highest :: Int
highest = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ((String, Int) -> Int
forall a b. (a, b) -> b
snd ((String, Int) -> Int) -> [(String, Int)] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(String, Int)]
bins)
    nbLbls :: Int
nbLbls  = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min ([(String, Int)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(String, Int)]
bins) (Int
height Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` 3)

    widthYAxis :: Int
widthYAxis = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ((Int -> Int -> Int
forall a. Num a => a -> a -> a
+3) (Int -> Int) -> (String -> Int) -> String -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (String -> Int) -> [String] -> [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
lbls)

    mkBar :: (String, i) -> Bar
mkBar (label :: String
label, x :: i
x) = Bar :: Int -> Int -> String -> Bar
Bar
        { $sel:width:Bar :: Int
width  = 1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- 2 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
widthYAxis) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` [(String, Int)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(String, Int)]
bins
        , $sel:height:Bar :: Int
height = Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
round (Int -> Double
forall i. Real i => i -> Double
double Int
height Double -> Double -> Double
forall a. Num a => a -> a -> a
* i -> Double
forall i. Real i => i -> Double
double i
x Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall i. Real i => i -> Double
double Int
highest)
        , String
$sel:label:Bar :: String
label :: String
label
        }

    lbls :: [String]
lbls = [String] -> [String]
forall a. Eq a => [a] -> [a]
nub [ Integer -> String
forall a. Show a => a -> String
show (Integer -> String) -> Integer -> String
forall a b. (a -> b) -> a -> b
$ Double -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
round (Int -> Double
forall i. Real i => i -> Double
double Int
highest Double -> Double -> Double
forall a. Num a => a -> a -> a
* Int -> Double
forall i. Real i => i -> Double
double Int
i Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall i. Real i => i -> Double
double Int
nbLbls)
               | Int
i <- [ 1 .. Int
nbLbls ]
               ]

    ys :: [Axis]
ys = Int -> String -> Axis
Axis (Double -> Int
forall a b. (RealFrac a, Integral b) => a -> b
round (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Double
forall i. Real i => i -> Double
double Int
height Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Int -> Double
forall i. Real i => i -> Double
double ([String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
lbls) Double -> Double -> Double
forall a. Num a => a -> a -> a
- 1) (String -> Axis) -> [String] -> [Axis]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [String]
lbls

-- | Represent a labelled vertical bar.
data Bar = Bar
    { Bar -> Int
width  :: Int
    , Bar -> Int
height :: Int
    , Bar -> String
label  :: String
    }

bar :: Bar -> Plot
bar :: Bar -> Plot
bar (Bar w :: Int
w h :: Int
h l :: String
l)
    | Int
h Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 0 =
        [String] -> Plot
Plot ([String] -> Plot) -> [String] -> Plot
forall a b. (a -> b) -> a -> b
$ Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
w '─' String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
lbl

    | Bool
otherwise =
        [String] -> Plot
Plot ([String] -> Plot) -> [String] -> Plot
forall a b. (a -> b) -> a -> b
$ [ Int -> Char -> ShowS
forall e. Int -> e -> [e] -> [e]
padMiddle Int
w '─' "┌┐" ]
            [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ ( Int -> Char -> ShowS
forall e. Int -> e -> [e] -> [e]
padMiddle Int
w ' ' "││" String -> Int -> [String]
forall e. e -> Int -> [e]
.* (Int
h Int -> Int -> Int
forall a. Num a => a -> a -> a
- 1) )
            [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [String]
lbl
  where
    lbl :: [String]
lbl = [ Int -> Char -> ShowS
forall e. Int -> e -> [e] -> [e]
padBoth Int
w ' ' String
l ]

-- | Represent one mark on the y-axis.
data Axis = Axis
    { Axis -> Int
height :: Int
    , Axis -> String
label  :: String
    }

axis :: Int -> [Axis] -> Plot
axis :: Int -> [Axis] -> Plot
axis _ [] = [String] -> Plot
Plot []
axis w :: Int
w es :: [Axis]
es = [String] -> Plot
Plot ((Axis -> [String]) -> [Axis] -> [String]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Axis -> [String]
each ([Axis] -> [Axis]
forall a. [a] -> [a]
reverse [Axis]
es) [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ [ Int -> Char -> String
forall a. Int -> a -> [a]
replicate Int
w ' ' ])
  where
    each :: Axis -> [String]
    each :: Axis -> [String]
each (Axis h :: Int
h l :: String
l) =
        Int -> Char -> ShowS
forall e. Int -> e -> [e] -> [e]
padLeft Int
w ' ' (String
l String -> ShowS
forall a. Semigroup a => a -> a -> a
<> " ┤ ") String -> [String] -> [String]
forall a. a -> [a] -> [a]
: Int -> String -> [String]
forall a. Int -> a -> [a]
replicate Int
h (Int -> Char -> ShowS
forall e. Int -> e -> [e] -> [e]
padLeft Int
w ' ' "│ ")

--
-- Plot
--

-- | A plot represents a partial plot on the console which is modeled as a 2D
-- grid (a list of list of chars). Plots can be combined horizontally and merged
-- into a bigger plot. Edges are smoothen into more appropriate charaters.
--
-- For any plots of length N & M, the plot (N <> M) is of length (N + M - 1).
newtype Plot = Plot { Plot -> [String]
getPlot :: [String] }

instance Semigroup Plot where
    Plot xs :: [String]
xs <> :: Plot -> Plot -> Plot
<> Plot ys :: [String]
ys
        | [String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
xs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
ys =
            let ys' :: [String]
ys' = Int -> [String] -> [String]
fill (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
y) [String]
ys in [String] -> Plot
Plot ([String] -> Plot) -> [String] -> Plot
forall a b. (a -> b) -> a -> b
$ (String -> ShowS) -> [String] -> [String] -> [String]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith String -> ShowS
merge [String]
xs [String]
ys'

        | Bool
otherwise =
            let xs' :: [String]
xs' = Int -> [String] -> [String]
fill (Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
x) [String]
xs in [String] -> Plot
Plot ([String] -> Plot) -> [String] -> Plot
forall a b. (a -> b) -> a -> b
$ (String -> ShowS) -> [String] -> [String] -> [String]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith String -> ShowS
merge [String]
xs' [String]
ys
      where
        (x :: Int
x, y :: Int
y) = ([String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
xs, [String] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [String]
ys)

        merge :: String -> ShowS
merge [] ys :: String
ys = String
ys
        merge xs :: String
xs [] = String
xs
        merge xs :: String
xs ys :: String
ys
            | String -> Bool
match "││" = ShowS
into "│"
            | String -> Bool
match "│┌" = ShowS
into "├"
            | String -> Bool
match "│ " = ShowS
into "│"
            | String -> Bool
match "┐│" = ShowS
into "┤"
            | String -> Bool
match "┐┌" = ShowS
into "┬"
            | String -> Bool
match "┐ " = ShowS
into "┐"
            | String -> Bool
match " │" = ShowS
into "│"
            | String -> Bool
match " ┌" = ShowS
into "┌"
            | String -> Bool
match "─│" = ShowS
into "┘"
            | String -> Bool
match "│─" = ShowS
into "└"
            | String -> Bool
match "──" = ShowS
into "─"
            | String -> Bool
match "┬─" = ShowS
into "┬"
            | String -> Bool
match "┐─" = ShowS
into "─"
            | String -> Bool
match "─┌" = ShowS
into "─"
            | Bool
otherwise  = ShowS
into " "
          where
            match :: String -> Bool
match [a :: Char
a,b :: Char
b] = String -> Char
forall a. [a] -> a
last String
xs Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
a Bool -> Bool -> Bool
&& String -> Char
forall a. [a] -> a
head String
ys Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
b
            into :: ShowS
into s :: String
s = ShowS
forall a. [a] -> [a]
init String
xs String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
s String -> ShowS
forall a. [a] -> [a] -> [a]
++ ShowS
forall a. [a] -> [a]
tail String
ys

instance Monoid Plot where
    mempty :: Plot
mempty = [String] -> Plot
Plot []

--
-- Helpers
--

infixl 5 .*

(.*) :: e -> Int -> [e]
.* :: e -> Int -> [e]
(.*) = (Int -> e -> [e]) -> e -> Int -> [e]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Int -> e -> [e]
forall a. Int -> a -> [a]
replicate

double :: Real i => i -> Double
double :: i -> Double
double = Rational -> Double
forall a. Fractional a => Rational -> a
fromRational (Rational -> Double) -> (i -> Rational) -> i -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> Rational
forall a. Real a => a -> Rational
toRational

padRight :: Int -> e -> [e] -> [e]
padRight :: Int -> e -> [e] -> [e]
padRight n :: Int
n e :: e
e es :: [e]
es =
    [e]
es [e] -> [e] -> [e]
forall a. [a] -> [a] -> [a]
++ Int -> e -> [e]
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- [e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [e]
es) e
e

padLeft :: Int -> e -> [e] -> [e]
padLeft :: Int -> e -> [e] -> [e]
padLeft n :: Int
n e :: e
e es :: [e]
es =
    Int -> e -> [e]
forall a. Int -> a -> [a]
replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- [e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [e]
es) e
e [e] -> [e] -> [e]
forall a. [a] -> [a] -> [a]
++ [e]
es

padBoth :: Int -> e -> [e] -> [e]
padBoth :: Int -> e -> [e] -> [e]
padBoth n :: Int
n e :: e
e es :: [e]
es
    | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n         = [e]
es
    | Int
len Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` 2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 0 = Int -> e -> [e] -> [e]
forall e. Int -> e -> [e] -> [e]
padBoth Int
n e
e ([e] -> [e]) -> [e] -> [e]
forall a b. (a -> b) -> a -> b
$ Int -> e -> [e] -> [e]
forall e. Int -> e -> [e] -> [e]
padRight (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) e
e [e]
es
    | Bool
otherwise        = Int -> e -> [e] -> [e]
forall e. Int -> e -> [e] -> [e]
padBoth Int
n e
e ([e] -> [e]) -> [e] -> [e]
forall a b. (a -> b) -> a -> b
$ Int -> e -> [e] -> [e]
forall e. Int -> e -> [e] -> [e]
padLeft  (Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) e
e [e]
es
  where
    len :: Int
len = [e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [e]
es

padMiddle :: Int -> e -> [e] -> [e]
padMiddle :: Int -> e -> [e] -> [e]
padMiddle n :: Int
n e :: e
e es :: [e]
es =
    Int -> [e] -> [e]
forall a. Int -> [a] -> [a]
take Int
half [e]
es [e] -> [e] -> [e]
forall a. [a] -> [a] -> [a]
++ Int -> e -> [e]
forall a. Int -> a -> [a]
replicate Int
δ e
e [e] -> [e] -> [e]
forall a. [a] -> [a] -> [a]
++ Int -> [e] -> [e]
forall a. Int -> [a] -> [a]
drop Int
half [e]
es
  where
    δ :: Int
δ    = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- [e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [e]
es
    half :: Int
half = [e] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [e]
es Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` 2

fill :: Int -> [String] -> [String]
fill :: Int -> [String] -> [String]
fill n :: Int
n []    = Int -> String -> [String]
forall a. Int -> a -> [a]
replicate Int
n String
forall a. Monoid a => a
mempty
fill n :: Int
n (h :: String
h:q :: [String]
q) = Int -> String -> [String]
forall a. Int -> a -> [a]
replicate Int
n (Int -> Char -> String
forall a. Int -> a -> [a]
replicate (String -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
h) ' ') [String] -> [String] -> [String]
forall a. [a] -> [a] -> [a]
++ (String
hString -> [String] -> [String]
forall a. a -> [a] -> [a]
:[String]
q)