{-# LANGUAGE OverloadedStrings #-}

module DataFrame.Display.Terminal.PrettyPrint where

import qualified Data.Text as T
import qualified Data.Vector as V

{- | Output format for 'showTable'. 'Plain' renders a terminal-style table with
ASCII borders; 'Markdown' renders a GitHub-flavoured pipe table suitable for
notebooks.
-}
data RenderFormat = Plain | Markdown
    deriving (Int -> RenderFormat -> ShowS
[RenderFormat] -> ShowS
RenderFormat -> String
(Int -> RenderFormat -> ShowS)
-> (RenderFormat -> String)
-> ([RenderFormat] -> ShowS)
-> Show RenderFormat
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> RenderFormat -> ShowS
showsPrec :: Int -> RenderFormat -> ShowS
$cshow :: RenderFormat -> String
show :: RenderFormat -> String
$cshowList :: [RenderFormat] -> ShowS
showList :: [RenderFormat] -> ShowS
Show, RenderFormat -> RenderFormat -> Bool
(RenderFormat -> RenderFormat -> Bool)
-> (RenderFormat -> RenderFormat -> Bool) -> Eq RenderFormat
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: RenderFormat -> RenderFormat -> Bool
== :: RenderFormat -> RenderFormat -> Bool
$c/= :: RenderFormat -> RenderFormat -> Bool
/= :: RenderFormat -> RenderFormat -> Bool
Eq)

-- Utility functions to show a DataFrame as a Markdown-ish table.

-- Fill functions, adapted from:
-- https://stackoverflow.com/questions/5929377/format-list-output-in-haskell
type Filler = Int -> T.Text -> T.Text

-- A description of one table column.
data ColDesc t = ColDesc
    { forall t. ColDesc t -> Filler
colTitleFill :: Filler
    , forall t. ColDesc t -> Text
colTitle :: T.Text
    , forall t. ColDesc t -> Filler
colValueFill :: Filler
    }

-- Fill text to a given width with a pad character, aligning left, right, or center.
fillLeft :: Char -> Int -> T.Text -> T.Text
fillLeft :: Char -> Filler
fillLeft Char
c Int
n Text
s = Text
s Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Filler
T.replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
T.length Text
s) (Char -> Text
T.singleton Char
c)

fillRight :: Char -> Int -> T.Text -> T.Text
fillRight :: Char -> Filler
fillRight Char
c Int
n Text
s = Filler
T.replicate (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
T.length Text
s) (Char -> Text
T.singleton Char
c) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
s

fillCenter :: Char -> Int -> T.Text -> T.Text
fillCenter :: Char -> Filler
fillCenter Char
c Int
n Text
s =
    Filler
T.replicate Int
l (Char -> Text
T.singleton Char
c) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
s Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Filler
T.replicate Int
r (Char -> Text
T.singleton Char
c)
  where
    x :: Int
x = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Text -> Int
T.length Text
s
    l :: Int
l = Int
x Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
    r :: Int
r = Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l

-- Fill with spaces.
left :: Int -> T.Text -> T.Text
left :: Filler
left = Char -> Filler
fillLeft Char
' '

right :: Int -> T.Text -> T.Text
right :: Filler
right = Char -> Filler
fillRight Char
' '

center :: Int -> T.Text -> T.Text
center :: Filler
center = Char -> Filler
fillCenter Char
' '

{- | Render a table from column-major data. @columns@ has one 'V.Vector' per
column; widths are computed in one pass per column (no row-major transpose),
and row lines are built by indexing each column at row @i@.
-}
showTable ::
    RenderFormat ->
    [T.Text] ->
    [T.Text] ->
    [V.Vector T.Text] ->
    T.Text
showTable :: RenderFormat -> [Text] -> [Text] -> [Vector Text] -> Text
showTable RenderFormat
fmt [Text]
header [Text]
types [Vector Text]
columns =
    let isMarkdown :: Bool
isMarkdown = RenderFormat
fmt RenderFormat -> RenderFormat -> Bool
forall a. Eq a => a -> a -> Bool
== RenderFormat
Markdown
        esc :: Text -> Text
esc = if Bool
isMarkdown then Text -> Text
escapeMarkdownCell else Text -> Text
forall a. a -> a
id
        hdr :: [Text]
hdr = (Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Text
esc [Text]
header
        tys :: [Text]
tys = (Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Text
esc [Text]
types
        cols :: [Vector Text]
cols = (Vector Text -> Vector Text) -> [Vector Text] -> [Vector Text]
forall a b. (a -> b) -> [a] -> [b]
map ((Text -> Text) -> Vector Text -> Vector Text
forall a b. (a -> b) -> Vector a -> Vector b
V.map Text -> Text
esc) [Vector Text]
columns
        consolidatedHeader :: [Text]
consolidatedHeader =
            if Bool
isMarkdown
                then (Text -> Text -> Text) -> [Text] -> [Text] -> [Text]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Text
h Text
t -> Text
h Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"<br>" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
t) [Text]
hdr [Text]
tys
                else [Text]
hdr
        cs :: [ColDesc t]
cs = (Text -> ColDesc t) -> [Text] -> [ColDesc t]
forall a b. (a -> b) -> [a] -> [b]
map (\Text
h -> Filler -> Text -> Filler -> ColDesc t
forall t. Filler -> Text -> Filler -> ColDesc t
ColDesc Filler
center Text
h Filler
left) [Text]
consolidatedHeader
        nRows :: Int
nRows = case [Vector Text]
cols of
            (Vector Text
c : [Vector Text]
_) -> Vector Text -> Int
forall a. Vector a -> Int
V.length Vector Text
c
            [] -> Int
0
        columnMaxWidth :: Vector Text -> Int
columnMaxWidth Vector Text
col
            | Vector Text -> Bool
forall a. Vector a -> Bool
V.null Vector Text
col = Int
0
            | Bool
otherwise = (Int -> Text -> Int) -> Int -> Vector Text -> Int
forall a b. (a -> b -> a) -> a -> Vector b -> a
V.foldl' (\Int
acc Text
x -> Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
acc (Text -> Int
T.length Text
x)) Int
0 Vector Text
col
        widths :: [Int]
widths =
            (Text -> Text -> Vector Text -> Int)
-> [Text] -> [Text] -> [Vector Text] -> [Int]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3
                (\Text
h Text
t Vector Text
col -> Text -> Int
T.length Text
h Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Text -> Int
T.length Text
t Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` Vector Text -> Int
columnMaxWidth Vector Text
col)
                [Text]
consolidatedHeader
                [Text]
tys
                [Vector Text]
cols
        dashesOf :: Int -> Text
dashesOf Int
w = Filler
T.replicate Int
w Text
"-"
        border :: Text
border = Text -> [Text] -> Text
T.intercalate Text
"---" ((Int -> Text) -> [Int] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Text
dashesOf [Int]
widths)
        separator :: Text
separator = Text -> [Text] -> Text
T.intercalate Text
"-|-" ((Int -> Text) -> [Int] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Text
dashesOf [Int]
widths)
        fillCells :: (ColDesc t -> Int -> c -> Text) -> [c] -> Text
fillCells ColDesc t -> Int -> c -> Text
fill [c]
cells =
            Text -> [Text] -> Text
T.intercalate Text
" | " ((ColDesc t -> Int -> c -> Text)
-> [ColDesc t] -> [Int] -> [c] -> [Text]
forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 ColDesc t -> Int -> c -> Text
fill [ColDesc t]
forall {t}. [ColDesc t]
cs [Int]
widths [c]
cells)
        rowCells :: Int -> [Text]
rowCells Int
i = (Vector Text -> Text) -> [Vector Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Vector Text -> Int -> Text
forall a. Vector a -> Int -> a
V.! Int
i) [Vector Text]
cols
        rowLines :: [Text]
rowLines = [(ColDesc Any -> Filler) -> [Text] -> Text
forall {t} {c}. (ColDesc t -> Int -> c -> Text) -> [c] -> Text
fillCells ColDesc Any -> Filler
forall t. ColDesc t -> Filler
colValueFill (Int -> [Text]
rowCells Int
i) | Int
i <- [Int
0 .. Int
nRows Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        wrapMd :: Text -> Text
wrapMd Text
t = [Text] -> Text
T.concat [Text
"| ", Text
t, Text
" |"]
        outputLines :: [Text]
outputLines =
            if Bool
isMarkdown
                then
                    Text -> Text
wrapMd ((ColDesc Any -> Filler) -> [Text] -> Text
forall {t} {c}. (ColDesc t -> Int -> c -> Text) -> [c] -> Text
fillCells ColDesc Any -> Filler
forall t. ColDesc t -> Filler
colTitleFill [Text]
consolidatedHeader)
                        Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Text -> Text
wrapMd Text
separator
                        Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: (Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Text
wrapMd [Text]
rowLines
                else
                    Text
border
                        Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: (ColDesc Any -> Filler) -> [Text] -> Text
forall {t} {c}. (ColDesc t -> Int -> c -> Text) -> [c] -> Text
fillCells ColDesc Any -> Filler
forall t. ColDesc t -> Filler
colTitleFill [Text]
consolidatedHeader
                        Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Text
separator
                        Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: (ColDesc Any -> Filler) -> [Text] -> Text
forall {t} {c}. (ColDesc t -> Int -> c -> Text) -> [c] -> Text
fillCells ColDesc Any -> Filler
forall t. ColDesc t -> Filler
colTitleFill [Text]
tys
                        Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: Text
separator
                        Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
: [Text]
rowLines
     in [Text] -> Text
T.unlines [Text]
outputLines

{- | Escape a value for a GitHub-flavoured Markdown table cell: a bare @|@ would
end the cell and a newline would end the row, so both are neutralised (the pipe
backslash-escaped, the newline turned into a @\<br\>@).
-}
escapeMarkdownCell :: T.Text -> T.Text
escapeMarkdownCell :: Text -> Text
escapeMarkdownCell = HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
T.replace Text
"\n" Text
"<br>" (Text -> Text) -> (Text -> Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
T.replace Text
"|" Text
"\\|"